import os
from dotenv import load_dotenv
import streamlit as st
import pandas as pd
import plotly.express as px
from e2b_code_interpreter import Sandbox
from groq import Groq
import re
from streamlit_lottie import st_lottie
import requests

# Load environment variables
load_dotenv()

# Retrieve API keys from the environment
E2B_API_KEY = os.getenv("E2B_API_KEY")
GROQ_API_KEY = os.getenv("GROQ_API_KEY")

# Groq client and model setup
client = Groq(api_key=GROQ_API_KEY)
MODEL_NAME = 'llama-3.1-70b-versatile'

def load_lottie_url(url: str):
    r = requests.get(url)
    if r.status_code != 200:
        return None
    return r.json()

def upload_and_display_csv():
    uploaded_file = st.file_uploader("Upload your CSV file", type="csv")
    if uploaded_file:
        try:
            dataframe = pd.read_csv(uploaded_file)
            st.dataframe(dataframe.head(), use_container_width=True)
            return dataframe
        except Exception as e:
            st.error(f"Error loading CSV: {str(e)}")
            return None
    return None

def plot_interactive_graph(dataframe, x_column, y_column):
    fig = px.scatter(dataframe, x=x_column, y=y_column, title=f"{x_column} vs {y_column}", template="plotly_dark")
    fig.update_layout(font=dict(color="white"), title=dict(font=dict(size=20, color="white")))
    st.plotly_chart(fig, use_container_width=True)

st.set_page_config(page_title="Theaimart", layout="wide", page_icon="favicon.ico")

# Custom CSS to fix selectbox issues
st.markdown("""
    <style>
    /* Apply dark theme to entire app */
    .stApp {
        background-color: #0e1117;
        color: white;
    }
    
    /* General fix for inputs and text */
    .stTextInput, .stSelectbox, .stFileUploader {
        background-color: #1f2937 !important;
        color: white !important;
        border-radius: 5px;
        border: 1px solid #2563eb !important;
    }
    
    /* Ensure all inner elements inherit colors */
    .stSelectbox div[data-baseweb="select"] * {
        color: white !important;
        background-color: #1f2937 !important;
    }
    
    /* Fix dropdown text visibility */
    .stSelectbox div[data-baseweb="select"] div {
        color: white !important;
    }

    /* Dropdown hover and focus */
    .stSelectbox div[data-baseweb="select"]:hover {
        background-color: #2d3748 !important;
    }
    .stSelectbox div[data-baseweb="select"]:focus {
        background-color: #2d3748 !important;
        color: white !important;
    }

    /* Ensure SVG (arrow icon) visibility */
    .stSelectbox svg {
        fill: white !important;
    }
    </style>
    """, unsafe_allow_html=True)

# Main UI with two columns
st.title("GSV: AI-Powered Data Analysis App")

# Load Lottie animation
lottie_url = "https://assets5.lottiefiles.com/packages/lf20_qp1q7mct.json"
lottie_json = load_lottie_url(lottie_url)
st_lottie(lottie_json, height=200)

col1, col2 = st.columns([2, 1])

with col1:
    st.header("📈 Data Visualization")
    dataframe = upload_and_display_csv()
    if dataframe is not None:
        col_x, col_y = st.columns(2)
        with col_x:
            x_column = st.selectbox("Select X-axis", dataframe.columns)
        with col_y:
            y_column = st.selectbox("Select Y-axis", dataframe.columns)

        if st.button("Generate Interactive Plot"):
            plot_interactive_graph(dataframe, x_column, y_column)

with col2:
    st.header("🤖 Chat with CSV")
    user_query = st.text_area("Ask a question about your data")
    if st.button("Ask AI Assistant"):
        st.write("AI is processing...")

st.markdown("---")
st.markdown("Created with ❤️ by theaimart")