"""Speaker-neutral grammar constraints for MOSS transcript generation.

The constraint enforces only the wire syntax.  It never chooses a speaker,
timestamp, or text token on the model's behalf.  Token transitions operate on
decoded token pieces so merged BPE tokens such as ``][`` and ``[S`` remain
valid when they cross grammar boundaries.
"""

from __future__ import annotations

import copy
import hashlib
import time
from dataclasses import dataclass
from typing import Iterable

import torch
from transformers import LogitsProcessor


SCHEMA = "moss-transcript-grammar-v1"


@dataclass(frozen=True)
class GrammarState:
    phase: str = "start"


def _digit(value: str) -> bool:
    return "0" <= value <= "9"


def advance_character(state: GrammarState, char: str) -> GrammarState | None:
    """Advance one Unicode character, or return ``None`` for invalid syntax."""

    phase = state.phase
    if phase == "start":
        return GrammarState("start_time_first") if char == "[" else None
    if phase == "start_time_first":
        return GrammarState("start_time_int") if _digit(char) else None
    if phase == "start_time_int":
        if _digit(char):
            return state
        if char == ".":
            return GrammarState("start_time_frac_first")
        if char == "]":
            return GrammarState("after_start_time")
        return None
    if phase == "start_time_frac_first":
        return GrammarState("start_time_frac") if _digit(char) else None
    if phase == "start_time_frac":
        if _digit(char):
            return state
        return GrammarState("after_start_time") if char == "]" else None
    if phase == "after_start_time":
        if char.isspace():
            return state
        return GrammarState("speaker_s") if char == "[" else None
    if phase == "speaker_s":
        return GrammarState("speaker_digit_1") if char == "S" else None
    if phase == "speaker_digit_1":
        return GrammarState("speaker_digit_2") if _digit(char) else None
    if phase == "speaker_digit_2":
        return GrammarState("speaker_close") if _digit(char) else None
    if phase == "speaker_close":
        return GrammarState("text_empty") if char == "]" else None
    if phase == "text_empty":
        if char in "[]":
            return None
        return state if char.isspace() else GrammarState("text")
    if phase == "text":
        if char == "[":
            return GrammarState("end_time_first")
        if char == "]":
            return None
        return state
    if phase == "end_time_first":
        return GrammarState("end_time_int") if _digit(char) else None
    if phase == "end_time_int":
        if _digit(char):
            return state
        if char == ".":
            return GrammarState("end_time_frac_first")
        if char == "]":
            return GrammarState("after_end_time")
        return None
    if phase == "end_time_frac_first":
        return GrammarState("end_time_frac") if _digit(char) else None
    if phase == "end_time_frac":
        if _digit(char):
            return state
        return GrammarState("after_end_time") if char == "]" else None
    if phase == "after_end_time":
        if char.isspace():
            return state
        return GrammarState("start_time_first") if char == "[" else None
    raise ValueError(f"unknown transcript grammar phase: {phase}")


def advance_piece(state: GrammarState, piece: str) -> GrammarState | None:
    if not piece:
        return None
    for char in piece:
        state = advance_character(state, char)
        if state is None:
            return None
    return state


def accepting(state: GrammarState) -> bool:
    return state.phase == "after_end_time"


class TranscriptGrammarVocabulary:
    """Tokenizer-specific transition cache shared across utterances."""

    def __init__(self, tokenizer):
        self.tokenizer = tokenizer
        self.vocab_size = len(tokenizer)
        eos = tokenizer.eos_token_id
        self.eos_token_ids = {int(value) for value in (eos if isinstance(eos, list) else [eos])}
        self.special_ids = {int(value) for value in getattr(tokenizer, "all_special_ids", [])}
        self.pieces = tuple(
            tokenizer.decode(
                [token_id], skip_special_tokens=False, clean_up_tokenization_spaces=False
            )
            for token_id in range(self.vocab_size)
        )
        digest = hashlib.sha256()
        for token_id, piece in enumerate(self.pieces):
            digest.update(str(token_id).encode("ascii"))
            digest.update(b"\0")
            digest.update(piece.encode("utf-8", errors="surrogatepass"))
            digest.update(b"\n")
        self.token_surface_sha256 = digest.hexdigest()
        self._allowed_cpu: dict[GrammarState, tuple[int, ...]] = {}
        self._allowed_device: dict[tuple[GrammarState, str], torch.Tensor] = {}

    def consume(self, token_ids: Iterable[int]) -> GrammarState:
        state = GrammarState()
        for token_id in token_ids:
            token_id = int(token_id)
            if token_id in self.eos_token_ids:
                if not accepting(state):
                    raise ValueError("EOS before a complete end timestamp")
                continue
            if token_id in self.special_ids or not 0 <= token_id < self.vocab_size:
                raise ValueError(f"invalid special/out-of-vocabulary token in transcript: {token_id}")
            next_state = advance_piece(state, self.pieces[token_id])
            if next_state is None:
                raise ValueError(
                    f"token {token_id} piece={self.pieces[token_id]!r} violates phase={state.phase}"
                )
            state = next_state
        return state

    def allowed_token_ids(self, state: GrammarState, device: torch.device) -> torch.Tensor:
        if state not in self._allowed_cpu:
            allowed = []
            for token_id, piece in enumerate(self.pieces):
                if token_id in self.special_ids:
                    continue
                if advance_piece(state, piece) is not None:
                    allowed.append(token_id)
            # An empty response is not a valid transcript for deployment and
            # the parser cannot score it.  EOS therefore becomes available
            # only after the model has completed a full segment.
            if accepting(state):
                allowed.extend(self.eos_token_ids)
            if not allowed:
                raise RuntimeError(f"transcript grammar has no continuation from {state.phase}")
            self._allowed_cpu[state] = tuple(sorted(set(allowed)))
        key = (state, str(device))
        if key not in self._allowed_device:
            self._allowed_device[key] = torch.tensor(
                self._allowed_cpu[state], dtype=torch.long, device=device
            )
        return self._allowed_device[key]

    def audit(self) -> dict:
        return {
            "schema": SCHEMA,
            "speaker_policy": "model_selected_exact_two_digit_tag",
            "timestamp_policy": "model_selected_nonnegative_number",
            "text_policy": "model_selected_nonempty_no_brackets",
            "eos_policy": "complete_segment_only",
            "vocab_size": self.vocab_size,
            "token_surface_sha256": self.token_surface_sha256,
            "eos_token_ids": sorted(self.eos_token_ids),
        }


