"""Hy3 chat app on gradio.Server.

Drop-in Gradio backend for the OpenRouter Quick Start for `tencent/hy3:free`:
  - OpenAI-compatible API via OpenRouter
  - reasoning enabled with `extra_body={"reasoning": {"enabled": True}}`
  - streaming responses
  - reasoning_details preserved across turns so the model can keep thinking

Run:
    export OPENROUTER_API_KEY=sk-or-v1-...
    python app.py
Then open http://127.0.0.1:7860
"""

import json
import os
from typing import Any, Generator

from fastapi.responses import HTMLResponse
from openai import OpenAI

from gradio import Server

OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY", "")

client = OpenAI(
    base_url="https://openrouter.ai/api/v1",
    api_key=OPENROUTER_API_KEY,
)

app = Server()


def _delta_fields(delta: Any) -> dict:
    """Return all known + OpenRouter-extra fields on a streaming delta.

    The OpenAI SDK keeps unknown fields (OpenRouter's `reasoning`,
    `reasoning_content`) in Pydantic extra storage, so we merge `model_dump`
    with `__pydantic_extra__` to be robust across SDK versions.
    """
    fields: dict = {}
    if hasattr(delta, "model_dump"):
        try:
            fields.update(delta.model_dump())
        except Exception:
            pass
    extra = getattr(delta, "__pydantic_extra__", None) or {}
    fields.update(extra)
    return fields


def _reasoning_text(fields: dict) -> str:
    """Pull a reasoning chunk from whatever field OpenRouter used."""
    r = fields.get("reasoning") or fields.get("reasoning_content")
    if r:
        return r
    # Rare: some flows surface reasoning as a details array.
    details = fields.get("reasoning_details")
    if isinstance(details, list):
        return "".join(
            d.get("content")
            or (d.get("summary", [{}])[0].get("text") if d.get("summary") else "")
            or ""
            for d in details
            if isinstance(d, dict)
        )
    return ""


def _stream_chat(
    messages: list,
    model: str,
    temperature: float,
    max_tokens: int,
) -> Generator[str, None, None]:
    """Stream a chat completion from OpenRouter, yielding JSON events.

    Events (one JSON string per yielded value):
      {"type": "reasoning", "text": "..."}   live thinking tokens
      {"type": "content",   "text": "..."}   live answer tokens
      {"type": "error",     "text": "..."}   something went wrong
      {"type": "done", "content": "...", "reasoning": "...",
       "reasoning_details": [...]}            final, with continuation payload
    """
    if not OPENROUTER_API_KEY:
        yield json.dumps(
            {
                "type": "error",
                "text": (
                    "OPENROUTER_API_KEY is not set on the server. "
                    "Run: export OPENROUTER_API_KEY=sk-or-v1-...  then restart."
                ),
            }
        )
        return

    # Rebuild the message list for OpenRouter. We pass the assistant
    # reasoning_details back through unmodified so the model can continue
    # reasoning from where it left off (per the OpenRouter quick start).
    api_messages: list[dict] = []
    for m in messages or []:
        msg = {"role": m.get("role", "user"), "content": m.get("content", "") or ""}
        if m.get("reasoning_details"):
            msg["reasoning_details"] = m["reasoning_details"]
        api_messages.append(msg)

    reasoning_buf: list[str] = []
    content_buf: list[str] = []

    try:
        stream = client.chat.completions.create(
            model=model,
            messages=api_messages,
            stream=True,
            temperature=temperature,
            max_tokens=max_tokens,
            extra_body={"reasoning": {"enabled": True}},
        )
    except Exception as e:  # noqa: BLE001 - surface any error to the UI
        yield json.dumps({"type": "error", "text": f"Request failed: {e}"})
        return

    try:
        for chunk in stream:
            choices = getattr(chunk, "choices", None) or []
            if not choices:
                continue
            delta = choices[0].delta
            if delta is None:
                continue
            fields = _delta_fields(delta)

            r = _reasoning_text(fields)
            if r:
                reasoning_buf.append(r)
                yield json.dumps({"type": "reasoning", "text": r})

            c = fields.get("content")
            if c:
                content_buf.append(c)
                yield json.dumps({"type": "content", "text": c})
    except Exception as e:  # noqa: BLE001
        yield json.dumps({"type": "error", "text": f"Stream interrupted: {e}"})
        return

    reasoning_full = "".join(reasoning_buf)
    content_full = "".join(content_buf)

    # Reconstruct reasoning_details for multi-turn continuation. The streamed
    # response only gives us reasoning text; we wrap it in the summary shape so
    # the next turn can carry it forward like the non-streaming quick start.
    reasoning_details = (
        [{"type": "reasoning", "summary": [{"type": "summary_text", "text": reasoning_full}]}]
        if reasoning_full
        else None
    )

    yield json.dumps(
        {
            "type": "done",
            "content": content_full,
            "reasoning": reasoning_full,
            "reasoning_details": reasoning_details,
        }
    )


@app.api()
def chat(
    messages: list,
    model: str = "tencent/hy3:free",
    temperature: float = 0.9,
    max_tokens: int = 2048,
) -> Generator[str, None, None]:
    """Hy3 chat with reasoning.

    Pass the full message history; assistant turns should include their
    `reasoning_details` so the model continues reasoning across turns.
    Yields JSON event strings (see _stream_chat).
    """
    for event in _stream_chat(messages, model, temperature, max_tokens):
        yield event


@app.get("/", response_class=HTMLResponse)
async def root() -> HTMLResponse:
    """Serve the custom chat frontend (vanilla HTML/CSS/JS)."""
    html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")
    with open(html_path, "r", encoding="utf-8") as f:
        return HTMLResponse(content=f.read())


if __name__ == "__main__":
    app.launch(show_error=True)