import streamlit as st
import pandas as pd
import sqlite3
import os
import json
from pathlib import Path
from datetime import datetime, timezone
from crewai import Agent, Crew, Process, Task
from crewai_tools import tool
from langchain_core.prompts import ChatPromptTemplate
from langchain_groq import ChatGroq
from langchain.schema.output import LLMResult
from langchain_core.callbacks.base import BaseCallbackHandler
from langchain_community.tools.sql_database.tool import (
    InfoSQLDatabaseTool,
    ListSQLDatabaseTool,
    QuerySQLCheckerTool,
    QuerySQLDataBaseTool,
)
from langchain_community.utilities.sql_database import SQLDatabase
import tempfile

# Setup GROQ API Key
os.environ["GROQ_API_KEY"] = st.secrets.get("GROQ_API_KEY", "")

# Callback handler for logging LLM responses
class Event:
    def __init__(self, event, text):
        self.event = event
        self.timestamp = datetime.now(timezone.utc).isoformat()
        self.text = text

class LLMCallbackHandler(BaseCallbackHandler):
    def __init__(self, log_path: Path):
        self.log_path = log_path

    def on_llm_start(self, serialized, prompts, **kwargs):
        with self.log_path.open("a", encoding="utf-8") as file:
            file.write(json.dumps({"event": "llm_start", "text": prompts[0], "timestamp": datetime.now().isoformat()}) + "\n")

    def on_llm_end(self, response: LLMResult, **kwargs):
        generation = response.generations[-1][-1].message.content
        with self.log_path.open("a", encoding="utf-8") as file:
            file.write(json.dumps({"event": "llm_end", "text": generation, "timestamp": datetime.now().isoformat()}) + "\n")

# LLM Setup
llm = ChatGroq(
    temperature=0,
    model_name="mixtral-8x7b-32768",
    callbacks=[LLMCallbackHandler(Path("prompts.jsonl"))],
)

# App Header
st.title("SQL-RAG with CrewAI 🚀")
st.write("Provide your query, and the app will extract, analyze, and summarize the data dynamically.")

# File Upload for Dataset
uploaded_file = st.file_uploader("Upload your dataset (CSV file)", type=["csv"])

if uploaded_file:
    st.success("File uploaded successfully!")
    
    # Temporary directory for SQLite DB
    temp_dir = tempfile.TemporaryDirectory()
    db_path = os.path.join(temp_dir.name, "data.db")
    
    # Create SQLite database
    df = pd.read_csv(uploaded_file)
    connection = sqlite3.connect(db_path)
    df.to_sql("data_table", connection, if_exists="replace", index=False)
    
    db = SQLDatabase.from_uri(f"sqlite:///{db_path}")

    # Tools
    @tool("list_tables")
    def list_tables() -> str:
        return ListSQLDatabaseTool(db=db).invoke("")
    
    @tool("tables_schema")
    def tables_schema(tables: str) -> str:
        return InfoSQLDatabaseTool(db=db).invoke(tables)

    @tool("execute_sql")
    def execute_sql(sql_query: str) -> str:
        return QuerySQLDataBaseTool(db=db).invoke(sql_query)

    @tool("check_sql")
    def check_sql(sql_query: str) -> str:
        return QuerySQLCheckerTool(db=db, llm=llm).invoke({"query": sql_query})

    # Agents
    sql_dev = Agent(
        role="Senior Database Developer",
        goal="Extract data from the database based on user query",
        llm=llm,
        tools=[list_tables, tables_schema, execute_sql, check_sql],
        allow_delegation=False,
    )

    data_analyst = Agent(
        role="Senior Data Analyst",
        goal="Analyze the database response and provide insights",
        llm=llm,
        allow_delegation=False,
    )

    report_writer = Agent(
        role="Senior Report Editor",
        goal="Summarize the analysis into a short report",
        llm=llm,
        allow_delegation=False,
    )

    # Tasks
    extract_data = Task(
        description="Extract data required for the query: {query}.",
        expected_output="Database result for the query",
        agent=sql_dev,
    )

    analyze_data = Task(
        description="Analyze the data and generate insights for: {query}.",
        expected_output="Detailed analysis text",
        agent=data_analyst,
        context=[extract_data],
    )

    write_report = Task(
        description="Summarize the analysis into a concise executive report.",
        expected_output="Markdown report",
        agent=report_writer,
        context=[analyze_data],
    )

    # Crew
    crew = Crew(
        agents=[sql_dev, data_analyst, report_writer],
        tasks=[extract_data, analyze_data, write_report],
        process=Process.sequential,
        verbose=2,
        memory=False,
    )

    # User Input Query
    query = st.text_input("Enter your query:")
    if query:
        with st.spinner("Processing your query..."):
            inputs = {"query": query}
            result = crew.kickoff(inputs=inputs)
            st.markdown("### Analysis Report:")
            st.markdown(result)

    # Clean up
    temp_dir.cleanup()