"""PROTOCOL template definitions - cryptographic-style transformations.

Three template families define how input x maps to output core(x):
- T1: XOR -> AFFINE -> ROTL
- T2: XOR -> AFFINE -> XOR-ROUND
- T3: XOR -> SBOX -> ROTR

Each template uses parameters derived deterministically from (secret_key, public_seed).
"""

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from purefn.games.protocol.prng import DomainPRNG


class TemplateType(Enum):
    """Template type identifier."""

    T1 = "T1"
    T2 = "T2"
    T3 = "T3"


def rotl8(value: int, shift: int) -> int:
    """Rotate left 8-bit value by shift positions.

    Args:
        value: 8-bit value (0-255)
        shift: Rotation amount (0-7)

    Returns:
        Rotated 8-bit value
    """
    shift = shift & 7  # Ensure shift is in [0, 7]
    return ((value << shift) | (value >> (8 - shift))) & 0xFF


def rotr8(value: int, shift: int) -> int:
    """Rotate right 8-bit value by shift positions.

    Args:
        value: 8-bit value (0-255)
        shift: Rotation amount (0-7)

    Returns:
        Rotated 8-bit value
    """
    shift = shift & 7  # Ensure shift is in [0, 7]
    return ((value >> shift) | (value << (8 - shift))) & 0xFF


@dataclass(frozen=True)
class T1Params:
    """Parameters for Template T1 (XOR -> AFFINE -> ROTL).

    Attributes:
        K: XOR key (0x00-0xFF)
        A: Affine multiplier (odd, 0x01-0xFF)
        B: Affine addend (0x00-0xFF)
        R: Rotation amount (0-7)
    """

    K: int
    A: int
    B: int
    R: int

    def __post_init__(self) -> None:
        """Validate parameters."""
        if not (0 <= self.K <= 0xFF):
            raise ValueError(f"K must be in [0, 255], got {self.K}")
        if not (0x01 <= self.A <= 0xFF) or (self.A & 1) == 0:
            raise ValueError(f"A must be odd and in [1, 255], got {self.A}")
        if not (0 <= self.B <= 0xFF):
            raise ValueError(f"B must be in [0, 255], got {self.B}")
        if not (0 <= self.R <= 7):
            raise ValueError(f"R must be in [0, 7], got {self.R}")


@dataclass(frozen=True)
class T2Params:
    """Parameters for Template T2 (XOR -> AFFINE -> XOR-ROUND).

    Attributes:
        K: XOR key (0x00-0xFF)
        A: Affine multiplier (odd, 0x01-0xFF)
        B: Affine addend (0x00-0xFF)
        R: Rotation amount for K (0-7)
    """

    K: int
    A: int
    B: int
    R: int

    def __post_init__(self) -> None:
        """Validate parameters."""
        if not (0 <= self.K <= 0xFF):
            raise ValueError(f"K must be in [0, 255], got {self.K}")
        if not (0x01 <= self.A <= 0xFF) or (self.A & 1) == 0:
            raise ValueError(f"A must be odd and in [1, 255], got {self.A}")
        if not (0 <= self.B <= 0xFF):
            raise ValueError(f"B must be in [0, 255], got {self.B}")
        if not (0 <= self.R <= 7):
            raise ValueError(f"R must be in [0, 7], got {self.R}")


@dataclass(frozen=True)
class T3Params:
    """Parameters for Template T3 (XOR -> SBOX -> ROTR).

    Attributes:
        K: XOR key (0x00-0xFF)
        B: XOR addend after SBOX (0x00-0xFF)
        R: Rotation amount (0-7)
        S: SBOX index (0-7)
    """

    K: int
    B: int
    R: int
    S: int

    def __post_init__(self) -> None:
        """Validate parameters."""
        if not (0 <= self.K <= 0xFF):
            raise ValueError(f"K must be in [0, 255], got {self.K}")
        if not (0 <= self.B <= 0xFF):
            raise ValueError(f"B must be in [0, 255], got {self.B}")
        if not (0 <= self.R <= 7):
            raise ValueError(f"R must be in [0, 7], got {self.R}")
        if not (0 <= self.S <= 7):
            raise ValueError(f"S must be in [0, 7], got {self.S}")


# Type alias for any template parameters
TemplateParams = T1Params | T2Params | T3Params


def core_t1(x: int, params: T1Params) -> int:
    """Compute Template T1: XOR -> AFFINE -> ROTL.

    Args:
        x: Input byte (0-255)
        params: T1 parameters

    Returns:
        Output byte (0-255)

    Stages:
        m = x XOR K
        a = (A * m + B) mod 256
        r = ROTL8(a, R)
        core(x) = r
    """
    m = x ^ params.K
    a = (params.A * m + params.B) & 0xFF
    r = rotl8(a, params.R)
    return r


