"""Model loading and the generation loop of the demo (independent of Gradio)."""

from __future__ import annotations

import os
import time
from dataclasses import dataclass
from typing import Iterator

from .render import Detokenizer, replay_html, stats_html, tokens_html

# display name -> Hub repo with the heads (base model and tokenizer are pinned in its config)
MODELS = {
    "Qwen3-8B-Base": "FlorianVal/DSSD-Qwen3-8B-Base",
    "Llama-3.1-8B": "FlorianVal/Llama-3.1-8B-DSSD",
}
MAX_PROMPT_TOKENS = 1536
MAX_NEW_TOKENS = 256


def _graphs() -> bool | str:
    # CUDA graphs are captured once per process; under ZeroGPU each call may run in a fresh GPU
    # worker, so the eager engine is the safe default there. DSSD_GRAPHS=1 forces them on.
    v = os.environ.get("DSSD_GRAPHS")
    if v is not None:
        return v == "1"
    return False if os.environ.get("SPACES_ZERO_GPU") else "auto"


def load_models(backend: str | None = None) -> dict:
    """{display name: DSSD}. backend "hub" (default): the published heads and their base models
    on CUDA; "tiny": a random tiny model on CPU for local UI work."""
    backend = backend or os.environ.get("DSSD_DEMO_BACKEND", "hub")
    if backend == "tiny":
        from .tiny import tiny_engine
        return {"Tiny random model (local test)": tiny_engine()}
    import torch
    from dssd import DSSD
    names = [n.strip() for n in os.environ.get("DSSD_DEMO_MODELS", ",".join(MODELS)).split(",") if n.strip()]
    device = "cuda" if (torch.cuda.is_available() or os.environ.get("SPACES_ZERO_GPU")) else "cpu"
    token = os.environ.get("HF_TOKEN")
    return {n: DSSD.from_pretrained(MODELS[n], device=device, graphs=_graphs(), token=token) for n in names}


@dataclass
class Update:
    dssd: str           # HTML of the DSSD output (coloured by source)
    ar: str             # HTML of the autoregressive output ("" when not compared)
    stats: str          # HTML stat cards (final update) or a progress line
    replay: str         # hidden replay payload (final update)
    done: bool = False


def generate_updates(engine, prompt: str, domain: str, max_new: int, compare: bool,
                     min_interval: float = 0.05) -> Iterator[Update]:
    """Stream DSSD (then AR when compare) as HTML updates, at most one every min_interval seconds,
    then a final update with stats and the replay timelines."""
    if not prompt or not prompt.strip():
        raise ValueError("Please enter a prompt.")
    n = len(engine.tokenizer.encode(prompt, add_special_tokens=False))
    if n > MAX_PROMPT_TOKENS:
        raise ValueError(f"The prompt is too long ({n} tokens; at most {MAX_PROMPT_TOKENS}).")
    max_new = max(1, min(int(max_new), MAX_NEW_TOKENS))
    detok = {"dssd": Detokenizer(engine.tokenizer), "ar": Detokenizer(engine.tokenizer)}
    pieces = {"dssd": [], "ar": []}
    timeline = {"dssd": [], "ar": []}
    last = 0.0

    def view(done=False, stats="", replay=""):
        return Update(tokens_html(pieces["dssd"], cursor=not done and not pieces["ar"]),
                      tokens_html(pieces["ar"], colored=False, cursor=not done and bool(pieces["ar"]))
                      if compare else "",
                      stats, replay, done)

    for ev in engine.stream(prompt, max_new, domain=domain, ar=compare):
        if ev.done:
            for ph in ("dssd", "ar"):
                new = detok[ph].flush()
                pieces[ph] += new
                if timeline[ph]:
                    timeline[ph] += [[timeline[ph][-1][0], t, s] for t, s in new]
            yield view(True, stats_html(ev.result), replay_html(timeline))
            return
        new = detok[ev.phase].feed(ev.token_ids, ev.sources)
        pieces[ev.phase] += new
        timeline[ev.phase] += [[round(ev.t, 4), t, s] for t, s in new]
        now = time.perf_counter()
        if now - last >= min_interval:
            last = now
            phase = "Generating with DSSD…" if ev.phase == "dssd" else "Generating autoregressively for comparison…"
            yield view(stats=f'<div class="progress">{phase}</div>')