import os
from typing import List, Tuple
import gradio as gr
from assistant import OAIAssistant
import utils
from dotenv import load_dotenv

load_dotenv()


class VisaPolicyAgent:
    def __init__(self) -> None:
        assistant_id = os.environ.get("ASSISTANT_ID")
        vector_store_id = os.environ.get("VECTORSTORE_ID")

        self.assistant = OAIAssistant(
            assistant_id=assistant_id, vectorstore_id=vector_store_id
        )
        self.thread_id = self.assistant.create_thread().id

    def chat(self, message: str, history: List[str]):
        response = self.assistant.chat(message, self.thread_id)
        history.append((message, response["answer"]))
        return (
            "",
            history,
        )

    def create_thread(self):
        try:
            self.assistant.delete_thread(self.thread_id)
            self.thread_id = self.assistant.create_thread().id
            gr.Info(message="New thread created. Start as New!")
        except Exception as e:
            gr.Error(message=f"Unable to create new thread. Error: {e}")

    def clear_all(self):
        # create new thread
        # delete previous files
        # create new chatbot

        self.create_thread()
        gr.Info(message="Created new thread")
        try:
            assistant_file_ids = self.assistant.get_files_list()
            for file_id in assistant_file_ids:
                self.assistant.remove_file(file_id=file_id)

            gr.Info(message="Deleted files in assistant")

        except Exception as e:
            gr.Error(message=f"Unable to delete files. Error: {e}")

        gr.Info("Chat is cleared.")
        
        return [("Clear Chatbot", "Chatbot cleared.")]

    def add_file(self, file: gr.File):
        self.assistant.add_file(file)

    def bot_response(self, message, history: List[Tuple]) -> List[Tuple]:
        
        response = self.assistant.chat(message, self.thread_id)
        history.append((message, response))
        return "", history
    
    def create_interface(self):
        with gr.Blocks() as demo:
            chatbot = gr.Chatbot(height=400)
            message = gr.Textbox(
                placeholder="Ask me anything",
                scale=7,
                label="Input message",
            )
            new_thread_button = gr.Button(value="Create New Thread")
            clear_button = gr.Button(value="Clear All")

            _ = gr.ClearButton([message, chatbot])

            message.submit(
                self.bot_response,
                inputs=[message, chatbot],
                outputs=[message, chatbot],
            )
            new_thread_button.click(self.create_thread)
            clear_button.click(self.clear_all, outputs=chatbot)
        
        return demo


if __name__ == "__main__":
    vp_chatbot = VisaPolicyAgent()

    interface = vp_chatbot.create_interface()
    interface.launch()