"""Explicit ActionCodec2Sparse token constraints for Hugging Face generate()."""

from __future__ import annotations

import torch
from transformers import LogitsProcessor

from ...backends.actioncodec2.cache import BoundedCache


class ActionCodec2SparseLogitsProcessor(LogitsProcessor):
    """Constrain generated action IDs to the sparse grammar; allow EOS once complete."""

    def __init__(
        self, codec, horizon, *, fps, prompt_length, eos_token_id, token_offset=0, pad_token_id=None
    ):
        for name, value in {
            "prompt_length": prompt_length,
            "eos_token_id": eos_token_id,
            "token_offset": token_offset,
        }.items():
            if type(value) is not int or value < 0:
                raise ValueError(f"{name} must be a nonnegative integer")
        if pad_token_id is not None and (type(pad_token_id) is not int or pad_token_id < 0):
            raise ValueError("pad_token_id must be a nonnegative integer")
        self.grammar = codec.grammar(horizon, fps=fps)
        self.prompt_length, self.token_offset = prompt_length, token_offset
        self.eos_token_id, self.pad_token_id = eos_token_id, pad_token_id
        stop = token_offset + self.grammar.vocab_size
        if token_offset <= eos_token_id < stop:
            raise ValueError("eos_token_id must be outside the action token interval")
        if pad_token_id is not None and token_offset <= pad_token_id < stop:
            raise ValueError("pad_token_id must be outside the action token interval")
        self._mask_cache = BoundedCache(64)

    def _state(self, row):
        if self.eos_token_id in row:
            end = row.index(self.eos_token_id)
            if any(t not in (self.eos_token_id, self.pad_token_id) for t in row[end:]):
                raise ValueError("only EOS/padding may follow completed action generation")
            if not self.grammar.is_complete([t - self.token_offset for t in row[:end]]):
                raise ValueError("EOS occurred before the action chunk was complete")
            return "done"
        state = self.grammar._consume([t - self.token_offset for t in row])
        if state is None:
            raise ValueError(
                "generated prefix violates the action grammar; check prompt_length and token_offset"
            )
        return state

    def _mask(self, state, scores):
        key = (state, str(scores.device), scores.shape[-1])
        mask = self._mask_cache.get(key)
        if mask is None:
            mask = torch.zeros(scores.shape[-1], device=scores.device, dtype=torch.bool)
            if state == "done":
                mask[self.eos_token_id] = True
                if self.pad_token_id is not None:
                    mask[self.pad_token_id] = True
            else:
                allowed = self.grammar._mask_from_state(state)
                stop = self.token_offset + len(allowed)
                mask[self.token_offset : stop] = torch.as_tensor(allowed, device=scores.device)
                mask[self.eos_token_id] = self.grammar._state_complete(state)
            self._mask_cache[key] = mask
        return mask

    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("expected input_ids(B,L) and scores(B,V) with the same batch size")
        if input_ids.shape[1] < self.prompt_length:
            raise ValueError("prompt_length exceeds the supplied input length")
        required = max(
            self.token_offset + self.grammar.vocab_size,
            self.eos_token_id + 1,
            0 if self.pad_token_id is None else self.pad_token_id + 1,
        )
        if scores.shape[1] < required:
            raise ValueError("model vocabulary does not cover configured action/EOS/padding IDs")
        self.grammar.reserve_states(int(input_ids.shape[0]))
        rows = input_ids[:, self.prompt_length :].detach().cpu().tolist()
        output = scores.clone()
        groups = {}
        for index, row in enumerate(rows):
            groups.setdefault(self._state(row), []).append(index)
        for state, indices in groups.items():
            output[indices] = scores[indices].masked_fill(~self._mask(state, scores), -torch.inf)
        if not torch.isfinite(output).any(dim=1).all():
            raise ValueError(
                "all legal action logits were suppressed; check other generation processors"
            )
        return output