"""Label-free JSON-shaped 137-slot input; no training, targets or decoding.

Literal JSON punctuation is fixed context. Every predicted value is masked,
including inactive fields. This template is not itself a valid JSON response.
"""

from __future__ import annotations

from dataclasses import dataclass

MASK = "<|fim_pad|>"
FIELD_WIDTHS = (
    ("action", 3), ("coordinate.x", 4), ("coordinate.y", 4),
    ("keys", 16), ("pixels", 4), ("time", 5), ("status", 1), ("text", 100),
)
OUTPUT_SLOTS = 137
# Each literal precedes one value block. keys brackets are context, not targets.
_LITERALS = (
    '{"action":"', '","coordinate":[', ',', '],"keys":[',
    '],"pixels":', ',"time":', ',"status":"', '","text":"', '"}',
)


def render_template() -> str:
    """Return the exact compact JSON-shaped suffix, without trailing newline."""
    return "".join(literal + MASK * width for literal, (_, width)
                   in zip(_LITERALS[:-1], FIELD_WIDTHS, strict=True)) + _LITERALS[-1]


@dataclass(frozen=True)
class ParallelInput:
    input_ids: tuple[int, ...]
    slot_positions: tuple[int, ...]
    field_positions: tuple[tuple[str, tuple[int, ...]], ...]
    prefix_length: int

    @property
    def prediction_positions(self) -> tuple[int, ...]:
        """Existing parallel-model alignment: head at j-1 predicts slot j."""
        return tuple(position - 1 for position in self.slot_positions)


def build_parallel_input(tokenizer, prefix_ids=()) -> ParallelInput:
    """Append the fixed suffix to already-rendered observation token IDs.

    Tokenize literals independently and insert native mask IDs explicitly, as
    in the existing fixed-table codecs. No reference action is accepted.
    This constructs CPU-side IDs only; attention masks and zero slot embeddings
    remain responsibilities of future model integration.
    """
    prefix = tuple(prefix_ids)
    vocab_size = tokenizer.get_vocab_size()
    if any(type(value) is not int or not 0 <= value < vocab_size for value in prefix):
        raise ValueError("prefix IDs must be in-vocabulary integers")
    mask_id = tokenizer.token_to_id(MASK)
    if type(mask_id) is not int or tokenizer.encode(MASK, add_special_tokens=False).ids != [mask_id]:
        raise ValueError("the original single-token mask is required")
    ids, positions, fields = list(prefix), [], []
    for literal, (name, width) in zip(_LITERALS[:-1], FIELD_WIDTHS, strict=True):
        ids.extend(tokenizer.encode(literal, add_special_tokens=False).ids)
        field = tuple(range(len(ids), len(ids) + width))
        fields.append((name, field))
        positions.extend(field)
        ids.extend([mask_id] * width)
    ids.extend(tokenizer.encode(_LITERALS[-1], add_special_tokens=False).ids)
    if len(positions) != OUTPUT_SLOTS or any(ids[p] != mask_id for p in positions):
        raise ValueError("parallel input slot contract differs")
    return ParallelInput(tuple(ids), tuple(positions), tuple(fields), len(prefix))