"""PROTOCOL game engine - core game logic.

A deterministic CLI game where teams of agents must infer a hidden
cryptographic-style rule from partial observations.

This module provides the core engine logic. The full game state management,
CLI commands, and logging are handled separately.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from purefn.games.protocol.drift import (
    DriftEvent,
    DriftParam,
    DriftSchedule,
    apply_drift_operation,
    generate_drift_schedule,
    get_pending_drift_events,
)
from purefn.games.protocol.encoding import encode_byte
from purefn.games.protocol.prng import DomainPRNG
from purefn.games.protocol.profiles import PROFILE_CONFIGS, Profile, ProfileConfig
from purefn.games.protocol.templates import (
    T1Params,
    T2Params,
    T3Params,
    TemplateParams,
    TemplateType,
    core_t1,
    core_t2,
    core_t3,
    generate_sbox,
    sample_t1_params,
    sample_t2_params,
    sample_t3_params,
    sample_template_type,
    stage_a_t1t2,
    stage_a_t3,
    stage_m,
)

# Re-export for backward compatibility
__all__ = [
    "PROFILE_CONFIGS",
    "Profile",
    "ProfileConfig",
    "ProtocolCore",
    "ProtocolEngine",
    "ProtocolState",
    "create_protocol_core",
    "create_protocol_core_at_turn",
    "serialize_oracle",
    "deserialize_oracle",
    "get_protocol_core_from_oracle",
]


@dataclass
class ProtocolState:
    """Internal protocol state (includes hidden values).

    This is the full state including hidden parameters that are NOT
    exposed to agents. Only a filtered view is serialized to the state file.
    """

    # Public fields (visible in state file)
    seed: int
    profile: Profile
    turn: int
    turn_limit: int
    team: dict[str, str]  # scout_id -> station_id
    # stations, map, gates_required, gates_unlocked defined elsewhere

    # Hidden fields (oracle-only)
    template_type: TemplateType
    params: TemplateParams
    sboxes: list[list[int]]  # 8 S-boxes for T3
    drift_schedule: list[int]  # Turn numbers when drift occurs
    current_epoch: int  # Current epoch (drift count)


@dataclass
class ProtocolCore:
    """Core protocol computation engine.

    Handles template evaluation and parameter derivation.
    Does not manage game state, maps, or movement.
    """

    prng: DomainPRNG
    template_type: TemplateType
    params: TemplateParams
    sboxes: list[list[int]] = field(default_factory=list)

    def compute_core(self, x: int) -> int:
        """Compute core(x) for the current template and parameters.

        Args:
            x: Input byte (0-255)

        Returns:
            Output byte (0-255)
        """
        if self.template_type == TemplateType.T1:
            assert isinstance(self.params, T1Params)
            return core_t1(x, self.params)
        elif self.template_type == TemplateType.T2:
            assert isinstance(self.params, T2Params)
            return core_t2(x, self.params)
        elif self.template_type == TemplateType.T3:
            assert isinstance(self.params, T3Params)
            sbox = self.sboxes[self.params.S]
            return core_t3(x, self.params, sbox)
        else:
            raise ValueError(f"Unknown template type: {self.template_type}")

    def compute_gate_output(self, x: int, salt: int) -> int:
        """Compute gate output y = core(x) XOR salt.

        Args:
            x: Challenge input byte
            salt: Per-gate salt byte

        Returns:
            Expected response byte
        """
        return self.compute_core(x) ^ salt

    def compute_stage_m(self, x: int) -> int:
        """Compute stage 1 output (m = x XOR K).

        Args:
            x: Input byte

        Returns:
            Stage 1 output
        """
        if isinstance(self.params, (T1Params, T2Params)):
            return stage_m(x, self.params.K)
        elif isinstance(self.params, T3Params):
            return stage_m(x, self.params.K)
        raise ValueError(f"Unknown params type: {type(self.params)}")

    def compute_stage_a(self, x: int) -> int:
        """Compute stage 2 output.

        Args:
            x: Input byte

        Returns:
            Stage 2 output
        """
        m = self.compute_stage_m(x)
        if isinstance(self.params, T1Params):
            return stage_a_t1t2(m, self.params.A, self.params.B)
        elif isinstance(self.params, T2Params):
            return stage_a_t1t2(m, self.params.A, self.params.B)
        elif isinstance(self.params, T3Params):
            sbox = self.sboxes[self.params.S]
            return stage_a_t3(m, sbox, self.params.B)
        raise ValueError(f"Unknown params type: {type(self.params)}")

    def compute_stage_r(self, x: int) -> int:
        """Compute stage 3 output (final core output).

        Args:
            x: Input byte

        Returns:
            Stage 3 output (same as core(x))
        """
        return self.compute_core(x)

    def get_params_dict(self) -> dict[str, Any]:
        """Get parameters as a dictionary for logging.

        Returns:
            Dictionary with encoded parameter values
        """
        if isinstance(self.params, T1Params):
            return {
                "K": encode_byte(self.params.K),
                "A": encode_byte(self.params.A),
                "B": encode_byte(self.params.B),
                "R": self.params.R,
            }
        elif isinstance(self.params, T2Params):
            return {
                "K": encode_byte(self.params.K),
                "A": encode_byte(self.params.A),
                "B": encode_byte(self.params.B),
                "R": self.params.R,
            }
        elif isinstance(self.params, T3Params):
            return {
                "K": encode_byte(self.params.K),
                "B": encode_byte(self.params.B),
                "R": self.params.R,
                "S": self.params.S,
            }
        raise ValueError(f"Unknown params type: {type(self.params)}")


def create_protocol_core(secret_key: str, public_seed: int, profile: Profile) -> ProtocolCore:
    """Create a protocol core from secret key and seed.

    Args:
        secret_key: Secret key string
        public_seed: Public seed integer
        profile: Difficulty profile

    Returns:
        Initialized ProtocolCore
    """
    prng = DomainPRNG(secret_key, public_seed)
    config = PROFILE_CONFIGS[profile]

    # Sample template type
    template_type = sample_template_type(prng, config.allowed_templates)

    # Sample parameters based on template type
    if template_type == TemplateType.T1:
        params: TemplateParams = sample_t1_params(prng)
    elif template_type == TemplateType.T2:
        params = sample_t2_params(prng)
    else:  # T3
        params = sample_t3_params(prng)

    # Generate S-boxes (needed for T3, but generate all 8 for consistency)
    sboxes = [generate_sbox(prng, i) for i in range(8)]

    return ProtocolCore(
        prng=prng,
        template_type=template_type,
        params=params,
        sboxes=sboxes,
    )


def _apply_drift_t1(params: T1Params, drift_param: DriftParam, delta: int) -> T1Params:
    """Apply drift to T1 parameters."""
    vals = {"K": params.K, "A": params.A, "B": params.B, "R": params.R}
    if drift_param.value in vals:
        vals[drift_param.value] = apply_drift_operation(drift_param, vals[drift_param.value], delta)
    return T1Params(K=vals["K"], A=vals["A"], B=vals["B"], R=vals["R"])


def _apply_drift_t2(params: T2Params, drift_param: DriftParam, delta: int) -> T2Params:
    """Apply drift to T2 parameters."""
    vals = {"K": params.K, "A": params.A, "B": params.B, "R": params.R}
    if drift_param.value in vals:
        vals[drift_param.value] = apply_drift_operation(drift_param, vals[drift_param.value], delta)
    return T2Params(K=vals["K"], A=vals["A"], B=vals["B"], R=vals["R"])


def _apply_drift_t3(params: T3Params, drift_param: DriftParam, delta: int) -> T3Params:
    """Apply drift to T3 parameters."""
    vals = {"K": params.K, "B": params.B, "R": params.R, "S": params.S}
    if drift_param.value in vals:
        vals[drift_param.value] = apply_drift_operation(drift_param, vals[drift_param.value], delta)
    return T3Params(K=vals["K"], B=vals["B"], R=vals["R"], S=vals["S"])


def _apply_drift_to_params(
    params: TemplateParams,
    drift_param: DriftParam,
    delta: int,
) -> TemplateParams:
    """Apply a single drift operation to parameters.

    Creates a new params object with the updated value.

    Args:
        params: Current parameters
        drift_param: Which parameter to change
        delta: Delta value

    Returns:
        New parameters with drift applied
    """
    if isinstance(params, T1Params):
        return _apply_drift_t1(params, drift_param, delta)
    elif isinstance(params, T2Params):
        return _apply_drift_t2(params, drift_param, delta)
    elif isinstance(params, T3Params):
        return _apply_drift_t3(params, drift_param, delta)
    raise ValueError(f"Unknown params type: {type(params)}")


def create_protocol_core_at_turn(
    secret_key: str,
    public_seed: int,
    profile: Profile,
    current_turn: int,
    turn_limit: int,
) -> ProtocolCore:
    """Create a protocol core with drift applied up to current turn.

    This creates a fresh PRNG, samples initial parameters, generates the
    drift schedule, and applies all drift events up to the current turn.

    Args:
        secret_key: Secret key string
        public_seed: Public seed integer
        profile: Difficulty profile
        current_turn: Current game turn (drift applied up to this turn)
        turn_limit: Maximum turn limit (for drift scheduling)

    Returns:
        ProtocolCore with drift-adjusted parameters
    """
    # Create base core with initial parameters
    core = create_protocol_core(secret_key, public_seed, profile)

    # Generate drift schedule (must use fresh PRNG to match init)
    # Note: PRNG state has advanced after parameter sampling, so drift
    # schedule generation uses continued PRNG state
    drift_schedule = generate_drift_schedule(core.prng, profile, core.template_type, turn_limit)

    # Apply all drift events up to current turn
    params = core.params
    for event in get_pending_drift_events(drift_schedule, 0, current_turn):
        for drift_param in event.params_changed:
            delta = event.deltas[drift_param]
            params = _apply_drift_to_params(params, drift_param, delta)

    # Return core with updated parameters
    return ProtocolCore(
        prng=core.prng,
        template_type=core.template_type,
        params=params,
        sboxes=core.sboxes,
    )


def serialize_oracle(
    core: ProtocolCore,
    drift_schedule: DriftSchedule,
) -> dict[str, Any]:
    """Serialize ProtocolCore and drift schedule to JSON-serializable dict.

    This creates an "oracle" that can be embedded in the state file,
    allowing agents to play without needing the secret key.

    Args:
        core: Protocol core with template type, params, and S-boxes
        drift_schedule: Complete drift schedule for the game

    Returns:
        Dictionary containing all hidden state needed for gameplay
    """
    # Serialize parameters based on template type
    if isinstance(core.params, T1Params):
        params_data = {
            "K": core.params.K,
            "A": core.params.A,
            "B": core.params.B,
            "R": core.params.R,
        }
    elif isinstance(core.params, T2Params):
        params_data = {
            "K": core.params.K,
            "A": core.params.A,
            "B": core.params.B,
            "R": core.params.R,
        }
    else:  # T3Params
        params_data = {
            "K": core.params.K,
            "B": core.params.B,
            "R": core.params.R,
            "S": core.params.S,
        }

    # Serialize drift events
    drift_events = []
    for event in drift_schedule.events:
        drift_events.append(
            {
                "turn": event.turn,
                "params_changed": [p.value for p in event.params_changed],
                "deltas": {p.value: delta for p, delta in event.deltas.items()},
            }
        )

    return {
        "template_type": core.template_type.value,
        "params": params_data,
        "sboxes": core.sboxes,
        "drift_events": drift_events,
    }


def deserialize_oracle(oracle_data: dict[str, Any]) -> tuple[ProtocolCore, DriftSchedule]:
    """Deserialize oracle data back to ProtocolCore and drift schedule.

    Args:
        oracle_data: Dictionary from serialize_oracle()

    Returns:
        Tuple of (ProtocolCore, DriftSchedule)
    """
    # Reconstruct template type
    template_type = TemplateType(oracle_data["template_type"])

    # Reconstruct parameters
    params_data = oracle_data["params"]
    if template_type == TemplateType.T1:
        params: TemplateParams = T1Params(
            K=params_data["K"],
            A=params_data["A"],
            B=params_data["B"],
            R=params_data["R"],
        )
    elif template_type == TemplateType.T2:
        params = T2Params(
            K=params_data["K"],
            A=params_data["A"],
            B=params_data["B"],
            R=params_data["R"],
        )
    else:  # T3
        params = T3Params(
            K=params_data["K"],
            B=params_data["B"],
            R=params_data["R"],
            S=params_data["S"],
        )

    # Reconstruct S-boxes
    sboxes = oracle_data["sboxes"]

    # Create a dummy PRNG (not used after deserialization)
    # We use a placeholder since we won't be generating new random values
    dummy_prng = DomainPRNG("", 0)

    core = ProtocolCore(
        prng=dummy_prng,
        template_type=template_type,
        params=params,
        sboxes=sboxes,
    )

    # Reconstruct drift schedule
    drift_events = []
    for event_data in oracle_data.get("drift_events", []):
        params_changed = [DriftParam(p) for p in event_data["params_changed"]]
        deltas = {DriftParam(k): v for k, v in event_data["deltas"].items()}
        drift_events.append(
            DriftEvent(
                turn=event_data["turn"],
                params_changed=params_changed,
                deltas=deltas,
            )
        )

    drift_schedule = DriftSchedule(events=drift_events)

    return core, drift_schedule


def get_protocol_core_from_oracle(
    oracle_data: dict[str, Any],
    current_turn: int,
) -> ProtocolCore:
    """Get ProtocolCore at a specific turn using embedded oracle data.

    This is the oracle-free equivalent of create_protocol_core_at_turn().
    It deserializes the oracle and applies drift events up to the current turn.

    Args:
        oracle_data: Dictionary from serialize_oracle() embedded in state file
        current_turn: Current game turn (drift applied up to this turn)

    Returns:
        ProtocolCore with drift-adjusted parameters
    """
    core, drift_schedule = deserialize_oracle(oracle_data)

    # Apply all drift events up to current turn
    params = core.params
    for event in get_pending_drift_events(drift_schedule, 0, current_turn):
        for drift_param in event.params_changed:
            delta = event.deltas[drift_param]
            params = _apply_drift_to_params(params, drift_param, delta)

    # Return core with updated parameters
    return ProtocolCore(
        prng=core.prng,
        template_type=core.template_type,
        params=params,
        sboxes=core.sboxes,
    )


class ProtocolEngine:
    """Envelope-based adapter for the protocol engine.

    Implements the GameEngine protocol for use with the global registry and
    the `purefn run` CLI entrypoint.

    Note: The full game engine (with state management, maps, movement, etc.)
    will be implemented in later phases. This provides the foundation.
    """

    def dispatch(self, op: str, global_token: str, args: dict[str, Any]) -> Any:
        """Dispatch an operation based on envelope fields.

        Args:
            op: Operation name
            global_token: Current state token
            args: Operation arguments

        Returns:
            PureResult with operation outcome
        """
        # Placeholder - full implementation in later phases
        raise NotImplementedError(f"Protocol operation '{op}' not yet implemented")
