"""Strict ``bpe.json`` I/O for both tokenization schemes.

The serialization layer owns only the artifact schema. Set-BPE spatial and temporal
merge semantics stay in :mod:`actioncodec2.setbpe` and plain-BPE merge
semantics in :mod:`actioncodec2.tokenization.bpe_vocabulary`, so training and deployment
cannot drift through two independent implementations of the same algorithm.

``config.set_bpe.scheme`` selects which serializer an artifact carries:
``"set_bpe"`` (the production default) or ``"bpe"``, the FAST-style baseline
that flattens the quantized cell matrix row-major before merging.
"""

from __future__ import annotations

import json
from collections.abc import Mapping, Sequence
from numbers import Integral
from pathlib import Path

import numpy as np

from ..common._validation import active_dimensions
from .bpe_vocabulary import BPEVocabulary
from ..physical.configuration import ActionCodec2Config, ConfigError
from .setbpe_spatial import SpatialVocabulary
from .setbpe_temporal import TemporalVocabulary

BPE_NAME = "bpe.json"


def _integer(value: object, path: str, *, positive: bool = False) -> int:
    if isinstance(value, bool) or not isinstance(value, int):
        raise ConfigError(f"{path} must be an integer")
    minimum = 1 if positive else 0
    if value < minimum:
        qualifier = "positive" if positive else "non-negative"
        raise ConfigError(f"{path} must be {qualifier}")
    return int(value)


def _rules(payload: object, width: int, path: str) -> tuple[tuple[int, ...], ...]:
    """Parse an ordered merge list without accepting lossy coercions."""
    if not isinstance(payload, Sequence) or isinstance(
        payload, (str, bytes, bytearray)
    ):
        raise ConfigError(f"{path} must be a list")
    output = []
    for rank, raw_rule in enumerate(payload):
        rule_path = f"{path}[{rank}]"
        if not isinstance(raw_rule, Sequence) or isinstance(
            raw_rule, (str, bytes, bytearray)
        ):
            raise ConfigError(f"{rule_path} must be a list of {width} integers")
        if len(raw_rule) != width:
            raise ConfigError(f"{rule_path} must contain exactly {width} integers")
        output.append(
            tuple(
                _integer(value, f"{rule_path}[{index}]")
                for index, value in enumerate(raw_rule)
            )
        )
    return tuple(output)


def _unique_json_object(pairs: list[tuple[str, object]]) -> dict[str, object]:
    output: dict[str, object] = {}
    for key, value in pairs:
        if key in output:
            raise ConfigError(f"{BPE_NAME} contains duplicate field {key!r}")
        output[key] = value
    return output


def _validate_rule_references(
    *,
    atom_count: int,
    lookahead: int,
    spatial_merges: tuple[tuple[int, int], ...],
    temporal_merges: tuple[tuple[int, int, int], ...],
) -> None:
    """Give artifact-level errors before constructing vocabulary metadata."""
    for rank, (left, right) in enumerate(spatial_merges):
        child = atom_count + rank
        if left >= child or right >= child:
            raise ConfigError(
                f"spatial_merges[{rank}] must reference already-created tokens"
            )
    spatial_size = atom_count + len(spatial_merges)
    for rank, (left, right, offset) in enumerate(temporal_merges):
        child = spatial_size + rank
        if left >= child or right >= child:
            raise ConfigError(
                f"temporal_merges[{rank}] must reference already-created tokens"
            )
        if not 1 <= offset < lookahead:
            raise ConfigError(
                f"temporal_merges[{rank}] offset must be positive and smaller than lookahead"
            )


def _read_payload(directory: str | Path) -> dict[str, object]:
    """Parse ``bpe.json`` from an artifact directory, rejecting duplicate keys."""
    path = Path(directory) / BPE_NAME
    try:
        return json.loads(
            path.read_text(encoding="utf-8"), object_pairs_hook=_unique_json_object
        )
    except OSError:
        raise
    except json.JSONDecodeError as error:
        raise ConfigError(f"cannot parse {path}: {error}") from error


