import streamlit as st
import os
import pandas as pd
import random
from os.path import join
from src import preprocess_and_load_df, load_agent, ask_agent, decorate_with_code, show_response, get_from_user, load_smart_df, ask_question
from dotenv import load_dotenv
from langchain_groq.chat_models import ChatGroq

load_dotenv("Groq.txt")
Groq_Token = os.environ["GROQ_API_KEY"]
models = {"mixtral": "mixtral-8x7b-32768", "llama": "llama2-70b-4096", "gemma": "gemma-7b-it"}

self_path = os.path.dirname(os.path.abspath(__file__))

# Using HTML and CSS to center the title
st.write(
    """
    <style>
    .title {
        text-align: center;
        color: #17becf;
    }
""",
    unsafe_allow_html=True,
)

# Displaying the centered title
st.markdown("<h2 class='title'>VayuBuddy</h2>", unsafe_allow_html=True)

# os.environ["PANDASAI_API_KEY"] = "$2a$10$gbmqKotzJOnqa7iYOun8eO50TxMD/6Zw1pLI2JEoqncwsNx4XeBS2"

# with open(join(self_path, "context1.txt")) as f:
#     context = f.read().strip()

# agent = load_agent(join(self_path, "app_trial_1.csv"), context)
# df = preprocess_and_load_df(join(self_path, "Data.csv"))
# inference_server = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.2"
# inference_server = "https://api-inference.huggingface.co/models/codellama/CodeLlama-13b-hf"
# inference_server = "https://api-inference.huggingface.co/models/pandasai/bamboo-llm"

model_name = st.sidebar.selectbox("Select LLM:", ["mixtral", "llama", "gemma"])

questions = ('Custom Prompt',
             'Plot the monthly average PM2.5 for the year 2023.',
             'Which month has the highest average PM2.5 overall?',
             'Which month has the highest PM2.5 overall?',
             'Which month has the highest average PM2.5 in 2023 for Mumbai?',
             'Plot and compare monthly timeseries of pollution for Mumbai and Bengaluru.',
             'Plot the yearly average PM2.5.',
             'Plot the monthly average PM2.5 of Delhi',
             'Mumbai and Bengaluru for the year 2022.',
             'Which month has the highest pollution?',
             'Plot the monthly average PM2.5 of Delhi for the year 2022.',
             'Which city has the highest PM2.5 level in July 2022?',
             'Plot and compare monthly timeseries of PM2.5 for Mumbai and Bengaluru.',
             'Plot and compare the monthly average PM2.5 of Delhi, Mumbai and Bengaluru for the year 2022.',
             'Plot the monthly average PM2.5.',
             'Plot the monthly average PM10 for the year 2023.',
             'Which month has the highest PM2.5?',
             'Plot the monthly average PM2.5 of Delhi for the year 2022.',
             'Plot the monthly average PM2.5 of Bengaluru for the year 2022.',
             'Plot the monthly average PM2.5 of Mumbai for the year 2022.',
             'Which state has the highest average PM2.5?',
             'Plot monthly PM2.5 in Gujarat for 2023.',
             'What is the name of the month with the highest average PM2.5 overall?')

waiting_lines = ("Thinking...", "Just a moment...", "Let me think...", "Working on it...", "Processing...", "Hold on...", "One moment...", "On it...")

# agent = load_agent(df, context="", inference_server=inference_server, name=model_name)

# Initialize chat history
if "responses" not in st.session_state:
    st.session_state.responses = []

# Display chat responses from history on app rerun
for response in st.session_state.responses:
    if not response["no_response"]:
        show_response(st, response)

show = True

if prompt := st.sidebar.selectbox("Select a Prompt:", questions):

    # add a note "select custom prompt to ask your own question"
    st.sidebar.info("Select 'Custom Prompt' to ask your own question.")

    if prompt == 'Custom Prompt':
        show = False
        # React to user input
        prompt = st.chat_input("Ask me anything about air quality!", key=10)
        if prompt : show = True   
    if show :
        
        # Add user input to chat history
        response = get_from_user(prompt)
        response["no_response"] = False
        st.session_state.responses.append(response)
        
        # Display user input
        show_response(st, response)

        no_response = False

        # select random waiting line
        with st.spinner(random.choice(waiting_lines)):
            ran = False
            for i in range(5):
                llm = ChatGroq(model=models[model_name], api_key=os.getenv("GROQ_API"), temperature=0.1)

                df_check = pd.read_csv("Data.csv")
                df_check["Timestamp"] = pd.to_datetime(df_check["Timestamp"])
                df_check = df_check.head(5)

                new_line = "\n"

                template = f"""```python
import pandas as pd
import matplotlib.pyplot as plt

df = pd.read_csv("Data.csv")
df["Timestamp"] = pd.to_datetime(df["Timestamp"])

# df.dtypes
{new_line.join(map(lambda x: '# '+x, str(df_check.dtypes).split(new_line)))}

# {prompt.strip()}
# <your code here>
```
"""

                query = f"""I have a pandas dataframe data of PM2.5 and PM10.
                * Frequency of data is daily. 
                * `pollution` generally means `PM2.5`.
                * DOn't print, but save result in a variable `answer` and make it global.
                * If result is a plot make it in tight layout, save it and save path in `answer`. Example: `answer='plot.png'`
                * If result is not a plot, save it as a string in `answer`. Example: `answer='The city is Mumbai'`

                Complete the following code.

                {template}

                """
                
                answer = llm.invoke(query)
                code = f"""
                {template.split("```python")[1].split("```")[0]}
                {answer.content.split("```python")[1].split("```")[0]}
                """
                # update variable `answer` when code is executed
                try:
                    exec(code)
                    ran = True
                    no_response = False
                except Exception as e:
                    no_response = True
                    exception = e

                response = {"role": "assistant", "content": answer, "gen_code": code, "ex_code": code, "last_prompt": prompt, "no_response": no_response}

                # Get response from agent
                # response = ask_question(model_name=model_name, question=prompt)
                # response = ask_agent(agent, prompt)

                if ran:
                    break
            
        if no_response:
            st.error(f"Failed to generate right output due to the following error:\n\n{exception}")
        # Add agent response to chat history
        st.session_state.responses.append(response)
        
        # Display agent response
        if not no_response:
            show_response(st, response)

        del prompt