"""Prompt layout, sampling, and streaming decode for booper-chat.

booper-chat was SFT'd on the *pair* layout -- `<bos> prompt <sep> reply <eos>`,
loss on the reply only -- so that is what it has to be served with. Fed the
plain continuation layout it keeps writing more of the user's message instead
of answering it, which is the single easiest way to make this model look
broken. See the model card and babble's `BABBLE_SERVE_LAYOUT=pair` note.

Multi-turn is the same layout with the earlier turns newline-joined into the
prompt, tail-truncated to whatever is left of the 1024-token window after the
reply's reservation.
"""

from __future__ import annotations

import os
import time
from dataclasses import dataclass, field
from pathlib import Path

import torch

from .bpe import BPETokenizer
from .model import Babbler, ModelConfig

BLOCK_SIZE = 1024


@dataclass
class SamplingParams:
    """Defaults follow babble's shipped serving config: temperature 0.5 was
    measured to be the single biggest quality lever over 1.0, and the two
    repetition knobs are both on because they fail differently. A flat
    `repetition_penalty` cannot escape a loop once one is running (live booper
    emitted ~85 consecutive "boop" tokens with it nominally active), so the
    count-scaled `frequency_penalty` carries the actual loop protection and
    keeps growing every repeat; `no_repeat_ngram_size` is the hard backstop
    that no penalty can be talked around.
    """

    max_new_tokens: int = 96
    temperature: float = 0.7
    top_k: int = 40
    top_p: float = 0.9
    repetition_penalty: float = 1.15
    frequency_penalty: float = 0.12
    presence_penalty: float = 0.0
    no_repeat_ngram_size: int = 3
    seed: int | None = None

    def clamped(self) -> "SamplingParams":
        """Clamp every knob to a sane band -- these arrive from a public web
        form, so nothing here may be trusted to be a number in range."""

        def num(value, lo, hi, fallback):
            try:
                value = float(value)
            except (TypeError, ValueError):
                return fallback
            if value != value:  # NaN
                return fallback
            return min(hi, max(lo, value))

        return SamplingParams(
            max_new_tokens=int(num(self.max_new_tokens, 1, 512, 96)),
            temperature=num(self.temperature, 0.0, 2.0, 0.7),
            top_k=int(num(self.top_k, 0, 16384, 40)),
            top_p=num(self.top_p, 0.0, 1.0, 0.9),
            repetition_penalty=num(self.repetition_penalty, 1.0, 2.0, 1.15),
            frequency_penalty=num(self.frequency_penalty, 0.0, 2.0, 0.12),
            presence_penalty=num(self.presence_penalty, 0.0, 2.0, 0.0),
            no_repeat_ngram_size=int(num(self.no_repeat_ngram_size, 0, 8, 3)),
            seed=self.seed,
        )


@dataclass
class Turn:
    role: str  # "user" | "assistant"
    content: str