def _dense_cells(
    cells: np.ndarray,
    bins: Sequence[int],
    active_dims: Sequence[int] | None,
) -> tuple[np.ndarray, tuple[int, ...] | None]:
    """Validate one trajectory and expand it to dense ``[T,D_max]`` cells.

    Compact ``[T,D_active]`` input may provide explicit global ``active_dims``.
    Without it, compact input uses prefix dimensions. Dense ``[T,D_max]`` input
    may mark absent slots with ``-1``.
    """
    raw = np.asarray(cells)
    if raw.dtype.kind not in {"i", "u"}:
        raise ValueError("cells must contain integer indices")
    values = np.asarray(raw, dtype=np.int64)
    if values.ndim != 2 or len(values) < 1 or values.shape[1] < 1:
        raise ValueError("cells must have non-empty shape [T, D_active]")
    if active_dims is None:
        if values.shape[1] > len(bins):
            raise ValueError(f"cells width exceeds codebook dimension {len(bins)}")
        if values.shape[1] < len(bins):
            dimensions = tuple(range(values.shape[1]))
            dense = np.full((len(values), len(bins)), -1, dtype=np.int64)
            dense[:, : values.shape[1]] = values
            values = dense
        else:
            dimensions = None
    else:
        dimensions = active_dimensions(active_dims, len(bins))
        if len(dimensions) != values.shape[1]:
            raise ValueError(
                "active_dims must have one entry per compact cell coordinate"
            )
        if np.any(values < 0):
            raise ValueError("compact cells must not contain -1")
        dense = np.full((len(values), len(bins)), -1, dtype=np.int64)
        dense[:, dimensions] = values
        values = dense
    counts = np.asarray(tuple(int(value) for value in bins), dtype=np.int64)
    if np.any(values < -1) or np.any(values >= counts[None, :]):
        raise ValueError("cells contain an invalid active or absent index")
    return values, dimensions


