import os
import streamlit as st
import pandas as pd
from pandasai import Agent
from langchain_groq import ChatGroq

llm = ChatGroq(
    model=os.environ['MODEL_NAME'],
    temperature=0,
    max_retries=2,
)

# Initialize session state for URL and DataFrame
if 'pre_url' not in st.session_state:
    st.session_state['pre_url'] = ''
if 'df' not in st.session_state:
    st.session_state['df'] = None

# Streamlit app
st.title("Conversation Analysis")

# Introduction and Explanation
st.markdown("""
## Go beyond static visualizations!
This app lets you directly ask for the insights you need from Polis conversation data.  
Instead of being limited to a standard PCA chart, you can request custom plots and analyses using natural language.  
> :bulb: For example, try asking for: "**Show me a pie chart of the most common sentiment expressed in these comments.**"
""")

# Step 1: Choose conversation
opendata_options = [
    '15-per-hour-seattle', 'american-assembly.bowling-green', 'brexit-consensus',
    'canadian-electoral-reform', 'football-concussions', 'march-on.operation-marchin-orders',
    'scoop-hivemind.affordable-housing', 'scoop-hivemind.biodiversity',
    'scoop-hivemind.freshwater', 'scoop-hivemind.taxes', 'scoop-hivemind.ubi',
    'ssis.land-bank-farmland.2rumnecbeh.2021-08-01', 'vtaiwan.uberx'
]
selected_option = st.selectbox("Choose conversation", opendata_options)

url =  (
    f"https://raw.githubusercontent.com/compdemocracy/openData/master/{selected_option}/comments.csv"
)

# Load data only if URL changes
if st.session_state['pre_url'] != url:
    agent = Agent(
        pd.read_csv(url, index_col=0),
        config={"llm": llm}
    )
    st.session_state['agent'] = agent
    st.session_state['pre_url'] = url

# Step 2: Request for analysis or chart
with st.form('Analyze'):
    request = st.text_input(
        "Enter your analysis request:",
        "Plot a histogram about the distribution of agree, disagree and neutral comments against the topic")
    submit_button = st.form_submit_button(label="Analyze")

# Execute chat and display results
if submit_button:
    if st.session_state['agent'] is not None:
        file = st.session_state['agent'].chat(request)
        if os.path.exists(file):
            st.image(file)
        else:
            st.info(file)
    else:
        st.warning("Please select a conversation and load data first.")