class Booper:
    def __init__(self, checkpoint: str | Path, tokenizer_json: str | Path, *,
                 threads: int | None = None, quantize: bool | str | None = None) -> None:
        threads = threads if threads is not None else int(os.environ.get("BOOPER_THREADS", "0")) or None
        if threads:
            torch.set_num_threads(max(1, threads))
            try:
                torch.set_num_interop_threads(1)
            except RuntimeError:
                pass  # already set by an earlier init in this process
        if hasattr(torch.backends, "mkldnn"):
            torch.backends.mkldnn.enabled = True
        try:
            torch.set_flush_denormal(True)
        except Exception:
            pass

        self.tok = BPETokenizer.from_json(tokenizer_json)
        blob = torch.load(checkpoint, map_location="cpu", weights_only=False)
        config = ModelConfig.from_dict(blob["config"])
        if config.vocab_size != self.tok.vocab_size:
            raise ValueError(
                f"checkpoint vocab_size {config.vocab_size} != tokenizer "
                f"vocab_size {self.tok.vocab_size} -- wrong tokenizer.json"
            )
        model = Babbler(config)
        model.load_state_dict(blob["model"])
        model.eval()
        self.n_params = sum({id(p): p.numel() for p in model.parameters()}.values())

        if quantize is None:
            quantize = os.environ.get("BOOPER_QUANTIZE", "auto").strip().lower()
        if isinstance(quantize, bool):
            quantize = "1" if quantize else "0"

        self.quantized = False
        self.backend = "fp32"
        if quantize != "0":
            candidate = _quantize_dynamic(model, config, blob)
            if candidate is None:
                pass  # no quantized engine in this wheel; fp32 it is
            elif quantize == "auto":
                # int8 is a big win on x86 (fbgemm/VNNI) and a ~2x *loss* on
                # Apple qnnpack, so the answer depends entirely on what CPU
                # the Space actually landed on. Measuring costs a few hundred
                # ms once at startup and is the only way to be right on both.
                fp32_s = _decode_bench(model, config)
                int8_s = _decode_bench(candidate, config)
                if int8_s < fp32_s:
                    model, self.quantized = candidate, True
                self.backend = f"int8 ({int8_s * 1e3:.1f}ms/tok vs fp32 {fp32_s * 1e3:.1f})" if self.quantized \
                    else f"fp32 ({fp32_s * 1e3:.1f}ms/tok vs int8 {int8_s * 1e3:.1f})"
            else:
                model, self.quantized, self.backend = candidate, True, "int8 (forced)"

        self.model = model
        self.config = config
        self.block_size = config.block_size
        del blob

    # --- prompt layout ----------------------------------------------------

    def build_context(self, turns: list[Turn], reserve: int) -> tuple[list[int], int]:
        """`<bos> …history… <sep>`, filling the window from the newest turn back.

        Turns are added newest-first until the next one would not fit, so the
        model always sees the whole of the current message and as much prior
        conversation as the 1024-token window still has room for. Returns the
        context and how many turns of history survived the cut.
        """
        budget = max(1, self.block_size - 2 - reserve)  # bos + sep cost a slot each
        newline = self.tok.encode("\n")
        pieces: list[list[int]] = []
        used = 0
        kept = 0
        for index, turn in enumerate(reversed(turns)):
            ids = self.tok.encode(turn.content.strip())
            cost = len(ids) + (0 if index == 0 else len(newline))
            if used + cost > budget:
                if index == 0:
                    # The current message alone overflows: keep its tail, the
                    # part being replied to, same rule training truncated with.
                    pieces.append(ids[-budget:])
                    used = budget
                    kept = 1
                break
            pieces.append(ids if index == 0 else ids + newline)
            used += cost
            kept += 1

        body: list[int] = []
        for ids in reversed(pieces):
            body.extend(ids)
        return [self.tok.specials.bos, *body, self.tok.specials.sep], kept

    # --- sampling ---------------------------------------------------------

    def _next_token(self, logits: torch.Tensor, generated: torch.Tensor, p: SamplingParams,
                    generator: torch.Generator) -> int:
        logits = logits.clone()
        logits[self._banned] = float("-inf")

        if generated.numel():
            counts = torch.bincount(generated, minlength=logits.numel()).to(logits.dtype)
            present = counts > 0
            if p.repetition_penalty != 1.0:
                # HF-style: positive logits divided, negative multiplied, so a
                # seen token gets less likely without flipping sign.
                scaled = torch.where(logits < 0, logits * p.repetition_penalty,
                                     logits / p.repetition_penalty)
                logits = torch.where(present, scaled, logits)
            if p.frequency_penalty or p.presence_penalty:
                # Count-scaled, so the discount keeps growing every repeat --
                # this is the one that actually breaks a running loop.
                logits = logits - counts * p.frequency_penalty - present.to(logits.dtype) * p.presence_penalty
            if p.no_repeat_ngram_size > 1:
                n = p.no_repeat_ngram_size
                seq = generated.tolist()
                if len(seq) >= n - 1:
                    prefix = tuple(seq[-(n - 1):])
                    banned = {
                        seq[j + n - 1]
                        for j in range(len(seq) - n + 1)
                        if tuple(seq[j:j + n - 1]) == prefix
                    }
                    if banned:
                        logits[list(banned)] = float("-inf")

        if p.temperature <= 0:
            return int(logits.argmax())
        logits = logits / p.temperature
        if p.top_k and 0 < p.top_k < logits.numel():
            cutoff = torch.topk(logits, p.top_k).values[-1]
            logits = logits.masked_fill(logits < cutoff, float("-inf"))
        if 0.0 < p.top_p < 1.0:
            ordered, index = torch.sort(logits, descending=True)
            mass = torch.cumsum(torch.softmax(ordered, dim=-1), dim=-1)
            drop = mass > p.top_p
            drop[1:] = drop[:-1].clone()  # always keep the top token
            drop[0] = False
            logits = torch.full_like(logits, float("-inf")).scatter(
                0, index, ordered.masked_fill(drop, float("-inf"))
            )
        probs = torch.softmax(logits, dim=-1)
        return int(torch.multinomial(probs, 1, generator=generator))

    @property
    def _banned(self) -> list[int]:
        # Structural ids are inputs, never outputs -- a stray <sep> or <bos>
        # in the middle of a reply is always a decode bug, never a message.
        s = self.tok.specials
        return [s.pad, s.bos, s.sep]

    @torch.inference_mode()
    def stream(self, turns: list[Turn], params: SamplingParams | None = None):
        """Yield `("token", text)` chunks, then one `("done", stats)`.

        Bytes are emitted as soon as they form valid UTF-8 -- a multi-byte
        character (an emoji, most of what a Discord model likes to say) spans
        several BPE ids, and flushing half of one would render as a replacement
        glyph that never un-renders.
        """
        p = (params or SamplingParams()).clamped()
        generator = torch.Generator()
        generator.manual_seed(p.seed if p.seed is not None else int.from_bytes(os.urandom(8), "big") % (2**63))

        context, kept = self.build_context(turns, reserve=p.max_new_tokens)
        room = self.block_size - len(context)
        budget = max(1, min(p.max_new_tokens, room))

        started = time.perf_counter()
        cache = self.model.new_cache(1, max_len=len(context) + budget)
        logits = self.model(torch.tensor([context]), cache=cache)[0, -1]
        prefill_s = time.perf_counter() - started

        eos = self.tok.specials.eos
        produced: list[int] = []
        generated = torch.empty(0, dtype=torch.long)
        pending = bytearray()
        step = torch.empty((1, 1), dtype=torch.long)
        stop = "eos"

        for i in range(budget):
            nxt = self._next_token(logits, generated, p, generator)
            if nxt == eos:
                break
            produced.append(nxt)
            generated = torch.cat([generated, torch.tensor([nxt])])
            pending.extend(self.tok.vocab.get(nxt, b""))
            try:
                text = pending.decode("utf-8")
            except UnicodeDecodeError:
                pass  # mid-character; hold the bytes until it completes
            else:
                pending.clear()
                if text:
                    yield "token", text
            if i == budget - 1 or cache.length >= cache.max_len:
                stop = "length"
                break
            step[0, 0] = nxt
            logits = self.model(step, cache=cache)[0, -1]

        if pending:
            yield "token", pending.decode("utf-8", errors="replace")

        elapsed = time.perf_counter() - started
        yield "done", {
            "stop": stop,
            "tokens": len(produced),
            "prompt_tokens": len(context),
            "context_used": len(context) + len(produced),
            "context_limit": self.block_size,
            "turns_kept": kept,
            "prefill_ms": round(prefill_s * 1000, 1),
            "elapsed_ms": round(elapsed * 1000, 1),
            "tok_per_s": round(len(produced) / elapsed, 1) if elapsed > 0 and produced else 0.0,
        }

    def reply(self, turns: list[Turn], params: SamplingParams | None = None) -> str:
        parts = [text for kind, text in self.stream(turns, params) if kind == "token"]
        return "".join(parts).strip()


