import streamlit as st
import streamlit.components.v1 as components
import openai
import subprocess
import re
import os
import pandas as pd
import socket
import time

csv_name = "./Financial_Sample.csv"
try:
    df = pd.read_csv(csv_name)
    metadata = {
        "columns": df.columns.tolist(),
        "dtypes": df.dtypes.apply(lambda x: x.name).to_dict(),
        "shape": df.shape,
    }
except Exception as e:
    st.error(f"Error loading CSV file: {e}")
    st.stop()


def is_port_in_use(port: int) -> bool:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        return s.connect_ex(("0.0.0.0", port)) == 0


def find_free_port(starting_port: int) -> int:
    port = starting_port
    while is_port_in_use(port):
        port += 1
    return port


def runcode(code: str) -> tuple:
    with open("code.py", "w") as file:
        file.write(code)

    starting_port = 8501
    free_port = find_free_port(starting_port)

    try:
        process = subprocess.Popen(
            ["streamlit", "run", "code.py", "--server.address", "0.0.0.0","--server.port", str(free_port)],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )

        time.sleep(5)

        if process.poll() is None:
            return (
                free_port,
                f"Streamlit process started successfully on port {free_port} on Server Address 0.0.0.0 ",
            )
        else:
            stdout, stderr = process.communicate()
            return (None, f"Streamlit process exited unexpectedly. Error: {stderr}")

    except FileNotFoundError:
        return (None, "Error: Streamlit is not installed or not in PATH")
    except Exception as e:
        return (None, f"An unexpected error occurred: {str(e)}")


class Agent:
    def __init__(self, system_prompt="", known_actions=None):
        self.system = system_prompt
        self.messages = []
        self.known_actions = known_actions if known_actions is not None else {}

        self.client = openai.OpenAI(
            api_key= os.environ.get('TOGETHER_API_KEY'),
            base_url="https://api.together.xyz/v1",
        )
        self.messages.append({"role": "system", "content": self.system})

    def __call__(self, message):
        self.messages.append({"role": "user", "content": message})
        result = self.execute()
        self.messages.append({"role": "assistant", "content": result})
        return result

    def execute(self):
        try:
            response = self.client.chat.completions.create(
                model="meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
                stop=["PAUSE"],
                messages=self.messages,
            )
            return response.choices[0].message.content
        except Exception as e:
            return f"Error executing model: {str(e)}"

    def query(self, question, max_turns=5):
        i = 0
        next_prompt = question
        while i < max_turns:
            i += 1
            result = self(next_prompt)
            st.session_state.logs += f"Assistant: {result}\n"

            action_re = re.search(r"Action: (\w+): (.*)", result, re.DOTALL)
            if action_re:
                action = action_re.group(1)
                action_input = action_re.group(2).strip()
                st.session_state.logs += (
                    f"Action: {action}\nAction Input: {action_input}\n"
                )

                if action not in self.known_actions:
                    error_msg = f"Unknown action: {action}: {action_input}"
                    raise Exception(error_msg)

                st.session_state.logs += (
                    f" ------------------------\n running {action} {action_input}\n"
                )
                observation = self.known_actions[action](action_input)
                st.session_state.logs += f"Observation: {observation}\n"
                next_prompt = f"Observation: {observation}"
            else:
                print("No action found, ending query")
                return


known_actions = {"runcode": runcode}
prompt = f"""
You are an expert in creating advanced Interactive Streamlit Dashboards in python based on user Query and you run in a loop of Thought, Action, PAUSE, Observation.
At the end of the loop you output an Answer
Use Thought to describe your thoughts about the question you have been asked.
Use Action to run one of the actions available to you - then return PAUSE.
Observation will be the result of running those actions.

Always return just code no need of ```

Your Task is help user get result of query about below dataset,Decide based on user query to make Dashboard or Just Textual Answer.
Here is the metadata of the dataset and name of dataset is Financial_Sample.csv:

Columns: {metadata['columns']}
Dtypes: {metadata['dtypes']}
Shape: {metadata['shape']}
You can use this metadata to generate results.

Your available actions are:
runcode

How to use actions
Action : action_name: input_to_action

if input_to_action is code then don't use ``` just write code
Example session:

Question:  Give me a dashboard to visualize the people height and age

Thought: I need to run a create a dashboard where i can visualize the Subscriptions and Country data

Action: runcode: import streamlit as st
import pandas as pd
import plotly.express as px
from datetime import datetime

@st.cache_data
def load_data():
    df = pd.read_csv(f{csv_name})
    return df

df = load_data()
st.title("Height vs Age Visualization")
fig = px.scatter(df, x='age', y='height', opacity=0.6)
fig.update_layout(
    xaxis_title="Age (years)",
    yaxis_title="Height (cm)",
    title="Height vs Age Scatter Plot"
)

st.plotly_chart(fig, use_container_width=True)

PAUSE

Observation :  understand the output based its stdout and take necessary steps. 

Answer: Final Answer for User Request if its Dashboard send "Please visit link to view dashboard"or Textual Answer "Your Interpretation of Answer"
""".strip()

if "bot" not in st.session_state:
    st.session_state.bot = Agent(system_prompt=prompt, known_actions=known_actions)

st.set_page_config(layout="wide")
st.title("Customer Data Analysis")

if "logs" not in st.session_state:
    st.session_state.logs = ""

if "messages" not in st.session_state:
    st.session_state.messages = []

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

if query := st.chat_input("Enter your query about the dataset"):
    st.session_state.messages.append({"role": "user", "content": query})
    with st.chat_message("user"):
        st.markdown(query)
    st.session_state.logs = ""
    with st.spinner("Generating response..."):
        st.session_state.bot.query(query)

    answer_match = re.search(r"Answer:", st.session_state.logs)

    if answer_match:
        answer_end_index = answer_match.end()
        subsequent_text = st.session_state.logs[answer_end_index:].strip()
    else:
        st.warning("No answer found in the generated response.")

    with st.chat_message("assistant"):
        url_pattern = r"(https?://[^\s]+)"
        url = re.findall(url_pattern, subsequent_text)
        if url:
            components.iframe(src=url[0], width=800, height=600)
            st.write(subsequent_text)
        else:
            st.success(subsequent_text)
    st.session_state.messages.append({"role": "assistant", "content": subsequent_text})

with st.sidebar:
    st.title("Dataset Metadata")
    st.write("Columns:", metadata["columns"])
    st.write("Dtypes:", metadata["dtypes"])
    st.write("Shape:", metadata["shape"])
    st.write("Sample Data:")
    st.write(df.head())

    with st.expander("Logs"):
        st.code(st.session_state.logs)