import os
from typing import Optional, Tuple

import gradio as gr
from langchain.agents import create_pandas_dataframe_agent
from langchain.llms import Anthropic
import pandas as pd
from threading import Lock

def load_data():
    """Load the data you want to use for the agent."""
    return pd.read_parquet("data/part.17.parquet")

def load_agent(trans: pd.DataFrame) -> create_pandas_dataframe_agent:
    """Logic for loading the agent we want to use."""
    llm = Anthropic(temperature=0)
    trans = load_data()
    return create_pandas_dataframe_agent(llm, 
                         trans, 
                         verbose=True)

def set_anthropic_api_key(api_key: str):
    """Set the api key and return agent.
    If no api_key, then None is returned.
    """
    if api_key:
        os.environ["ANTHROPIC_API_KEY"] = api_key
        agent = load_agent()
        os.environ["ANTHROPIC_API_KEY"] = ""
        return agent


class ChatWrapper:

    def __init__(self):
        self.lock = Lock()
    def __call__(
        self, api_key: str, inp: str, history: Optional[Tuple[str, str]], agent: Optional[create_pandas_dataframe_agent]
    ):
        """Execute the chat functionality."""
        self.lock.acquire()
        try:
            history = history or []
            # If chain is None, that is because no API key was provided.
            if agent is None:
                history.append((inp, "Please paste your anthropic key to use"))
                return history, history
            # Set OpenAI key
            agent = set_anthropic_api_key(history)
            # Run chain and append input.
            output = agent.run(input=inp)
            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: lightgray}")

with block:
    with gr.Row():
        gr.Markdown("<h3><center>LangChain Demo</center></h3>")

        anthropic_api_key_textbox = gr.Textbox(
            placeholder="Paste your Anthropic 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="What was the average gas price on 2019-01-22?",
            lines=1,
        )
        submit = gr.Button(value="Send", variant="secondary").style(full_width=False)

    gr.Examples(
        examples=[
            "Which to_address spent the most gas?",
            "What was the average gas price on 2019-01-22?",
            "How many unique addresses were sending transactions on 2019-01-22?",
        ],
        inputs=message,
    )

    gr.HTML("Demo application of a LangChain agent.")

    gr.HTML(
        "<center>Powered by <a href='https://github.com/hwchase17/langchain'>LangChain 🦜️🔗</a></center>"
    )

    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])

    anthropic_api_key_textbox.change(
        set_anthropic_api_key,
        inputs=[anthropic_api_key_textbox],
        outputs=[agent_state],
    )

block.launch(debug=True)