"""
Ornith-1.0-9B — a gradio.Server chat app.

This Space demonstrates `deepreinforce-ai/Ornith-1.0-9B`, an agentic coding
(qwen3.5-based) reasoning model, served through `gradio.Server`.

`gradio.Server` extends FastAPI and adds Gradio's API engine on top
(queuing, concurrency control, SSE streaming, gradio_client compatibility).
Instead of Gradio's component UI we ship a custom HTML/CSS/JS chat page
(see `index.html`) served at `GET /`, while the model inference is exposed
as a streaming `@app.api()` endpoint at `/gradio_api/call/generate`.

Run locally (needs a CUDA GPU + transformers >= 5.8.1):
    python app.py
On a Hugging Face ZeroGPU Space, `@spaces.GPU` handles device allocation.
"""

import json
import os
from threading import Thread
from typing import Optional

import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TextIteratorStreamer,
)

from fastapi.responses import HTMLResponse
from gradio import Server

# `spaces` is only available on Hugging Face Spaces. Guard the import so the
# file still parses/imports on a plain machine (the GPU decorator becomes a
# no-op locally).
try:
    import spaces  # type: ignore

    GPU = spaces.GPU
except ImportError:  # local, non-Spaces environment
    def GPU(*args, **kwargs):
        def decorator(fn):
            return fn

        return decorator


MODEL_ID = "deepreinforce-ai/Ornith-1.0-9B"

# Recommended sampling for Ornith (see model card): temp=0.6, top_p=0.95,
# top_k=20. temp=1.0 reproduces the reported benchmark setup.
DEFAULT_TEMPERATURE = 0.6
DEFAULT_MAX_NEW_TOKENS = 2048

# Ornith is a reasoning model: the assistant turn opens with a chain-of-
# thought block before the final answer. We split the two so the UI can
# render the reasoning separately from the answer. The delimiters are read
# from the tokenizer config at runtime (Qwen3-style reasoning markers, as
# Ornith ships a qwen3 reasoning parser) rather than hardcoded, so the
# raw special-token bytes never need to live in source.
THINK_OPEN = None   # resolved lazily from tokenizer config in _split_think
THINK_CLOSE = None  # resolved lazily from tokenizer config in _split_think


# --------------------------------------------------------------------------- #

# --------------------------------------------------------------------------- #
# Model loading: at module level on cuda (per ZeroGPU docs)
# --------------------------------------------------------------------------- #
# ZeroGPU runs in PyTorch CUDA-emulation mode outside @spaces.GPU functions,
# so a real GPU is unavailable at import time. We load the model at startup
# and call .to("cuda") once - the CUDA emulation handles that safely. Real
# CUDA then becomes available inside the @spaces.GPU-decorated function.
print(f"[ornith] loading {MODEL_ID} ...", flush=True)
_tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
_model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True,
)
_model.to("cuda")
_model.eval()
print("[ornith] model ready on cuda", flush=True)


def _resolve_markers():
    open_tok = None
    close_tok = None
    try:
        added = getattr(_tokenizer, "added_tokens_decoder", {}) or {}
        for entry in added.values():
            content = getattr(entry, "content", None)
            special = getattr(entry, "special", False)
            if not (special and content):
                continue
            low = content.lower()
            if "think" in low and not content.startswith("</"):
                open_tok = content
            elif "think" in low and content.startswith("</"):
                close_tok = content
    except Exception:
        pass
    if not open_tok or not close_tok:
        # Qwen3 defaults (resolved from char codes so they survive this toolchain).
        open_tok = open_tok or "".join(chr(c) for c in (60, 116, 104, 105, 110, 107, 62))
        close_tok = close_tok or "".join(chr(c) for c in (60, 47, 116, 104, 105, 110, 107, 62))
    return open_tok, close_tok


def _split_think(text, open_tok, close_tok):
    if close_tok and close_tok in text:
        head, _, tail = text.partition(close_tok)
        reasoning = head.replace(open_tok or "", "").strip()
        return reasoning, tail.strip()
    return text.replace(open_tok or "", "").strip(), ""


# --------------------------------------------------------------------------- #
# App + streaming endpoint
# --------------------------------------------------------------------------- #
app = Server()


@app.api(name="generate", concurrency_limit=1, stream_every=0.5)
@GPU(duration=300)
def generate(
    message: str,
    history: list = None,
    max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
    temperature: float = DEFAULT_TEMPERATURE,
) -> str:
    # Stream an Ornith-1.0-9B completion; yields JSON strings with
    # {reasoning, answer, status, error}.
    if history is None:
        history = []
    open_tok, close_tok = _resolve_markers()

    try:
        messages = []
        for turn in history:
            role = turn.get("role")
            content = turn.get("content")
            if role in ("system", "user", "assistant") and content:
                messages.append({"role": role, "content": content})
        messages.append({"role": "user", "content": message or ""})

        prompt = _tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
        )
        inputs = _tokenizer(prompt, return_tensors="pt").to("cuda")

        streamer = TextIteratorStreamer(
            _tokenizer,
            skip_prompt=True,
            skip_special_tokens=True,
            timeout=120.0,
        )

        pad_token_id = getattr(_tokenizer, "eos_token_id", None)
        gen_kwargs = dict(
            **inputs,
            streamer=streamer,
            max_new_tokens=int(max_new_tokens),
            do_sample=True,
            temperature=float(temperature),
            top_p=0.95,
            top_k=20,
            pad_token_id=pad_token_id,
        )

        thread = Thread(target=_model.generate, kwargs=gen_kwargs)
        thread.start()

        accumulated = ""
        for text in streamer:
            accumulated += text
            reasoning, answer = _split_think(accumulated, open_tok, close_tok)
            yield json.dumps(
                {
                    "reasoning": reasoning,
                    "answer": answer,
                    "status": "generating",
                    "error": None,
                }
            )

        thread.join()
        reasoning, answer = _split_think(accumulated, open_tok, close_tok)
        yield json.dumps(
            {
                "reasoning": reasoning,
                "answer": answer,
                "status": "complete",
                "error": None,
            }
        )
    except Exception as exc:  # noqa: BLE001 - surface any failure to the UI
        yield json.dumps(
            {
                "reasoning": "",
                "answer": "",
                "status": "error",
                "error": str(exc),
            }
        )


@app.get("/", response_class=HTMLResponse)
async def homepage():
    """Serve the custom chat frontend; this overrides Gradio's default UI."""
    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 f.read()


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