import os
from typing import Optional, Tuple

import gradio as gr
from query_data import get_chain
from threading import Lock
import pinecone
from langchain.vectorstores import Chroma, Pinecone
from langchain.embeddings.openai import OpenAIEmbeddings

embeddings = OpenAIEmbeddings()

PINECONE_API_KEY = '6af52b8a-a3df-4189-899b-b21163027bb8'
PINECONE_API_ENV = 'asia-southeast1-gcp'

# initialize pinecone
pinecone.init(
    api_key=PINECONE_API_KEY,  # find at app.pinecone.io
    environment=PINECONE_API_ENV  # next to api key in console
)

index_name = "twimbit-answer"
vectorstore = Pinecone.from_existing_index(index_name=index_name, embedding=embeddings)

api_key = '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'


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 []
            # 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

            if api_key:
                os.environ["OPENAI_API_KEY"] = api_key
                chain = get_chain(vectorstore)
                os.environ["OPENAI_API_KEY"] = ""

            import openai
            openai.api_key = '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'
            # 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: #111827};footer "
                      "{visibility: hidden};")

with block:
    # with gr.Row():
    #     openai_api_key_textbox = gr.Textbox(
    #         placeholder="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",
    #         show_label=False,
    #         lines=1,
    #         type="password",
    #         value="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"
    #     )

    chatbot = gr.Chatbot().style(height=500)

    with gr.Row():
        message = gr.Textbox(
            label="What's your question?",
            placeholder="Ask questions about reports",
            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])
    message.submit(chat, inputs=[message, state, agent_state], outputs=[chatbot, state])

    # openai_api_key_textbox.change(
    #     set_openai_api_key,
    #     inputs=[openai_api_key_textbox],
    #     outputs=[agent_state],
    # )

# block.launch(debug=True)
block.launch(debug=True, auth=('admin', 'password'), auth_message='enter username password to proceed further')