def core_t2(x: int, params: T2Params) -> int:
    """Compute Template T2: XOR -> AFFINE -> XOR-ROUND.

    Args:
        x: Input byte (0-255)
        params: T2 parameters

    Returns:
        Output byte (0-255)

    Stages:
        m = x XOR K
        a = (A * m + B) mod 256
        r = a XOR ROTL8(K, R)
        core(x) = r
    """
    m = x ^ params.K
    a = (params.A * m + params.B) & 0xFF
    r = a ^ rotl8(params.K, params.R)
    return r


def core_t3(x: int, params: T3Params, sbox: list[int]) -> int:
    """Compute Template T3: XOR -> SBOX -> ROTR.

    Args:
        x: Input byte (0-255)
        params: T3 parameters
        sbox: The S-box permutation (256 bytes) for index S

    Returns:
        Output byte (0-255)

    Stages:
        m = x XOR K
        a = SBOX[S][m] XOR B
        r = ROTR8(a, R)
        core(x) = r
    """
    m = x ^ params.K
    a = sbox[m] ^ params.B
    r = rotr8(a, params.R)
    return r


def stage_m(x: int, k: int) -> int:
    """Compute stage 1 (m) for any template.

    Args:
        x: Input byte
        k: XOR key K

    Returns:
        m = x XOR K
    """
    return x ^ k


def stage_a_t1t2(m: int, a_param: int, b: int) -> int:
    """Compute stage 2 (a) for T1 and T2 templates.

    Args:
        m: Stage 1 output
        a_param: Affine multiplier A
        b: Affine addend B

    Returns:
        a = (A * m + B) mod 256
    """
    return (a_param * m + b) & 0xFF


def stage_a_t3(m: int, sbox: list[int], b: int) -> int:
    """Compute stage 2 (a) for T3 template.

    Args:
        m: Stage 1 output
        sbox: S-box permutation
        b: XOR addend B

    Returns:
        a = SBOX[m] XOR B
    """
    return sbox[m] ^ b


def generate_sbox(prng: DomainPRNG, index: int) -> list[int]:
    """Generate an S-box permutation using Fisher-Yates shuffle.

    Uses domain tag "sbox:<index>" for deterministic generation.

    Args:
        prng: Domain PRNG instance
        index: S-box index (0-7)

    Returns:
        Permutation of [0, 1, 2, ..., 255]
    """
    domain = f"sbox:{index}"
    arr = list(range(256))

    # Fisher-Yates shuffle: for j = 255 down to 1
    for j in range(255, 0, -1):
        k = prng.rand_int(domain, 0, j)
        arr[j], arr[k] = arr[k], arr[j]

    return arr


def sample_template_type(prng: DomainPRNG, allowed: set[TemplateType]) -> TemplateType:
    """Sample a template type uniformly from allowed set.

    Args:
        prng: Domain PRNG instance
        allowed: Set of allowed template types

    Returns:
        Selected template type
    """
    templates = sorted(allowed, key=lambda t: t.value)  # Deterministic order
    idx = prng.rand_int("template", 0, len(templates) - 1)
    return templates[idx]


def sample_t1_params(prng: DomainPRNG) -> T1Params:
    """Sample T1 parameters from PRNG.

    Args:
        prng: Domain PRNG instance

    Returns:
        T1Params instance
    """
    k = prng.rand_u8("params")
    a = prng.rand_u8("params") | 0x01  # Ensure odd
    b = prng.rand_u8("params")
    r = prng.rand_int("params", 0, 7)
    return T1Params(K=k, A=a, B=b, R=r)


def sample_t2_params(prng: DomainPRNG) -> T2Params:
    """Sample T2 parameters from PRNG.

    Args:
        prng: Domain PRNG instance

    Returns:
        T2Params instance
    """
    k = prng.rand_u8("params")
    a = prng.rand_u8("params") | 0x01  # Ensure odd
    b = prng.rand_u8("params")
    r = prng.rand_int("params", 0, 7)
    return T2Params(K=k, A=a, B=b, R=r)


def sample_t3_params(prng: DomainPRNG) -> T3Params:
    """Sample T3 parameters from PRNG.

    Args:
        prng: Domain PRNG instance

    Returns:
        T3Params instance
    """
    k = prng.rand_u8("params")
    b = prng.rand_u8("params")
    r = prng.rand_int("params", 0, 7)
    s = prng.rand_int("params", 0, 7)
    return T3Params(K=k, B=b, R=r, S=s)
