import re

import anthropic
import gradio as gr
from gradio.components import Checkbox, Dropdown, IOComponent, Textbox

from ai import AnthropicCustom
from const import ClaudeModels, ModelTokenLength, Prompts

# Define a global variable for the conversation history
conversation_history = ""


async def interact_with_ai(
    user_question, token, model, token_length, prompt, prompt_input, memory
):
    global conversation_history

    if memory:
        prompt = Prompts[prompt].value.format(
            memory=conversation_history, question=user_question
        )
    else:
        prompt = Prompts[prompt].value.format(memory="", question=user_question)

    if prompt_input != re.search(r"Human: (.*?) \n\nConversations:", prompt).group(1):
        prompt = re.sub(
            r"Human: (.*?) \n\nConversations:",
            f"Human: {prompt_input} \n\nConversations:",
            prompt,
        )

    # Create an instance of the custom class
    anth = AnthropicCustom(
        api_key=token, model=model, max_tokens=token_length, prompt=prompt
    )

    # Create a generator to stream the response
    response_accumulated = ""
    async for response in anth.get_anthropic_response_async():
        response_accumulated += response
        conversation_history = f"{conversation_history} {anthropic.HUMAN_PROMPT} {user_question} {anthropic.AI_PROMPT} {response_accumulated}"
        yield response_accumulated


async def chat_with_ai(
    message,
    history,
    token,
    model,
    token_length,
    prompt,
    prompt_input,
    memory,
):
    global conversation_history
    if memory:
        for conversation in history:
            user_question, response_accumulated = conversation
            conversation_history = f"{conversation_history} {anthropic.HUMAN_PROMPT} {user_question} {anthropic.AI_PROMPT} {response_accumulated}"
        prompt = Prompts[prompt].value.format(memory=history, question=message)
    else:
        prompt = Prompts[prompt].value.format(memory="", question=message)

    if prompt_input != re.search(r"Human: (.*?) \n\nConversations:", prompt).group(1):
        prompt = re.sub(
            r"Human: (.*?) \n\nConversations:",
            f"Human: {prompt_input} \n\nConversations:",
            prompt,
        )

    # Create an instance of the custom class
    anth = AnthropicCustom(
        api_key=token, model=model, max_tokens=token_length, prompt=prompt
    )

    # Create a generator to stream the response
    response_accumulated = ""
    async for response in anth.get_anthropic_response_async():
        response_accumulated += response
        yield response_accumulated


promptDropdown: IOComponent = Dropdown(
    choices=list(Prompts.__members__.keys()),
    label="Prompt",
    value=list(Prompts.__members__.keys())[0],
)
prompt_input: IOComponent = Textbox(
    label="Custom Prompt",
    placeholder="Enter a custom prompt here",
    lines=3,
    value=re.search(
        r"Human: (.*?) \n\nConversations:", Prompts[promptDropdown.value].value
    ).group(1),
)


iface = gr.Interface(
    fn=interact_with_ai,
    flagging_options=["Inappropriate", "Disrespectful", "Spam"],
    allow_flagging="auto",
    title="Claude Space",
    inputs=[
        Textbox(label="Question", placeholder="Enter a question here"),
        Textbox(label="Token", placeholder="Enter a token here", type="password"),
        Dropdown(
            choices=[model.value for model in ClaudeModels],
            label="Model",
            value=[model.value for model in ClaudeModels][0],
        ),
        Dropdown(
            choices=[token.value for token in ModelTokenLength],
            label="Token Length",
            value=[token.value for token in ModelTokenLength][0],
        ),
        promptDropdown,
        prompt_input,
        Checkbox(label="Memory", value=False),
    ],
    outputs="markdown",
    cache_examples=True,
)

promptDropdown: IOComponent = Dropdown(
    choices=list(Prompts.__members__.keys()),
    label="Prompt",
    value=list(Prompts.__members__.keys())[0],
)
prompt_input: IOComponent = Textbox(
    label="Custom Prompt",
    placeholder="Enter a custom prompt here",
    lines=3,
    value=re.search(
        r"Human: (.*?) \n\nConversations:", Prompts[promptDropdown.value].value
    ).group(1),
)

cface = gr.ChatInterface(
    fn=chat_with_ai,
    additional_inputs=[
        Textbox(label="Token", placeholder="Enter a token here", type="password"),
        Dropdown(
            choices=[model.value for model in ClaudeModels],
            label="Model",
            value=[model.value for model in ClaudeModels][0],
        ),
        Dropdown(
            choices=[token.value for token in ModelTokenLength],
            label="Token Length",
            value=[token.value for token in ModelTokenLength][0],
        ),
        promptDropdown,
        prompt_input,
        Checkbox(label="Memory", value=True),
    ],
)


if __name__ == "__main__":
    gd = gr.TabbedInterface(
        [iface, cface], tab_names=["Claude Space", "Claude Chat"], title="Claude Space"
    )
    gd.queue(concurrency_count=75, max_size=100).launch(
        debug=True, share=False, server_name="0.0.0.0", server_port=7864
    )