import chainlit as cl
from utils import load_details_dataset, load_election_dataset, load_maha_election_dataset
from sqlite3 import connect
from typing import cast
from utils.load_llm import load_llm
from dotenv import load_dotenv
from utils.query_generator import sql_generator, sql_formatter, analyze_results
from langchain.schema.runnable import Runnable
from utils.sql_runtime import SQLRuntime

load_dotenv()

# global variables
db_path = './data/elections.db'
sql_runtime = SQLRuntime(dbname=db_path)

# Load the dataset
@cl.action_callback("Load Datasets")
async def on_action(action: cl.Action):
    print("Loading datasets...")

    # save the datasets as tables
    conn = connect('./data/elections.db')
    
    load_details_dataset.load_data_from_csv_to_db('./data/details_of_assembly_segment_2019.csv', conn)
    load_election_dataset.load_data_from_csv_to_db('./data/eci_data_2024.csv', conn)
    load_maha_election_dataset.load_data_from_csv_to_db('./data/maha_results_2019.csv', conn)

    return "Datasets loaded successfully."

@cl.action_callback("Execute Query")
async def on_action(action: cl.Action):
    res = await cl.AskUserMessage(content="Enter Query to run Manually", timeout=20).send()
    actions = [
        cl.Action(name="Execute Query", description="Execute the query on the dataset", value="Execute Query")
    ]
    if res:
        query = res['output']
        res = sql_runtime.execute(query)
        print(res)
        if res["code"] == 0:
            data = ""
            if res["data"]:
                for row in res["data"]:
                    data += str(row) + "\n"

            elements = [
                cl.Text(name="Result", content=data, display="inline"),
            ]
            await cl.Message(
                content=f"Query: {query}",
                elements=elements,
                actions=actions,
            ).send()
        else:
            error = res["msg"]["traceback"]
            elements = [   
                cl.Text(name="Error", content=error, display="inline"),
            ]
            await cl.Message(
                content=f"Query: {query}",
                elements=elements,
                actions=actions,
            ).send()

    # return "Query executed successfully."

@cl.on_chat_start
async def start():
    # Sending an action button within a chatbot message
    actions = [
        cl.Action(name="Load Datasets", description="Load the datasets into the database", value="Load Datasets")
    ]

    chain = sql_generator | sql_formatter | analyze_results

    cl.user_session.set("chain", chain)
    cl.user_session.set("db_path", './data/elections.db')

    await cl.Message(content="I am your personal political expert. I can help you analyze the election data. Click the button below to load the datasets.", actions=actions).send()

@cl.on_message
async def on_message(message: cl.Message):
    chain = cast(Runnable, cl.user_session.get("chain"))
    db_path = cl.user_session.get("db_path")

    actions = [
        cl.Action(name="Execute Query", description="Execute the query on the dataset", value="Execute Query")
    ]

    print(message.content)

    try:
        res = chain.invoke({
            "query": message.content,
            "db_path": db_path
        })
    except Exception as e:
        print(e)
        await cl.Message(content="An error occurred while processing the query. Please try again.").send()
        return

    queries = "\n".join(res.queries)

    errors = "".join(res.errors)

    elements = [
        cl.Text(name='results', content=res.summary, display="inline"),
        cl.Text(name="queries", content=queries, display="inline"),
    ]

    if errors:
        elements.append(cl.Text(name="errors", content=errors, display="inline"))

    await cl.Message(
        content="Let's analyze the results of the query",
        elements=elements,
        actions=actions
    ).send()