"""
ThinkingCap-Qwen3.6-27B — Hugging Face Space
============================================

Custom chat frontend (vanilla HTML/CSS/JS in index.html) talking to a
Gradio backend built with `gradio.Server`.

`gradio.Server` extends FastAPI. We expose one queued, ZeroGPU-backed
endpoint, `generate`, via `@app.api()` so it goes through Gradio's queue /
concurrency engine and is callable from the Gradio JS Client. A plain
`@app.get("/")` serves the static index.html.

Run locally:  python app.py
On Spaces:    select ZeroGPU hardware + xlarge GPU (96 GB) in Space settings
             (see README). The 27B bf16 model needs the larger VRAM tier.
"""

import os
import re
import sys
import threading
from typing import Iterator, Optional

import spaces
import torch
from PIL import Image
from fastapi.responses import HTMLResponse
from gradio import Server
from gradio.data_classes import FileData
from transformers import AutoModelForImageTextToText, AutoProcessor
from transformers import TextIteratorStreamer

MODEL_ID = os.environ.get("MODEL_ID", "bottlecapai/ThinkingCap-Qwen3.6-27B")

# Defaults mirror the model card's evaluation settings.
DEFAULT_MAX_NEW_TOKENS = 32768
DEFAULT_TEMPERATURE = 1.0
DEFAULT_TOP_P = 0.95
DEFAULT_TOP_K = 20

# Qwen3 reasoning-trace tags + EOS markers. Built by concatenation so the
# literal markers never appear together in source.
_THINK_OPEN = "<" + "think" + ">"
_THINK_CLOSE = "<" + "/think" + ">"
_THINK_RE = re.compile(
    re.escape(_THINK_OPEN) + r"(.*?)" + re.escape(_THINK_CLOSE), re.DOTALL
)
_IM_END = "<" + "|im_end|" + ">"
_EOT = "<" + "|endoftext|" + ">"

# ---------------------------------------------------------------------------
# Model load (root level so ZeroGPU's CUDA emulation is active during startup;
# real CUDA is used inside @spaces.GPU).
# ---------------------------------------------------------------------------
print(f"[app] loading {MODEL_ID} ...", flush=True)
try:
    processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
    model = AutoModelForImageTextToText.from_pretrained(
        MODEL_ID,
        dtype=torch.bfloat16,
        trust_remote_code=True,
    )
    model.to("cuda").eval()
    print("[app] model ready.", flush=True)
except Exception as e:  # pragma: no cover - surface clearly on Spaces logs
    print(f"[app] FAILED to load model: {e}", file=sys.stderr, flush=True)
    raise


# ---------------------------------------------------------------------------
# Backend
# ---------------------------------------------------------------------------
app = Server()


def _load_image(image: Optional[FileData]) -> Optional[Image.Image]:
    """Accept a Gradio FileData (dict-like or object) and return a PIL image."""
    if image is None:
        return None
    path = image["path"] if isinstance(image, dict) else image.path
    if not path or not os.path.exists(path):
        return None
    return Image.open(path).convert("RGB")


def _build_messages(prompt: str, image: Optional[Image.Image]) -> list:
    """Build a Qwen-VL style chat message list with an optional image."""
    if image is None:
        return [{"role": "user", "content": prompt}]
    return [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": prompt},
            ],
        }
    ]


def _split_thinking(text: str) -> tuple[str, str]:
    """Split a Qwen3 reasoning trace into (thinking, answer)."""
    matches = list(_THINK_RE.finditer(text))
    if not matches:
        return "", text.strip()
    last = matches[-1]
    thinking = last.group(1).strip()
    answer = text[last.end():].strip()
    if not answer:
        answer = text[: matches[0].start()].strip()
    return thinking, answer


def _tidy(text: str) -> str:
    """Drop leftover special tokens (eos) and an unclosed think block."""
    text = text.replace(_IM_END, "").replace(_EOT, "").strip()
    if _THINK_OPEN in text and _THINK_CLOSE not in text:
        text = text.split(_THINK_OPEN, 1)[0].strip()
    return text


@app.api()
@spaces.GPU(size="xlarge", duration=300)
def generate(
    prompt: str,
    image: Optional[FileData] = None,
    history: Optional[list] = None,
    max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
    temperature: float = DEFAULT_TEMPERATURE,
    top_p: float = DEFAULT_TOP_P,
    top_k: int = DEFAULT_TOP_K,
    enable_thinking: bool = True,
) -> Iterator[dict]:
    """Stream a response from ThinkingCap-Qwen3.6-27B.

    Yields ``{"text": <full text so far>}`` dicts as tokens are produced
    (so the frontend can render the reasoning trace live), then a final dict
    with parsed ``thinking`` / ``answer`` and a thinking-token count.
    Parameters mirror the model card's recommended sampling.
    """
    history = history or []
    image_pil = _load_image(image)

    messages: list = []
    for turn in history:
        role = turn.get("role", "user")
        content = turn.get("content", "")
        if role == "assistant":
            _, ans = _split_thinking(content)
            messages.append({"role": "assistant", "content": ans or content})
        else:
            messages.append({"role": "user", "content": content})
    messages.append(_build_messages(prompt, image_pil)[0])

    try:
        text = processor.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
            enable_thinking=enable_thinking,
        )
    except TypeError:
        text = processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )

    inputs = processor(
        text=[text],
        images=[image_pil] if image_pil is not None else None,
        return_tensors="pt",
    ).to(model.device)

    streamer = TextIteratorStreamer(
        processor.tokenizer,
        skip_prompt=True,
        skip_special_tokens=False,
        timeout=60.0,
    )

    gen_kwargs = dict(
        **inputs,
        max_new_tokens=int(max_new_tokens),
        do_sample=float(temperature) > 0,
        temperature=max(float(temperature), 1e-5),
        top_p=float(top_p),
        top_k=int(top_k),
        streamer=streamer,
    )

    thread = threading.Thread(target=lambda: model.generate(**gen_kwargs))
    thread.start()

    accumulated = ""
    for piece in streamer:
        accumulated += piece
        # Echo the full text-so-far so the client can render the live trace.
        yield {"text": accumulated}

    thread.join()

    full = _tidy(accumulated)
    thinking, answer = _split_thinking(full)
    thinking_tokens = len(processor.tokenizer.encode(thinking)) if thinking else 0

    yield {
        "text": full,
        "thinking": thinking,
        "answer": answer,
        "thinking_tokens": thinking_tokens,
    }


# ---------------------------------------------------------------------------
# Static frontend
# ---------------------------------------------------------------------------
_STATIC_DIR = os.path.dirname(os.path.abspath(__file__))


@app.get("/", response_class=HTMLResponse)
async def homepage():
    html_path = os.path.join(_STATIC_DIR, "index.html")
    with open(html_path, "r", encoding="utf-8") as f:
        return HTMLResponse(content=f.read())


@app.get("/health")
async def health():
    return {"ok": True, "model": MODEL_ID}


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