class SetBPESerializer:
    """Lossless mapping between ``[T,D]`` physical cells and token IDs.

    Merge lists are interpreted in creation order. Token IDs are therefore
    stable across physical save/load cycles, and decoding can infer any positive
    trajectory length directly from stream coverage.
    """

    scheme = "set_bpe"

    def __init__(
        self,
        *,
        bins: Sequence[int],
        lookahead: int,
        spatial_merges: Sequence[Sequence[int]],
        temporal_merges: Sequence[Sequence[int]],
    ) -> None:
        self.lookahead = _integer(lookahead, "lookahead", positive=True)
        self.bins = tuple(
            _integer(value, f"bins[{index}]", positive=True)
            for index, value in enumerate(bins)
        )
        if not self.bins:
            raise ConfigError("bins must contain at least one cell slot")
        self.spatial_merges = _rules(spatial_merges, 2, "spatial_merges")
        self.temporal_merges = _rules(temporal_merges, 3, "temporal_merges")
        _validate_rule_references(
            atom_count=sum(self.bins),
            lookahead=self.lookahead,
            spatial_merges=self.spatial_merges,
            temporal_merges=self.temporal_merges,
        )
        try:
            self.spatial = SpatialVocabulary(self.bins, self.spatial_merges)
            self.vocabulary = TemporalVocabulary(
                self.spatial, self.lookahead, self.temporal_merges
            )
        except ValueError as error:
            raise ConfigError(str(error)) from error
        self.patterns = self.vocabulary.patterns
        self.root = self.vocabulary.root
        self.vocab_size = self.vocabulary.vocab_size
        self.token_cell_counts = np.asarray(
            [
                sum(
                    int(self.spatial.support[spatial_token]).bit_count()
                    for _, spatial_token in self.patterns[token]
                )
                for token in range(self.vocab_size)
            ],
            dtype=np.int64,
        )

    def encode(
        self,
        cells: np.ndarray,
        *,
        active_dims: Sequence[int] | None = None,
    ) -> list[int]:
        """Serialize one non-empty trajectory without emitting absent slots.

        Compact ``[T,D_active]`` input may provide explicit global
        ``active_dims``. Without it, compact input uses prefix dimensions.
        Dense ``[T,D_max]`` input may mark absent slots with ``-1``.
        """
        values, _ = _dense_cells(cells, self.bins, active_dims)
        return self.vocabulary.encode(values)

    def decode(
        self,
        tokens: Sequence[int],
        *,
        active_dims: Sequence[int] | None = None,
    ) -> np.ndarray:
        """Recover dense cells, writing ``-1`` in semantically absent slots."""
        if len(tokens) < 1:
            raise ValueError("token stream must not be empty")
        checked = []
        for index, token in enumerate(tokens):
            if isinstance(token, bool) or not isinstance(token, Integral):
                raise ValueError(f"tokens[{index}] must be an integer")
            value = int(token)
            if not 0 <= value < self.vocab_size:
                raise ValueError(f"tokens[{index}] is outside the vocabulary")
            checked.append(value)
        return self.vocabulary.decode(checked, active_dims=active_dims)

    def to_dict(self) -> dict[str, object]:
        """Return the intentionally minimal, path-independent BPE artifact."""
        return {
            "spatial_merges": [list(rule) for rule in self.spatial_merges],
            "temporal_merges": [list(rule) for rule in self.temporal_merges],
        }

    @classmethod
    def from_dict(
        cls,
        payload: object,
        *,
        bins: Sequence[int],
        lookahead: int,
    ) -> SetBPESerializer:
        if not isinstance(payload, Mapping) or not all(
            isinstance(key, str) for key in payload
        ):
            raise ConfigError(f"{BPE_NAME} must contain one JSON object")
        allowed = {"spatial_merges", "temporal_merges"}
        unknown = sorted(set(payload) - allowed)
        if unknown:
            raise ConfigError(
                f"{BPE_NAME} contains unknown field(s): {', '.join(unknown)}"
            )
        missing = sorted(allowed - set(payload))
        if missing:
            raise ConfigError(f"{BPE_NAME} is missing field(s): {', '.join(missing)}")
        return cls(
            bins=bins,
            lookahead=lookahead,
            spatial_merges=payload["spatial_merges"],
            temporal_merges=payload["temporal_merges"],
        )

    @classmethod
    def from_pretrained(
        cls,
        directory: str | Path,
        *,
        bins: Sequence[int],
        lookahead: int,
    ) -> SetBPESerializer:
        return cls.from_dict(_read_payload(directory), bins=bins, lookahead=lookahead)

    def save_pretrained(self, directory: str | Path) -> Path:
        """Write ``bpe.json`` and return its path."""
        path = Path(directory) / BPE_NAME
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(json.dumps(self.to_dict(), indent=2) + "\n", encoding="utf-8")
        return path


