from data_sources import process_data_upload

import gradio as gr
import json

from haystack.dataclasses import ChatMessage
from haystack.components.generators.chat import OpenAIChatGenerator

import os
from getpass import getpass
from dotenv import load_dotenv

load_dotenv()

if "OPENAI_API_KEY" not in os.environ:
    os.environ["OPENAI_API_KEY"] = getpass("Enter OpenAI API key:")

chat_generator = OpenAIChatGenerator(model="gpt-4o")
response = None
messages = [
    ChatMessage.from_system(
        "You are a helpful and knowledgeable agent who has access to an SQL database which has a table called 'data_source'"
    )
]

def chatbot_with_fc(message, history):
    print("CHATBOT FUNCTIONS")
    from functions import sqlite_query_func
    from pipelines import rag_pipeline_func
    import tools
    import importlib
    importlib.reload(tools)

    available_functions = {"sql_query_func": sqlite_query_func, "rag_pipeline_func": rag_pipeline_func}
    messages.append(ChatMessage.from_user(message))
    response = chat_generator.run(messages=messages, generation_kwargs={"tools": tools.tools})

    while True:
        # if OpenAI response is a tool call
        if response and response["replies"][0].meta["finish_reason"] == "tool_calls":
            function_calls = json.loads(response["replies"][0].content)
            for function_call in function_calls:
                ## Parse function calling information
                function_name = function_call["function"]["name"]
                function_args = json.loads(function_call["function"]["arguments"])

                ## Find the correspoding function and call it with the given arguments
                function_to_call = available_functions[function_name]
                function_response = function_to_call(**function_args)
                ## Append function response to the messages list using `ChatMessage.from_function`
                messages.append(ChatMessage.from_function(content=function_response['reply'], name=function_name))
                response = chat_generator.run(messages=messages, generation_kwargs={"tools": tools.tools})

        # Regular Conversation
        else:
            messages.append(response["replies"][0])
            break
    return response["replies"][0].content

css= ".file_marker .large{min-height:50px !important;}"

with gr.Blocks(css=css) as demo:
    title = gr.HTML("<h1 style='text-align:center;'>Virtual Data Analyst</h1>")
    description = gr.HTML("<p style='text-align:center;'>Upload a CSV file and chat with our virtual data analyst to get insights on your data set</p>")
    file_output = gr.File(label="CSV File", show_label=True, elem_classes="file_marker", file_types=['.csv'])

    @gr.render(inputs=file_output)
    def data_options(filename):
        print(filename)
        if filename:
            bot = gr.Chatbot(type='messages', label="CSV Chat Window", show_label=True, render=False, visible=True, elem_classes="chatbot")
            chat = gr.ChatInterface(
                                fn=chatbot_with_fc,
                                type='messages',
                                chatbot=bot,
                                title="Chat with your data file",
                                examples=[
                                    ["Describe the dataset"],
                                    ["List the columns in the dataset"],
                                    ["What could this data be used for?"],
                                ],
                                ) 
            
            process_upload(filename)
    
    def process_upload(upload_value):
        if upload_value:
            print("UPLOAD VALUE")
            print(upload_value)
            process_data_upload(upload_value)
        return [], []