import os
from typing import Optional, Tuple
from threading import Lock
import pickle

import gradio as gr
from query_data import get_chain



class ChatWrapper:
    def __init__(self):
        self.lock = Lock()

    def set_openai_api_key(self, api_key: str):
        with open("vectorstore.pkl", "rb") as f:
            vectorstore = pickle.load(f)
        """Set the api key and return chain.
        If no api_key, then None is returned.
        """
        if api_key:
            os.environ["OPENAI_API_KEY"] = api_key
            chain = get_chain(vectorstore)
            os.environ["OPENAI_API_KEY"] = ""
            return chain

    def __call__(self, inp: str, history: Optional[Tuple[str, str]]):
        self.lock.acquire()
        api_key = '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'
        chain = self.set_openai_api_key(api_key)
        try:
            history = history or []
            # If chain is None, that is because no API key was provided.
            if chain is None:
                history.append((inp, "Please paste your OpenAI key to use"))
                return inp, history
            # Set OpenAI key
            #import openai
            #openai.api_key = api_key
            # Run chain and append input.
            output = chain({"question": inp, "chat_history": history})["answer"]
            history.append((inp, output))
            chatResult = (output, history)
        except Exception as e:
            raise e
        finally:
            self.lock.release()
        return chatResult

chat = ChatWrapper()
state = gr.outputs.State()

gradio_interface = gr.Interface(chat, inputs=["text", state], outputs=["text", state])
gradio_interface.launch(debug=True)