from __future__ import annotations

from collections.abc import Generator

import gradio as gr
from huggingface_hub import InferenceClient

MODEL_ID = "MarkChenX/lfm2-quantum-128m-sft"


def respond(
    message: str,
    history: list[dict[str, str]],
    system_message: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
    hf_token: gr.OAuthToken | None,
) -> Generator[str, None, None]:
    """Generate and stream a response from the Hugging Face model."""

    if hf_token is None:
        yield "Please sign in with Hugging Face before sending a message."
        return

    messages = [
        {
            "role": "system",
            "content": system_message.strip() or "You are a friendly chatbot.",
        },
        *history,
        {"role": "user", "content": message},
    ]

    client = InferenceClient(
        model=MODEL_ID,
        token=hf_token.token,
    )

    response = ""

    try:
        for chunk in client.chat_completion(
            messages=messages,
            max_tokens=int(max_tokens),
            temperature=float(temperature),
            top_p=float(top_p),
            stream=True,
        ):
            if not chunk.choices:
                continue

            token = chunk.choices[0].delta.content or ""
            response += token
            yield response

    except Exception as exc:
        # Keep the existing partial response visible if generation fails.
        error_message = f"\n\nGeneration failed: {exc}"
        yield response + error_message


with gr.Blocks(title="LFM2 Quantum Chat") as demo:
    with gr.Sidebar():
        gr.Markdown("### Authentication")
        gr.LoginButton()

    gr.ChatInterface(
        fn=respond,
        additional_inputs=[
            gr.Textbox(
                value="You are a friendly chatbot.",
                label="System message",
                lines=3,
            ),
            gr.Slider(
                minimum=1,
                maximum=2048,
                value=512,
                step=1,
                label="Max new tokens",
            ),
            gr.Slider(
                minimum=0.1,
                maximum=4.0,
                value=0.7,
                step=0.1,
                label="Temperature",
            ),
            gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.95,
                step=0.05,
                label="Top-p",
            ),
        ],
    )


if __name__ == "__main__":
    # Spaces already exposes the application publicly.
    # Disabling SSR avoids launching the experimental Node.js SSR process.
    demo.launch(
        share=False,
        ssr_mode=False,
    )