def _quantize_dynamic(model: Babbler, config: ModelConfig, blob: dict) -> Babbler | None:
    """A dynamic-int8 copy of `model`, or None if this build has no engine.

    `quantize_dynamic` needs `lm_head` to own its weight, but booper ties it to
    the token embedding, so the copy unties first. That untying is decode-only
    and never reaches a checkpoint -- which is why this builds a fresh model
    rather than mutating the one that was passed in.
    """
    import torch.nn as nn

    try:
        if torch.backends.quantized.engine == "none":
            # macOS/ARM wheels leave this unset and every prepack then fails
            # with NoQEngine; x86 Linux picks fbgemm on its own.
            engines = [e for e in torch.backends.quantized.supported_engines if e != "none"]
            if not engines:
                return None
            torch.backends.quantized.engine = engines[0]
        clone = Babbler(config)
        clone.load_state_dict(blob["model"])
        clone.lm_head.weight = nn.Parameter(clone.tok_emb.weight.detach().clone())
        clone.eval()
        return torch.ao.quantization.quantize_dynamic(clone, {nn.Linear}, dtype=torch.qint8)
    except Exception:
        return None


@torch.inference_mode()
def _decode_bench(model, config: ModelConfig, steps: int = 12) -> float:
    """Median seconds per cached single-token forward -- the thing that decides
    how fast a reply feels. Median, not mean, so one scheduler hiccup on a
    shared vCPU does not pick the wrong backend."""
    cache = model.new_cache(1, max_len=steps + 8)
    step = torch.zeros((1, 1), dtype=torch.long)
    model(torch.zeros((1, 4), dtype=torch.long), cache=cache)  # warm the kernels
    times = []
    for _ in range(steps):
        started = time.perf_counter()
        model(step, cache=cache)
        times.append(time.perf_counter() - started)
    times.sort()
    return times[len(times) // 2]