import os
import cohere
import gradio as gr

try:
    import spaces

    @spaces.GPU
    def _zerogpu_dummy():
        """Satisfies ZeroGPU startup check if Space is assigned ZeroGPU hardware."""
        pass
except Exception:
    pass

SYSTEM_PROMPT = "You are Simplify, a friendly tutor who turns complex ideas into clear, simple explanations that anyone can understand."


def chat(message, history, user_api_key):
    effective_key = (user_api_key or "").strip() or os.environ.get("COHERE_API_KEY")
    if not effective_key:
        yield "⚠️ Please enter your Cohere API key in the configuration box below to start chatting.\n\n👉 [Get your free Cohere API key here](https://dashboard.cohere.com/api-keys)"
        return

    co = cohere.ClientV2(api_key=effective_key)
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for msg in history:
        messages.append({"role": msg["role"], "content": msg["content"]})
    messages.append({"role": "user", "content": message})

    try:
        response = co.chat_stream(
            model="command-a-plus-05-2026",
            messages=messages,
            max_tokens=1500,
            temperature=0.3,
        )
        text = ""
        for event in response:
            if event.type == "content-delta":
                chunk = getattr(event.delta.message.content, "text", None)
                if chunk:
                    text += chunk
                    yield text
    except Exception as exc:
        yield f"⚠️ An error occurred: {exc}"


api_key_input = gr.Textbox(
    label="Cohere API Key",
    placeholder="Paste your Cohere API key here...",
    type="password",
    info="Don't have an API key? Get one for free at https://dashboard.cohere.com/api-keys",
)

message_input = gr.Textbox(
    placeholder="Type a message...",
    show_label=False,
    autofocus=True,
    submit_btn=True,
    stop_btn=True,
    elem_id="chat-input",
)

demo = gr.ChatInterface(
    fn=chat,
    textbox=message_input,
    title="Simplify",
    description="Breaks down complex concepts into their simplest, most intuitive form. Need a key? [Get a free Cohere API key here](https://dashboard.cohere.com/api-keys).",
    examples=[
        ["Explain derivatives in the simplest way"],
        ["What is quantum physics?"],
        ["Explain integrals simply"],
    ],
    additional_inputs=[api_key_input],
    additional_inputs_accordion=gr.Accordion(label="🔑 API Key Configuration", open=True),
    cache_examples=False,
)

if __name__ == "__main__":
    demo.launch(
        js=r"""
    (() => {
        const selector = "#chat-input textarea";

        const focusWhenReady = () => {
            let attempts = 0;
            let sawDisabled = false;
            let timer;

            const tryFocus = () => {
                const input = document.querySelector(selector);
                if (!input) {
                    if (++attempts >= 300) clearInterval(timer);
                    return;
                }

                if (input.disabled) {
                    sawDisabled = true;
                }

                // Do not focus too early: Gradio disables the textbox during
                // submission and enables it again after the streamed reply.
                if (!input.disabled && (sawDisabled || attempts >= 5)) {
                    input.focus();
                    clearInterval(timer);
                } else if (++attempts >= 300) {
                    clearInterval(timer);
                }
            };

            // Let Gradio process the submit first. While the response streams,
            // the textbox is disabled; focus it as soon as it is enabled again.
            timer = setInterval(tryFocus, 100);
            setTimeout(tryFocus, 150);
        };

        document.addEventListener("click", (event) => {
            if (
                event.target instanceof Element &&
                event.target.closest("#chat-input [data-testid='submit-button']")
            ) {
                focusWhenReady();
            }
        }, true);

        document.addEventListener("keydown", (event) => {
            if (
                event.key === "Enter" &&
                !event.shiftKey &&
                event.target instanceof Element &&
                event.target.matches(selector)
            ) {
                focusWhenReady();
            }
        }, true);
    })();
    """,
    )