"""A secure Gradio demo for an OpenAI-compatible Qwen3.8-Flash-Next endpoint."""

import spaces  # ZeroGPU requires one decorated function, even for an API proxy.

import base64
import mimetypes
import os
from pathlib import Path
from typing import Any

import gradio as gr
from openai import OpenAI


MODEL_ID = os.getenv("MODEL_ID", "Qwen/Qwen3.8-Flash-Next")
API_KEY = os.getenv("OPENAI_API_KEY")
BASE_URL = os.getenv("OPENAI_BASE_URL")
MAX_IMAGE_BYTES = 10 * 1024 * 1024


@spaces.GPU(duration=1)
def _zerogpu_marker() -> None:
    """Declare ZeroGPU support without allocating a GPU for API-proxy requests."""
    return None


def image_part(image_path: str | None) -> dict[str, Any] | None:
    """Return a local image as a bounded data URL for the vision-capable model."""
    if not image_path:
        return None

    path = Path(image_path)
    if not path.is_file():
        raise ValueError("The uploaded image could not be read. Please try again.")
    if path.stat().st_size > MAX_IMAGE_BYTES:
        raise ValueError("Images must be 10 MB or smaller.")

    media_type = mimetypes.guess_type(path.name)[0] or "application/octet-stream"
    if not media_type.startswith("image/"):
        raise ValueError("Please upload an image file.")

    encoded = base64.b64encode(path.read_bytes()).decode("ascii")
    return {"type": "image_url", "image_url": {"url": f"data:{media_type};base64,{encoded}"}}


def api_messages(history: list[tuple[str, str | None]]) -> list[dict[str, Any]]:
    """Convert Gradio's tuple-based history into Chat Completions messages."""
    messages: list[dict[str, Any]] = []
    for user_content, assistant_content in history:
        if isinstance(user_content, str):
            messages.append({"role": "user", "content": user_content})
        if isinstance(assistant_content, str):
            messages.append({"role": "assistant", "content": assistant_content})
    return messages


def respond(
    message: str,
    image_path: str | None,
    history: list[tuple[str, str | None]],
    enable_thinking: bool,
    reasoning_effort: str,
):
    """Answer a question with Qwen3.8-Flash-Next, optionally using an uploaded image."""
    text = (message or "").strip()
    if not text and not image_path:
        yield history, "Add a question or an image before sending."
        return
    if not API_KEY or not BASE_URL:
        yield history, (
            "This Space needs the OPENAI_API_KEY and OPENAI_BASE_URL secrets. "
            "See the README setup instructions."
        )
        return

    try:
        content: list[dict[str, Any]] = []
        part = image_part(image_path)
        if part:
            content.append(part)
        if text:
            content.append({"type": "text", "text": text})

        user_message: dict[str, Any] = {"role": "user", "content": content}
        conversation = api_messages(history)
        conversation.append(user_message)
        display_text = text or "Please analyse this image."
        updated_history = [*history]

        client = OpenAI(api_key=API_KEY, base_url=BASE_URL, timeout=180)
        stream = client.chat.completions.create(
            model=MODEL_ID,
            messages=conversation,
            stream=True,
            reasoning_effort=reasoning_effort,
            extra_body={
                "chat_template_kwargs": {
                    "enable_thinking": enable_thinking,
                    "preserve_thinking": False,
                }
            },
        )

        answer = ""
        for chunk in stream:
            if not chunk.choices:
                continue
            delta = chunk.choices[0].delta
            # Deliberately do not display internal reasoning traces. Only final output is shown.
            if delta.content:
                answer += delta.content
                yield [*updated_history, (display_text, answer)], ""

        if not answer:
            answer = "The endpoint returned no final response. Try again or check its logs."
        yield [*updated_history, (display_text, answer)], ""
    except ValueError as error:
        yield history, str(error)
    except Exception:
        yield history, "The model endpoint could not complete this request. Please try again shortly."


with gr.Blocks(title="Qwen3.8-Flash-Next Demo") as demo:
    gr.Markdown(
        "# Qwen3.8-Flash-Next\n"
        "A private-key, OpenAI-compatible chat demo for the Qwen vision-language checkpoint. "
        "Your API credentials remain in Space secrets."
    )
    with gr.Row():
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(label="Conversation", height=520)
            with gr.Row():
                message = gr.Textbox(
                    label="Your message",
                    placeholder="Ask a question, request code, or describe the attached image…",
                    lines=2,
                    autofocus=True,
                )
                image = gr.Image(label="Optional image", type="filepath", sources=["upload"])
            with gr.Row():
                send = gr.Button("Send", variant="primary")
                clear = gr.Button("Clear conversation")
        with gr.Column(scale=1):
            gr.Markdown("### Generation settings")
            enable_thinking = gr.Checkbox(
                label="Enable model thinking",
                value=True,
                info="The model reasons internally; only its final answer is shown.",
            )
            reasoning_effort = gr.Radio(
                choices=["low", "medium", "xhigh"],
                value="xhigh",
                label="Reasoning effort",
                info="Higher effort can improve complex answers but takes longer.",
            )
            status = gr.Textbox(label="Status", interactive=False, lines=2)
            gr.Markdown(
                "**Model:** `Qwen/Qwen3.8-Flash-Next`  \n"
                "**Inputs:** text and one image (up to 10 MB)"
            )
            gr.Examples(
                examples=[
                    ["Write a concise Python function that merges two sorted linked lists."],
                    ["Explain the practical difference between sparse attention and full attention."],
                    ["Summarise this image and list details that could be important to a researcher."],
                ],
                inputs=message,
                label="Try an example",
            )

    inputs = [message, image, chatbot, enable_thinking, reasoning_effort]
    send.click(respond, inputs=inputs, outputs=[chatbot, status])
    message.submit(respond, inputs=inputs, outputs=[chatbot, status])
    clear.click(lambda: ([], ""), outputs=[chatbot, status])


if __name__ == "__main__":
    demo.queue().launch()