import streamlit as st
import pandas as pd
import sqlite3
import os
from crewai import Agent, Crew, Process, Task
from crewai.tools import tool
from langchain_groq import ChatGroq
from langchain_openai import ChatOpenAI
from langchain_community.tools.sql_database.tool import (
    InfoSQLDatabaseTool,
    ListSQLDatabaseTool,
    QuerySQLDataBaseTool,
)
from langchain_community.utilities.sql_database import SQLDatabase
from datasets import load_dataset
import tempfile

st.title("Blah App")
st.write("Analyze datasets using natural language queries.")

def initialize_llm(model_choice):
    groq_api_key = os.getenv("GROQ_API_KEY")
    openai_api_key = os.getenv("OPENAI_API_KEY")

    if model_choice == "llama-3.3-70b":
        if not groq_api_key:
            st.error("Groq API key is missing. Please set the GROQ_API_KEY environment variable.")
            return None
        return ChatGroq(groq_api_key=groq_api_key, model="groq/llama-3.3-70b-versatile")
    elif model_choice == "GPT-4o":
        if not openai_api_key:
            st.error("OpenAI API key is missing. Please set the OPENAI_API_KEY environment variable.")
            return None
        return ChatOpenAI(api_key=openai_api_key, model="gpt-4o")

model_choice = st.radio("Select LLM", ["GPT-4o", "llama-3.3-70b"], index=0, horizontal=True)
llm = initialize_llm(model_choice)

def load_dataset_into_session():
    input_option = st.radio("Select Dataset Input:", ["Use Hugging Face Dataset", "Upload CSV File"], index=0, horizontal=True)
    if input_option == "Use Hugging Face Dataset":
        dataset_name = st.text_input("Enter Hugging Face Dataset Name:", value="HUPD/hupd")
        if st.button("Load Dataset"):
            try:
                dataset = load_dataset(dataset_name, name="sample", split="train", trust_remote_code=True, uniform_split=True)
                st.session_state.df = pd.DataFrame(dataset)
                st.success(f"Dataset '{dataset_name}' loaded successfully!")
                st.dataframe(st.session_state.df.head(10))
            except Exception as e:
                st.error(f"Error: {e}")
    elif input_option == "Upload CSV File":
        uploaded_file = st.file_uploader("Upload CSV File:", type=["csv"])
        if uploaded_file:
            st.session_state.df = pd.read_csv(uploaded_file)
            st.success("File uploaded successfully!")
            st.dataframe(st.session_state.df.head(10))

if "df" not in st.session_state:
    st.session_state.df = None
load_dataset_into_session()

def initialize_database(df):
    temp_dir = tempfile.TemporaryDirectory()
    db_path = os.path.join(temp_dir.name, "patent_data.db")
    connection = sqlite3.connect(db_path)
    df.to_sql("patents", connection, if_exists="replace", index=False)
    db = SQLDatabase.from_uri(f"sqlite:///{db_path}")
    return db, temp_dir

def create_sql_tools(db):
    @tool("list_tables")
    def list_tables() -> str:
        """List all tables in the patent database."""
        return ListSQLDatabaseTool(db=db).invoke("")

    @tool("tables_schema")
    def tables_schema(tables: str) -> str:
        """Get schema and sample rows for given tables."""
        return InfoSQLDatabaseTool(db=db).invoke(tables)

    @tool("execute_sql")
    def execute_sql(sql_query: str) -> str:
        """Execute a SQL query against the patent database."""
        return QuerySQLDataBaseTool(db=db).invoke(sql_query)

    return list_tables, tables_schema, execute_sql

def initialize_agents(llm, tools):
    list_tables, tables_schema, execute_sql = tools

    sql_agent = Agent(
        role="Patent Data Analyst",
        goal="Extract patent data using optimized SQL queries.",
        backstory="Expert in optimized SQL for patent databases.",
        llm=llm,
        tools=[list_tables, tables_schema, execute_sql],
    )

    analyst_agent = Agent(
        role="Patent Data Analyst",
        goal="Analyze the data and produce insights.",
        backstory="Data analyst identifying trends.",
        llm=llm,
    )

    writer_agent = Agent(
        role="Patent Report Writer",
        goal="Summarize patent insights into a report.",
        backstory="Expert in clear, concise reporting.",
        llm=llm,
    )

    return sql_agent, analyst_agent, writer_agent

def setup_crew(sql_agent, analyst_agent, writer_agent):
    extract_task = Task(
        description="Extract patents related to the query: {query}.",
        expected_output="Patent data matching the query.",
        agent=sql_agent,
    )

    analyze_task = Task(
        description="Analyze the extracted patent data.",
        expected_output="Analysis text summarizing findings.",
        agent=analyst_agent,
        context=[extract_task],
    )

    report_task = Task(
        description="Summarize analysis into a report.",
        expected_output="Markdown report of insights.",
        agent=writer_agent,
        context=[analyze_task],
    )

    return Crew(
        agents=[sql_agent, analyst_agent, writer_agent],
        tasks=[extract_task, analyze_task, report_task],
        process=Process.sequential,
        verbose=True,
    )

if st.session_state.df is not None:
    db, temp_dir = initialize_database(st.session_state.df)
    tools = create_sql_tools(db)
    sql_agent, analyst_agent, writer_agent = initialize_agents(llm, tools)
    crew = setup_crew(sql_agent, analyst_agent, writer_agent)

    query = st.text_area("Enter Patent Analysis Query:", value="How many patents related to Machine Learning were filed after 2016?")
    if st.button("Submit Query"):
        if query.strip():
            with st.spinner("Processing your query..."):
                result = crew.kickoff(inputs={"query": query})
                st.markdown("### 📊 Patent Analysis Report")
                st.markdown(result)
        else:
            st.warning("Please enter a valid query.")
else:
    st.info("Please load a patent dataset to proceed.")