import os
from typing import Optional, Tuple

import gradio as gr
import pickle
from query_data import get_chain
from threading import Lock

with open("vectorstore.pkl", "rb") as f:
    vectorstore = pickle.load(f)



class ChatWrapper:

    def __init__(self):
        self.lock = Lock()
    def __call__(
        self, inp: str, history: Optional[Tuple[str, str]],chain
    ):
        """Execute the chat functionality."""
        self.lock.acquire()
        try:
            history = history or []
            chain = get_chain(vectorstore)
            # 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 history, history
            # Set OpenAI key
            import openai
            
            openai.api_key = os.environ["OPENAI_API_KEY"]
            # Run chain and append input.
            output = chain({"question": inp, "chat_history": history})["answer"]
            history.append((inp, output))
        except Exception as e:
            raise e
        finally:
            self.lock.release()
        return history, history

chat = ChatWrapper()

block = gr.Blocks(css=".gradio-container {background-color: #04A7E3}")

image_url = "https://www.thebotforge.io/wp-content/uploads/2022/01/the-bot-forge-logo-100px.png"
html = "<center><img src='%s' width='300'></center>" % image_url
with block:
    with gr.Row():
        gr.HTML(html)
        
    with gr.Row():
        gr.Markdown("<h2><center>The Bot Forge</h2></center><h3> <center>Chat about - State-of-the-Union 2023</center></h3>")

        # openai_api_key_textbox = gr.Textbox(
        #     placeholder="Paste your OpenAI API key (sk-...)",
        #     show_label=False,
        #     lines=1,
        #     type="password",
        # )

    chatbot = gr.Chatbot()

    with gr.Row():
        message = gr.Textbox(
            label="What's your question?",
            placeholder="Ask questions about the most recent state of the union",
            lines=1,
        )
        submit = gr.Button(value="Send", variant="secondary").style(full_width=False)

    gr.Examples(
        examples=[
            "What did the president say about Kentaji Brown Jackson",
            "Did he mention Stephen Breyer?",
            "What was his stance on Ukraine",
        ],
        inputs=message,
    )

  
    state = gr.State()
    agent_state = gr.State()

    submit.click(chat, inputs=[message, state, agent_state], outputs=[chatbot, state])
    

    
def echo(name, request: gr.Request):
    if request:
        print("Request headers dictionary:", request.headers)
        print("IP address:", request.client.host)
        print("Body", request.body)
    return name


block.launch(debug=True)