"""ActionCodec2-only tensor and token-row validation helpers."""

from __future__ import annotations

from collections.abc import Sequence

import numpy as np
import torch

from ._types import ActionIdxBatch, TokenIds


class BaseActionTokenizer:
    """Shared action tensor methods, without policy-specific sentinels."""

    @property
    def new_action_tokens(self) -> list[str]:
        """Token strings in id order — downstream registers these into its VLM vocab."""
        return list(self._new_action_tokens)

    def to_token_strs(self, token_ids: TokenIds) -> list[str]:
        """Render one id row as its token strings, e.g. ``<right_arm_0>`` / ``<action0042>``.

        This is exactly the text a VLM would see after registering
        :attr:`new_action_tokens` — useful for eyeballing sequences and for
        prompt-template debugging.

        Args:
            token_ids: One row of action-token ids in ``[0, vocab_size)``.

        Returns:
            list[str]: Token strings in row order.
        """
        row = self._to_id_row(token_ids).tolist()
        n = len(self._new_action_tokens)
        bad = [i for i in row if not 0 <= i < n]
        if bad:
            raise ValueError(f"token ids {bad[:5]} out of range [0, {n})")
        return [self._new_action_tokens[i] for i in row]

    def _check_encode_kwargs(self, encode_kwargs: dict | None) -> dict:
        """Normalize ``encode_kwargs`` to a dict, rejecting keys outside :attr:`encode_kwargs`."""
        kw = dict(encode_kwargs or {})
        unknown = sorted(set(kw) - self.encode_kwargs)
        if unknown:
            raise ValueError(
                f"{type(self).__name__} does not consume encode_kwargs {unknown}; it declares {sorted(self.encode_kwargs)}"
            )
        return kw

    @staticmethod
    def _to_id_row(token_ids: TokenIds) -> torch.Tensor:
        """Validate and normalize one integer id row to a 1-D long tensor."""
        if isinstance(token_ids, torch.Tensor):
            row = token_ids.detach().cpu()
            if (
                row.dtype == torch.bool
                or row.dtype.is_floating_point
                or row.dtype.is_complex
            ):
                raise ValueError("token ids must be integers")
            row = row.to(torch.long)
        else:
            values = np.asarray(token_ids)
            if values.dtype.kind not in {"i", "u"}:
                raise ValueError("token ids must be integers")
            row = torch.from_numpy(np.asarray(values, dtype=np.int64))
        if row.ndim != 1:
            raise ValueError(f"token id row must be 1D, got {tuple(row.shape)}")
        return row

    def _checked_decode_row(self, token_ids: TokenIds) -> torch.Tensor:
        """Return one validated row without repairing invalid model output."""
        row = self._to_id_row(token_ids)
        invalid = (row < 0) | (row >= int(self._action_vocab_size))
        if torch.any(invalid):
            bad = row[invalid][:5].tolist()
            raise ValueError(
                f"token ids {bad} out of range [0, {int(self._action_vocab_size)})"
            )
        return row

    def _decode_token_ids_to_actions_batch(
        self,
        action_token_ids: Sequence[TokenIds],
        *,
        time_horizon: int | None = None,
        action_dim: int | None = None,
        decode_kwargs: dict | None = None,
    ) -> torch.Tensor:
        """Decode a batch of id rows to actions ``(B, T, D)``.

        The decoder applies the configured action-space layout to each row.
        """
        batch_action_idx: ActionIdxBatch = [
            self._checked_decode_row(a).tolist() for a in action_token_ids
        ]
        (actions_np, _) = self._decode_action_indices(
            batch_action_idx,
            time_horizon=time_horizon,
            action_dim=action_dim,
            decode_kwargs=decode_kwargs,
        )
        return torch.from_numpy(actions_np)

    @property
    def vocab_size(self) -> int:
        return int(self._action_vocab_size)