import os
import time
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 Pinecone
from langchain.embeddings.openai import OpenAIEmbeddings


PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
PINECONE_API_ENV = os.getenv("PINECONE_API_ENV")
PINECONE_INDEX = os.getenv("PINECONE_INDEX")

# 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 = PINECONE_INDEX

embeddings = OpenAIEmbeddings()

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

vectorstore = Pinecone.from_existing_index(index_name=index_name, embedding=embeddings)

chain = get_chain(vectorstore)


class ChatWrapper:

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

    def __call__(
            self, inp: str, history: Optional[Tuple[str, str]]
    ):
        """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

            import openai
            openai.api_key = os.getenv("OPENAI_API_KEY")
            start_time = time.time()
            chain_obj = chain({"question": inp, "chat_history": history})
            print('=======time===== : ' + str(time.time() - start_time))

            output = chain_obj["answer"]

            history.append((inp, output))
            print(history)
            # print(chain_obj)
        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="",
    #         show_label=False,
    #         lines=1,
    #         type="password",
    #         value=""
    #     )

    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], outputs=[chatbot, state])
    message.submit(chat, inputs=[message, 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', 'Twimbit@2019'), auth_message='enter username password to proceed further')