"""Strict execution settings kept separate from Yunmo-Next model geometry.

The type records how one training process executes an already-defined model.
It intentionally has no context-to-stage inference and no production route
table; stage recipes must select an engine explicitly.
"""

from __future__ import annotations

from dataclasses import asdict, dataclass, fields
from typing import Any, Mapping


@dataclass(frozen=True)
class TrainingEngineConfig:
    """Stage-independent training execution schema."""

    feed_forward_chunk_tokens: int = 0
    feed_forward_serial_backward: bool = False
    recurrent_segment_tokens: int = 0
    recurrent_packed_documents: bool = False
    recurrent_official_forward_serial_backward: bool = False
    activation_checkpointing: bool = True
    activation_checkpoint_mode: str = "non_reentrant"
    cross_entropy_backend: str = "liger"

    def __post_init__(self) -> None:
        for name in ("feed_forward_chunk_tokens", "recurrent_segment_tokens"):
            value = getattr(self, name)
            if type(value) is not int:
                raise TypeError(f"{name} must be int, got {type(value).__name__}")
            if value < 0:
                raise ValueError(f"{name} must be non-negative")
        for name in (
            "feed_forward_serial_backward",
            "recurrent_packed_documents",
            "recurrent_official_forward_serial_backward",
            "activation_checkpointing",
        ):
            value = getattr(self, name)
            if type(value) is not bool:
                raise TypeError(f"{name} must be bool, got {type(value).__name__}")
        for name in ("activation_checkpoint_mode", "cross_entropy_backend"):
            value = getattr(self, name)
            if type(value) is not str:
                raise TypeError(f"{name} must be str, got {type(value).__name__}")

        if (
            self.feed_forward_serial_backward
            and self.feed_forward_chunk_tokens == 0
        ):
            raise ValueError(
                "serial feed-forward backward requires a positive chunk size"
            )
        if (
            self.recurrent_segment_tokens
            and self.recurrent_segment_tokens % 64
        ):
            raise ValueError("recurrent segment size must be a multiple of 64")
        if self.recurrent_packed_documents and self.recurrent_segment_tokens == 0:
            raise ValueError(
                "packed recurrent execution requires a positive segment size"
            )
        if (
            self.recurrent_official_forward_serial_backward
            and not self.recurrent_packed_documents
        ):
            raise ValueError(
                "official-forward serial-backward recurrent execution "
                "requires packed documents"
            )
        if self.activation_checkpoint_mode != "non_reentrant":
            raise ValueError(
                "Yunmo-Next training uses non-reentrant activation checkpointing"
            )
        if self.cross_entropy_backend not in {"liger", "pytorch"}:
            raise ValueError(
                "cross_entropy_backend must be liger or pytorch"
            )

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)

    @classmethod
    def from_dict(cls, value: Mapping[str, Any]) -> "TrainingEngineConfig":
        if not isinstance(value, Mapping):
            raise TypeError("TrainingEngineConfig must be loaded from a mapping")
        expected = {field.name for field in fields(cls)}
        actual = set(value)
        if actual != expected:
            raise ValueError(
                "TrainingEngineConfig field set changed: "
                f"missing={sorted(expected - actual)}, "
                f"extra={sorted(actual - expected)}"
            )
        return cls(**dict(value))


TRAINING_ENGINE_CONFIG_FIELDS = tuple(
    field.name for field in fields(TrainingEngineConfig)
)