import os
import logging

import gradio as gr
import anthropic
from dotenv import load_dotenv

load_dotenv(".env")

logging.basicConfig(level=logging.INFO)
logging.getLogger("gradio").setLevel(logging.INFO)
logging.getLogger("httpx").setLevel(logging.WARNING)


def generate_completion(
    input,
    history,
    api_key,
    model,
    system_prompt,
    temperature,
    max_tokens,
):
    if os.getenv("ANTHROPIC_API_KEY"):
        api_key = os.getenv("ANTHROPIC_API_KEY")

    if not api_key:
        # raise ValueError("API Key is required")
        yield "No API key provided"

    client = anthropic.Anthropic(
        api_key=api_key,
    )

    messages = []

    if history:
        for entry in history:
            if len(entry) == 2:
                messages.append(
                    {
                        "role": "user",
                        "content": entry[0],
                    }
                )
                messages.append(
                    {
                        "role": "assistant",
                        "content": entry[1],
                    }
                )

    messages.append(
        {
            "role": "user",
            "content": input,
        }
    )

    with client.messages.stream(
        model=model,
        max_tokens=max_tokens,
        temperature=temperature,
        system=system_prompt,
        messages=messages,
    ) as stream:
        answer_str = ""
        for text in stream.text_stream:
            # print(text, end="", flush=True)
            answer_str += text
            yield answer_str


api_key = gr.Textbox(label="API Key", type="password")
model = gr.Textbox(label="Model", value="claude-3-opus-20240229")
system_prompt = gr.Textbox(
    label="System Prompt",
    value="You are a world-class assistant.",
)
temperature = gr.Slider(label="Temperature", value=0.0, minimum=0.0, maximum=1.0)
max_tokens = gr.Slider(label="Max Tokens", value=4096, minimum=1, maximum=4096)


demo = gr.ChatInterface(
    fn=generate_completion,
    additional_inputs=[
        api_key,
        model,
        system_prompt,
        temperature,
        max_tokens,
    ],
    description="Claude Chatbot, Clone the space and add your own API key in the 'additional inputs' section. Get your key here: https://console.anthropic.com/login",
    fill_height=True,
)

if __name__ == "__main__":
    demo.queue()
    demo.launch()