"""Schema-constrained decoding of the `<|spec|>` JSON.

The grammar accepts exactly the canonical compact form written by `zenith.format.spec_json`:
`{"sim":"<sim>"}` for meta sims, otherwise `{"sim":"<sim>","op":"<op>"` followed by zero or more
`,"<param>":<number>` (params of that op, no repeats, numbers of at most 12 characters) and `}`.
Ranges and exclusive parameters are left to engine validation.
"""

from __future__ import annotations

import re

import torch

from zenith.engines.base import Param, SpecError
from zenith.engines.registry import META_SIMS
from zenith.model.generate import generate

_NUM_FULL = re.compile(r"-?(0|[1-9]\d*)(\.\d+)?(e-?\d+)?")
_NUM_PARTIAL = re.compile(r"-?|-?(0|[1-9]\d*)(\.\d*)?|-?(0|[1-9]\d*)(\.\d+)?e-?\d*")
_NUM_CHARS = set("-0123456789.e")
MAX_NUMBER_CHARS = 12

PREFIX, COMPLETE = "prefix", "complete"


class _Partial(Exception):
    """The text ended inside a valid construction."""


class _Invalid(Exception):
    """The text cannot be extended into a valid spec."""


class SpecGrammar:
    def __init__(self, schemas: dict[str, dict[str, list[Param]]]):
        self.ops = {sim: {op: [p.name for p in params] for op, params in ops.items()} for sim, ops in schemas.items()}
        self.sims = [*self.ops, *META_SIMS]

    def state(self, text: str) -> str | None:
        """PREFIX, COMPLETE, or None if `text` is not the start of a valid spec."""
        pos = 0

        def lit(s: str) -> None:
            nonlocal pos
            rest = text[pos:]
            if len(rest) < len(s):
                if s.startswith(rest):
                    raise _Partial
                raise _Invalid
            if not rest.startswith(s):
                raise _Invalid
            pos += len(s)

        def name(options) -> str:
            nonlocal pos
            rest = text[pos:]
            end = rest.find('"')
            if end < 0:
                if any(o.startswith(rest) for o in options):
                    raise _Partial
                raise _Invalid
            if rest[:end] not in options:
                raise _Invalid
            pos += end + 1
            return rest[:end]

        def number() -> None:
            nonlocal pos
            start = pos
            while pos < len(text) and text[pos] in _NUM_CHARS:
                pos += 1
            run = text[start:pos]
            if len(run) > MAX_NUMBER_CHARS:
                raise _Invalid
            if pos == len(text):
                # A run that is not yet a full number needs room to grow, or it could never be closed.
                if _NUM_FULL.fullmatch(run) or (_NUM_PARTIAL.fullmatch(run) and len(run) < MAX_NUMBER_CHARS):
                    raise _Partial
                raise _Invalid
            if not _NUM_FULL.fullmatch(run):
                raise _Invalid

        try:
            lit('{"sim":"')
            sim = name(self.sims)
            if sim in META_SIMS:
                lit("}")
            else:
                lit(',"op":"')
                params = self.ops[sim][name(list(self.ops[sim]))]
                used: set[str] = set()
                while True:
                    if pos == len(text):
                        raise _Partial
                    if text[pos] == "}":
                        pos += 1
                        break
                    remaining = [n for n in params if n not in used]
                    if not remaining:
                        raise _Invalid  # only "}" can follow once every parameter is used
                    lit(',"')
                    p = name(remaining)
                    used.add(p)
                    lit(":")
                    number()
            if pos != len(text):
                raise _Invalid
            return COMPLETE
        except _Partial:
            return PREFIX
        except _Invalid:
            return None

    def is_prefix(self, text: str) -> bool:
        return self.state(text) is not None

    def is_complete(self, text: str) -> bool:
        return self.state(text) == COMPLETE


class SpecConstraint:
    """`logits_hook` keeping only the `k` best-scoring tokens that extend a valid spec prefix."""

    def __init__(self, tok, grammar: SpecGrammar, prompt_len: int, k: int):
        self.tok, self.grammar, self.prompt_len, self.k = tok, grammar, prompt_len, k
        self.pieces = [tok.decode([i]) for i in range(tok.vocab_size)]

    def __call__(self, ids: list[int], logits: torch.Tensor) -> torch.Tensor:
        text = self.tok.decode(ids[self.prompt_len :])
        keep = []
        for i in torch.argsort(logits[0], descending=True).tolist():
            if logits[0, i] == float("-inf"):
                break
            if self.grammar.is_prefix(text + self.pieces[i]):
                keep.append(i)
                if len(keep) == self.k:
                    break
        if not keep:
            raise SpecError("no token can continue the spec")
        masked = torch.full_like(logits, float("-inf"))
        masked[0, keep] = logits[0, keep]
        return masked


_GRAMMARS: dict[int, SpecGrammar] = {}


def _grammar() -> SpecGrammar:
    from zenith.engines.registry import schemas

    if 0 not in _GRAMMARS:
        _GRAMMARS[0] = SpecGrammar(schemas())
    return _GRAMMARS[0]


def decode_spec(model, tok, prompt_ids: list[int], temperature: float, top_k: int | None, max_tokens: int = 96) -> str:
    """Generate a schema-valid spec after `prompt_ids` (which ends with `<|spec|>`)."""
    grammar = _grammar()
    n = len(prompt_ids)
    k = 1 if temperature == 0 else (top_k or 50)
    device = next(model.parameters()).device
    out = generate(model, torch.tensor([prompt_ids], device=device), max_tokens, temperature=temperature, top_k=top_k,
                   logits_hook=SpecConstraint(tok, grammar, n, k),
                   stop_fn=lambda ids: grammar.is_complete(tok.decode(ids[n:])))
    text = tok.decode(out[0, n:].tolist())
    if not grammar.is_complete(text):
        raise SpecError(f"spec not finished within {max_tokens} tokens: {text!r}")
    return text