import os
import gradio as gr
from cactus.agent import Cactus
from huggingface_hub import InferenceClient


MODEL_API_KEY_MAP = {
    "gpt-3.5-turbo": "OPENAI_API_KEY",
    "gpt-4": "OPENAI_API_KEY",
    "claude-3-haiku-20240307": "ANTHROPIC_API_KEY",
    "claude-3-opus-20240229": "ANTHROPIC_API_KEY",
    "claude-3-sonnet-20240229": "ANTHROPIC_API_KEY",
}


def agent_fn(prompt, api_key=None):
    env_key_name = MODEL_API_KEY_MAP.get(model_name)
    if env_key_name:
        if api_key:
            os.environ[env_key_name] = api_key
        else:
            return f"Error: API key not found for {model_name}. Please provide it."
    try:
        cactus_model = Cactus(model_name, model_type="api")
        response = cactus_model.run(prompt)
        return response
    except Exception as e:
        return f"Error: {e}"


with gr.Blocks() as iface:
    with gr.Row():
        model_name_dropdown = gr.Dropdown(
            list(MODEL_API_KEY_MAP.keys()), label="Select Model"
        )
        api_key_input = gr.Textbox(label="Enter API Key", type="password")
        prompt_input = gr.Textbox(
            lines=2, placeholder="Enter your cheminformatics question..."
        )

    response_output = gr.Textbox(label="Agent Response")

    # Event Triggers
    model_name_dropdown.change(
        lambda x: gr.update(visible=MODEL_API_KEY_MAP.get(x) is not None),
        model_name_dropdown,
        api_key_input,
    )
    api_key_input.change(
        agent_fn, [prompt_input, model_name_dropdown, api_key_input], response_output
    )
    prompt_input.change(
        agent_fn, [prompt_input, model_name_dropdown, api_key_input], response_output
    )

iface.launch()