import os
from collections.abc import Generator
from typing import Any

import gradio as gr
from dotenv import load_dotenv
from groq import Groq

# Hugging Face ZeroGPU validates that a Space contains at least one
# @spaces.GPU-decorated function. This app calls Groq remotely and does not
# perform local GPU inference, so the compatibility function below is never
# connected to the UI and never consumes GPU time.
try:
    import spaces
except ImportError:
    # Keep local execution working outside Hugging Face Spaces.
    class _SpacesFallback:
        @staticmethod
        def GPU(*args: Any, **kwargs: Any):
            def decorator(function):
                return function

            return decorator

    spaces = _SpacesFallback()

# Loads GROQ_API_KEY from a local .env file when running on your computer.
# On Hugging Face Spaces, add GROQ_API_KEY under Settings > Secrets.
load_dotenv()

DEFAULT_SYSTEM_PROMPT = (
    "You are a helpful, accurate, and friendly AI assistant. "
    "Answer clearly, use Markdown when helpful, and admit uncertainty when needed."
)

MODEL_CHOICES = [
    "openai/gpt-oss-20b",
    "openai/gpt-oss-120b",
]

CUSTOM_CSS = """
.gradio-container {
    max-width: 1050px !important;
    margin: 0 auto !important;
}
#app-header {
    text-align: center;
    padding: 12px 0 4px 0;
}
#app-subtitle {
    text-align: center;
    opacity: 0.8;
    margin-bottom: 10px;
}
footer {
    display: none !important;
}
"""


def _text_history(history: list[dict[str, Any]]) -> list[dict[str, str]]:
    """Keep only plain-text user and assistant messages for the Groq API."""
    cleaned: list[dict[str, str]] = []

    for item in history or []:
        role = item.get("role")
        content = item.get("content")

        if role in {"user", "assistant"} and isinstance(content, str):
            cleaned.append({"role": role, "content": content})

    return cleaned


@spaces.GPU(duration=1)
def zero_gpu_compatibility_check() -> str:
    """Allow startup on ZeroGPU; the Groq chat itself remains CPU/API based."""
    return "ZeroGPU compatibility ready"


def chat_with_groq(
    message: str,
    history: list[dict[str, Any]],
    model: str,
    system_prompt: str,
    temperature: float,
    max_tokens: int,
) -> Generator[str, None, None]:
    """Stream a Groq response to the Gradio chat interface."""
    api_key = os.getenv("GROQ_API_KEY")

    if not api_key:
        yield (
            "### Missing API key\n\n"
            "Add a Hugging Face Space secret named `GROQ_API_KEY`, then restart the Space."
        )
        return

    user_message = (message or "").strip()
    if not user_message:
        yield "Please enter a message."
        return

    selected_model = model if model in MODEL_CHOICES else MODEL_CHOICES[0]
    instructions = (system_prompt or "").strip() or DEFAULT_SYSTEM_PROMPT

    messages: list[dict[str, str]] = [
        {"role": "system", "content": instructions},
        *_text_history(history),
        {"role": "user", "content": user_message},
    ]

    try:
        client = Groq(api_key=api_key)
        stream = client.chat.completions.create(
            model=selected_model,
            messages=messages,
            temperature=float(temperature),
            max_completion_tokens=int(max_tokens),
            stream=True,
        )

        response = ""
        for chunk in stream:
            delta = chunk.choices[0].delta.content
            if delta:
                response += delta
                yield response

        if not response:
            yield "The model returned an empty response. Please try again."

    except Exception as error:
        # Show a useful message without exposing the API key or other secrets.
        error_name = type(error).__name__
        yield (
            "### Request failed\n\n"
            f"`{error_name}`: {error}\n\n"
            "Check your Groq API key, model access, account limits, and network connection."
        )


with gr.Blocks(title="Groq AI Chat") as demo:
    # ZeroGPU scans Gradio's registered event handlers during startup.
    # This hidden event is never invoked by the chat UI, so Groq requests do
    # not reserve or consume a GPU allocation.
    zero_gpu_trigger = gr.Button(visible=False)
    zero_gpu_status = gr.Textbox(visible=False)
    zero_gpu_trigger.click(
        fn=zero_gpu_compatibility_check,
        inputs=None,
        outputs=zero_gpu_status,
        api_visibility="private",
    )

    gr.Markdown("# ⚡ Groq AI Chat", elem_id="app-header")
    gr.Markdown(
        "A fast, streaming chatbot powered by Groq and built with Gradio.",
        elem_id="app-subtitle",
    )

    with gr.Accordion("Chat settings", open=False):
        model_input = gr.Dropdown(
            choices=MODEL_CHOICES,
            value=MODEL_CHOICES[0],
            label="Groq model",
            info="GPT-OSS 20B is faster; GPT-OSS 120B is stronger for complex tasks.",
        )
        system_prompt_input = gr.Textbox(
            value=DEFAULT_SYSTEM_PROMPT,
            label="System prompt",
            lines=3,
        )
        with gr.Row():
            temperature_input = gr.Slider(
                minimum=0.0,
                maximum=2.0,
                value=0.7,
                step=0.1,
                label="Temperature",
            )
            max_tokens_input = gr.Slider(
                minimum=128,
                maximum=4096,
                value=1024,
                step=128,
                label="Maximum response tokens",
            )

    chatbot = gr.Chatbot(
        label="Conversation",
        placeholder="Ask anything to begin the conversation.",
        height=520,
    )

    gr.ChatInterface(
        fn=chat_with_groq,
        chatbot=chatbot,
        additional_inputs=[
            model_input,
            system_prompt_input,
            temperature_input,
            max_tokens_input,
        ],
        examples=[
            [
                "Explain artificial intelligence in simple words.",
                MODEL_CHOICES[0],
                DEFAULT_SYSTEM_PROMPT,
                0.7,
                1024,
            ],
            [
                "Write a professional email requesting a meeting.",
                MODEL_CHOICES[0],
                DEFAULT_SYSTEM_PROMPT,
                0.7,
                1024,
            ],
            [
                "Create a beginner-friendly Python learning plan.",
                MODEL_CHOICES[1],
                DEFAULT_SYSTEM_PROMPT,
                0.5,
                1536,
            ],
        ],
        editable=True,
        save_history=True,
        flagging_mode="never",
        api_visibility="private",
        concurrency_limit=5,
        fill_height=True,
        fill_width=True,
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=5).launch(
        theme=gr.themes.Soft(),
        css=CUSTOM_CSS,
    )