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_groq import ChatGroq
from langchain_openai import ChatOpenAI
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
from datasets import load_dataset
import tempfile

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

st.title("Blah Blah App Using CrewAI 🚀")
st.write("Analyze datasets using natural language queries powered by SQL and CrewAI.")

# Initialize LLM
llm = None

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


# API Key Validation and LLM Initialization
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.")
        llm = None
    else:
        llm = 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.")
        llm = None
    else:
        llm = ChatOpenAI(api_key=openai_api_key, model="gpt-4o")
        
# Initialize session state for data persistence
if "df" not in st.session_state:
    st.session_state.df = None

# Dataset Input
input_option = st.radio("Select Dataset Input:", ["Use Hugging Face Dataset", "Upload CSV File"])
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:
            with st.spinner("Loading dataset..."):
                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())
        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())


if st.session_state.df is not None:
    # Database setup
    temp_dir = tempfile.TemporaryDirectory()
    db_path = os.path.join(temp_dir.name, "patent_data.db")
    connection = sqlite3.connect(db_path)
    st.session_state.df.to_sql("patents", connection, if_exists="replace", index=False)
    db = SQLDatabase.from_uri(f"sqlite:///{db_path}")

    # SQL Tools
    @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)

    # --- CrewAI Agents for Patent Analysis ---
    patent_sql_dev = Agent(
        role="Patent Data Analyst",
        goal="Extract patent data using optimized SQL queries.",
        backstory="An expert in writing optimized SQL queries for complex patent databases.",
        llm=llm,
        tools=[list_tables, tables_schema, execute_sql],
    )

    patent_data_analyst = Agent(
        role="Patent Data Analyst",
        goal="Analyze the data and produce insights.",
        backstory="A seasoned analyst who identifies trends and patterns in datasets.",
        llm=llm,
    )

    patent_report_writer = Agent(
        role="Patent Report Writer",
        goal="Summarize patent insights into a clear report.",
        backstory="Expert in summarizing patent data insights into comprehensive reports.",
        llm=llm,
    )

    # --- Crew Tasks ---
    extract_data = Task(
        description="Extract patents related to the query: {query}.",
        expected_output="Patent data matching the query.",
        agent=patent_sql_dev,
    )

    analyze_data = Task(
        description="Analyze the extracted patent data for query: {query}.",
        expected_output="Analysis text summarizing findings.",
        agent=patent_data_analyst,
        context=[extract_data],
    )

    write_report = Task(
        description="Summarize analysis into an executive report.",
        expected_output="Markdown report of insights.",
        agent=patent_report_writer,
        context=[analyze_data],
    )

    # Assemble Crew
    crew = Crew(
        agents=[patent_sql_dev, patent_data_analyst, patent_report_writer],
        tasks=[extract_data, analyze_data, write_report],
        process=Process.sequential,
        verbose=True,
    )

    #Query Input for Patent Analysis
    query = st.text_area("Enter Patent Analysis Query:", placeholder="e.g., 'How many patents related to Machine Learning were filed after 2016?'")
    if st.button("Submit Query"):
        with st.spinner("Processing your query..."):
            inputs = {"query": query}
            result = crew.kickoff(inputs=inputs)
            st.markdown("### 📊 Patent Analysis Report")
            st.markdown(result)

    temp_dir.cleanup()
else:
    st.info("Please load a patent dataset to proceed.")