class TranscriptGrammarLogitsProcessor(LogitsProcessor):
    """Mask tokens that cannot extend the canonical transcript grammar."""

    def __init__(
        self,
        tokenizer,
        prompt_length: int,
        *,
        vocabulary: TranscriptGrammarVocabulary | None = None,
    ):
        if prompt_length < 0:
            raise ValueError("prompt_length must be non-negative")
        self.prompt_length = int(prompt_length)
        self.vocabulary = vocabulary or TranscriptGrammarVocabulary(tokenizer)
        if self.vocabulary.tokenizer is not tokenizer:
            raise ValueError("constraint vocabulary belongs to another tokenizer instance")
        self.tokenizer = tokenizer
        self.vocab_size = self.vocabulary.vocab_size
        self.eos_token_ids = self.vocabulary.eos_token_ids

    def consume(self, token_ids: Iterable[int]) -> GrammarState:
        return self.vocabulary.consume(token_ids)

    def allowed_token_ids(self, state: GrammarState, device: torch.device) -> torch.Tensor:
        return self.vocabulary.allowed_token_ids(state, device)

    def audit(self) -> dict:
        return self.vocabulary.audit()

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        if input_ids.ndim != 2 or scores.ndim != 2 or input_ids.shape[0] != scores.shape[0]:
            raise ValueError("transcript grammar requires aligned 2-D token and score batches")
        constrained = torch.full_like(scores, -torch.inf)
        # Hugging Face expands a batch-one prompt to one row per beam.  Replay
        # each prefix independently so beam reordering cannot leak DFA state.
        for row in range(input_ids.shape[0]):
            generated = input_ids[row, self.prompt_length :].tolist()
            state = self.consume(generated)
            allowed = self.allowed_token_ids(state, scores.device)
            constrained[row].index_copy_(0, allowed, scores[row].index_select(0, allowed))
        return constrained


def generate_constrained_transcription(
    model,
    processor,
    messages,
    *,
    max_length: int = 131072,
    max_new_tokens: int,
    device: torch.device,
    dtype: torch.dtype,
    constraint_vocabulary: TranscriptGrammarVocabulary | None = None,
    num_beams: int = 1,
) -> dict:
    """Generate one transcript with syntax-only constraints."""

    from moss_transcribe_diarize.inference_utils import prepare_inputs

    preprocess_context = (
        torch.amp.autocast("cuda", dtype=dtype)
        if device.type == "cuda" and dtype in (torch.float16, torch.bfloat16)
        else torch.no_grad()
    )
    with preprocess_context:
        inputs = prepare_inputs(
            processor, messages, max_length=max_length, device=device
        ).to(device)
    prompt_length = int(inputs["attention_mask"][0].sum().item())
    constraint = TranscriptGrammarLogitsProcessor(
        processor.tokenizer,
        prompt_length,
        vocabulary=constraint_vocabulary,
    )
    generation_config = copy.deepcopy(model.generation_config)
    generation_config.max_new_tokens = int(max_new_tokens)
    generation_config.do_sample = False
    generation_config.num_beams = int(num_beams)
    generation_config.num_return_sequences = 1
    if num_beams > 1:
        generation_config.early_stopping = True
    generation_context = (
        torch.amp.autocast("cuda", dtype=dtype)
        if device.type == "cuda" and dtype in (torch.float16, torch.bfloat16)
        else torch.no_grad()
    )
    started = time.perf_counter()
    with torch.inference_mode(), generation_context:
        outputs = model.generate(
            input_ids=inputs["input_ids"],
            attention_mask=inputs["attention_mask"],
            input_features=inputs["input_features"],
            audio_feature_lengths=inputs["audio_feature_lengths"],
            audio_chunk_mapping=inputs["audio_chunk_mapping"],
            generation_config=generation_config,
            logits_processor=[constraint],
        )
    if device.type == "cuda":
        torch.cuda.synchronize(device)
    generated_ids = outputs[0][prompt_length:]
    non_eos = [
        int(token_id)
        for token_id in generated_ids.tolist()
        if int(token_id) not in constraint.eos_token_ids
    ]
    final_state = constraint.consume(non_eos)
    text = processor.tokenizer.decode(
        generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
    ).strip()
    return {
        "text": text,
        "prompt_len": prompt_length,
        "generated_tokens": int(generated_ids.numel()),
        "inference_seconds": time.perf_counter() - started,
        "constraint_complete": accepting(final_state),
        "constraint_final_phase": final_state.phase,
        "constraint_num_beams": int(num_beams),
        "constraint_audit": constraint.audit(),
    }