"""Constrained decoding: a per-sequence structure state machine + token sampling.

The state machine guarantees every generated media block is well formed: inside <image> only
visual codes are allowed and exactly G of them; inside a video, frames are
<frame><t_k> G codes </frame> with consecutive timestamps; text spans can only contain text
tokens.  Purely structural tokens are *forced* (no sampling); content tokens are sampled from
the model's distribution restricted to the legal set.
"""
from __future__ import annotations

from dataclasses import dataclass

import torch

from ..tokenizers.multimodal import MAX_FRAMES, UnifiedVocab


@dataclass
class Sampling:
    temperature: float = 1.0
    top_k: int = 0
    top_p: float = 1.0
    repetition_penalty: float = 1.0  # applied to previously generated *text* tokens only
    max_new_tokens: int = 256
    seed: int = 0
    cfg_scale: float = 1.0  # classifier-free guidance on visual tokens (1.0 = off)


class Masks:
    """Boolean legal-token sets, precomputed once per vocabulary/device."""

    def __init__(self, v: UnifiedVocab, device):
        V = v.size

        def mk(ids=(), ranges=()):
            m = torch.zeros(V, dtype=torch.bool, device=device)
            for lo, hi in ranges:
                m[lo:hi] = True
            if ids:
                m[list(ids)] = True
            return m

        self.text = mk([v.eos], [(0, v.text_size)])
        self.text_media = mk([v.eos, v.img_start, v.vid_start], [(0, v.text_size)])
        self.visual = mk(ranges=[(v.visual_offset, v.visual_offset + v.visual_size)])
        self.timestamp = mk(ranges=[(v.timestamp(0), v.timestamp(0) + MAX_FRAMES)])
        self.frame_or_end = mk([v.frm_start, v.vid_end])


class StructureFSM:
    """Tracks where one sequence is in the media grammar.

    mode:
      "text"   free text until <eos> (media tags disallowed unless allow_media)
      "image"  we are right after <image>: G codes, </image>, then <eos> if stop_after_media
      "video"  we are inside <video> right before the first <frame>
    n_frames: exact total frame count for videos (None = model decides, bounded by max_timestamp)
    first_ts: timestamp of the next frame to generate (None = sampled by the model)
    frames_done: frames already present in the prefix (prefix then ends right after a </frame>)
    codes_done: image codes already present after <image> (image completion)
    """

    def __init__(self, v: UnifiedVocab, mode: str = "text", allow_media: bool = False, stop_after_media: bool = True,
                 n_frames: int | None = None, first_ts: int | None = None, max_timestamp: int = 7, frames_done: int = 0, codes_done: int = 0):
        self.v, self.allow_media, self.stop_after_media = v, allow_media, stop_after_media
        self.n_frames, self.next_ts, self.max_ts = n_frames, first_ts, max_timestamp
        self.done = False
        self.frames = frames_done
        self.state = {"text": "text", "image": "img_codes", "video": "after_frame" if frames_done else "frame_start"}[mode]
        self.count = codes_done if mode == "image" else 0

    def constraint(self, masks: Masks) -> tuple[int | None, torch.Tensor | None]:
        """(forced_token, None) or (None, legal_mask)."""
        v, s = self.v, self.state
        if self.done:
            return v.pad, None
        if s == "text":
            return None, masks.text_media if self.allow_media else masks.text
        if s in ("img_codes", "frm_codes"):
            return None, masks.visual
        if s == "img_end":
            return v.img_end, None
        if s == "frm_end":
            return v.frm_end, None
        if s == "frame_start":
            return v.frm_start, None
        if s == "timestamp":
            return (v.timestamp(self.next_ts), None) if self.next_ts is not None else (None, masks.timestamp)
        if s == "after_frame":
            last = self.next_ts - 1
            if (self.n_frames is not None and self.frames >= self.n_frames) or last >= self.max_ts:
                return v.vid_end, None
            return (v.frm_start, None) if self.n_frames is not None else (None, masks.frame_or_end)
        if s == "eos":
            return v.eos, None
        raise AssertionError(s)

    def advance(self, tok: int) -> None:
        v, s = self.v, self.state
        if self.done:
            return
        G = v.tokens_per_frame
        if s == "text":
            if tok == v.eos:
                self.done = True
            elif tok == v.img_start:
                self.state, self.count = "img_codes", 0
            elif tok == v.vid_start:
                self.state = "frame_start"
        elif s == "img_codes":
            self.count += 1
            if self.count == G:
                self.state = "img_end"
        elif s == "img_end":
            self.state = "eos" if self.stop_after_media else "text"
        elif s == "frame_start":
            self.state = "timestamp"
        elif s == "timestamp":
            self.next_ts = tok - v.timestamp(0) + 1
            self.state, self.count = "frm_codes", 0
        elif s == "frm_codes":
            self.count += 1
            if self.count == G:
                self.state = "frm_end"
        elif s == "frm_end":
            self.frames += 1
            self.state = "after_frame"
        elif s == "after_frame":
            if tok == v.frm_start:
                self.state = "timestamp"
            else:  # </video>
                self.state = "eos" if self.stop_after_media else "text"
        elif s == "eos":
            self.done = True

    @property
    def in_visual(self) -> bool:
        return self.state in ("img_codes", "frm_codes")


def sample_next(logits: torch.Tensor, legal: torch.Tensor, sp: Sampling, gen: torch.Generator,
                history: torch.Tensor | None = None, text_size: int = 0) -> torch.Tensor:
    """logits [B,V] (float32), legal [B,V] bool -> sampled ids [B]."""
    logits = logits.float().clone()
    if sp.repetition_penalty != 1.0 and history is not None and history.numel():
        prev = history.clamp(min=0)
        is_text = prev < text_size
        vals = logits.gather(1, prev)
        pen = torch.where(vals > 0, vals / sp.repetition_penalty, vals * sp.repetition_penalty)
        logits.scatter_(1, prev, torch.where(is_text, pen, vals))
    logits = logits.masked_fill(~legal, float("-inf"))
    if sp.temperature <= 0:
        return logits.argmax(-1)
    logits = logits / sp.temperature
    if sp.top_k > 0:
        kth = logits.topk(min(sp.top_k, logits.size(-1)), dim=-1).values[:, -1:]
        logits = logits.masked_fill(logits < kth, float("-inf"))
    if sp.top_p < 1.0:
        srt, idx = logits.sort(dim=-1, descending=True)
        probs = srt.softmax(-1)
        drop = probs.cumsum(-1) - probs > sp.top_p  # keep the smallest prefix with mass >= top_p
        logits = logits.masked_fill(drop.scatter(1, idx, drop), float("-inf"))
    return torch.multinomial(logits.softmax(-1), 1, generator=gen).squeeze(1)