class BPESerializer:
    """Lossless mapping between ``[T,D]`` physical cells and plain-BPE ids.

    The FAST-style baseline: cells are serialized row-major over the active
    support and merged as an ordinary symbol sequence. It shares the primitive
    atoms, the artifact layout, and the total vocabulary budget with
    :class:`SetBPESerializer`, so the two can be fitted on one corpus and
    compared directly.

    ``lookahead`` bounds a token at ``lookahead * len(bins)`` primitive cells,
    the plain-BPE analogue of the Set-BPE rule that a temporal token spans at
    most ``lookahead`` control steps.
    """

    scheme = "bpe"

    def __init__(
        self,
        *,
        bins: Sequence[int],
        lookahead: int,
        merges: Sequence[Sequence[int]],
    ) -> None:
        self.lookahead = _integer(lookahead, "lookahead", positive=True)
        self.bins = tuple(
            _integer(value, f"bins[{index}]", positive=True)
            for index, value in enumerate(bins)
        )
        if not self.bins:
            raise ConfigError("bins must contain at least one cell slot")
        self.merges = _rules(merges, 2, "merges")
        atom_count = sum(self.bins)
        for rank, (left, right) in enumerate(self.merges):
            child = atom_count + rank
            if left >= child or right >= child:
                raise ConfigError(
                    f"merges[{rank}] must reference already-created tokens"
                )
        try:
            self.vocabulary = BPEVocabulary(self.bins, self.merges)
        except ValueError as error:
            raise ConfigError(str(error)) from error
        self.max_token_cells = self.lookahead * len(self.bins)
        self.token_cell_counts = self.vocabulary.token_cells
        oversized = int(np.count_nonzero(self.token_cell_counts > self.max_token_cells))
        if oversized:
            raise ConfigError(
                f"{oversized} merge rule(s) cover more than the {self.max_token_cells} "
                "primitive cells allowed by lookahead"
            )
        self.vocab_size = self.vocabulary.vocab_size

    def encode(
        self,
        cells: np.ndarray,
        *,
        active_dims: Sequence[int] | None = None,
    ) -> list[int]:
        """Serialize one non-empty trajectory without emitting absent slots.

        Input conventions match :meth:`SetBPESerializer.encode`: compact
        ``[T,D_active]`` may carry explicit global ``active_dims`` and
        otherwise uses prefix dimensions, while dense ``[T,D_max]`` marks
        absent slots with ``-1``.
        """
        values, _ = _dense_cells(cells, self.bins, active_dims)
        return self.vocabulary.encode(values)

    def decode(
        self,
        tokens: Sequence[int],
        *,
        active_dims: Sequence[int] | None = None,
    ) -> np.ndarray:
        """Recover dense cells, writing ``-1`` in semantically absent slots."""
        if len(tokens) < 1:
            raise ValueError("token stream must not be empty")
        checked = []
        for index, token in enumerate(tokens):
            if isinstance(token, bool) or not isinstance(token, Integral):
                raise ValueError(f"tokens[{index}] must be an integer")
            value = int(token)
            if not 0 <= value < self.vocab_size:
                raise ValueError(f"tokens[{index}] is outside the vocabulary")
            checked.append(value)
        return self.vocabulary.decode(checked, active_dims=active_dims)

    def to_dict(self) -> dict[str, object]:
        """Return the intentionally minimal, path-independent BPE artifact."""
        return {"merges": [list(rule) for rule in self.merges]}

    @classmethod
    def from_dict(
        cls,
        payload: object,
        *,
        bins: Sequence[int],
        lookahead: int,
    ) -> BPESerializer:
        if not isinstance(payload, Mapping) or not all(
            isinstance(key, str) for key in payload
        ):
            raise ConfigError(f"{BPE_NAME} must contain one JSON object")
        unknown = sorted(set(payload) - {"merges"})
        if unknown:
            raise ConfigError(
                f"{BPE_NAME} contains unknown field(s): {', '.join(unknown)}"
            )
        if "merges" not in payload:
            raise ConfigError(f"{BPE_NAME} is missing field(s): merges")
        return cls(bins=bins, lookahead=lookahead, merges=payload["merges"])

    @classmethod
    def from_pretrained(
        cls,
        directory: str | Path,
        *,
        bins: Sequence[int],
        lookahead: int,
    ) -> BPESerializer:
        return cls.from_dict(_read_payload(directory), bins=bins, lookahead=lookahead)

    def save_pretrained(self, directory: str | Path) -> Path:
        """Write ``bpe.json`` and return its path."""
        path = Path(directory) / BPE_NAME
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(json.dumps(self.to_dict(), indent=2) + "\n", encoding="utf-8")
        return path


SERIALIZERS: dict[str, type[SetBPESerializer] | type[BPESerializer]] = {
    SetBPESerializer.scheme: SetBPESerializer,
    BPESerializer.scheme: BPESerializer,
}


def serializer_class(
    config: ActionCodec2Config,
) -> type[SetBPESerializer] | type[BPESerializer]:
    """Return the serializer the configured tokenization scheme requires."""
    return SERIALIZERS[config.set_bpe.scheme]