from __future__ import annotations

from dataclasses import dataclass
from itertools import combinations, product
from random import sample
from typing import Sequence

from tea3.tea3 import Tea3

MASK32 = 0xFFFFFFFF


@dataclass(frozen=True)
class CubeResult:
    cube: tuple[int, ...]
    cube_size: int
    cube_sum: int
    output_byte: int
    output_bit: int


def get_output_bit(
    frame_number: int,
    key_register: Sequence[int],
    output_byte: int = 0,
    output_bit: int = 0,
) -> int:
    """
    Evaluate TEA3 as a black box and return one output bit from the keystream.
    """
    tea = Tea3(frame_number=frame_number, key_register=key_register)
    ks = tea.keystream(output_byte + 1)
    return (ks[output_byte] >> output_bit) & 1


def set_bits(value: int, bit_indices: Sequence[int], bits: Sequence[int]) -> int:
    """
    Set selected bit positions of `value` according to `bits`.
    Bit index 0 is the least significant bit.
    """
    if len(bit_indices) != len(bits):
        raise ValueError("bit_indices and bits must have the same length")

    x = value & MASK32
    for idx, bit in zip(bit_indices, bits):
        if bit not in (0, 1):
            raise ValueError("bits must be 0 or 1")
        if bit:
            x |= 1 << idx
        else:
            x &= ~(1 << idx)
    return x & MASK32


def cube_sum_tea3(
    cube_bits: Sequence[int],
    base_frame_number: int,
    key_register: Sequence[int],
    output_byte: int = 0,
    output_bit: int = 0,
) -> int:
    """
    Compute the cube sum directly on TEA3 by querying the cipher on all assignments of the chosen cube bits.

    The other IV/frame bits are taken from `base_frame_number`.
    """
    acc = 0
    for assignment in product((0, 1), repeat=len(cube_bits)):
        frame_number = set_bits(base_frame_number, cube_bits, assignment)
        acc ^= get_output_bit(
            frame_number=frame_number,
            key_register=key_register,
            output_byte=output_byte,
            output_bit=output_bit,
        )
    return acc & 1


def search_cubes_exhaustive_tea3(
    public_bits: Sequence[int],
    cube_size: int,
    base_frame_number: int,
    key_register: Sequence[int],
    output_byte: int = 0,
    output_bit: int = 0,
    limit: int = 20,
    keep_zero: bool = False,
) -> list[CubeResult]:
    if cube_size < 0:
        raise ValueError("cube_size must be non-negative")
    if cube_size > len(public_bits):
        raise ValueError("cube_size cannot exceed the number of public bits")

    results: list[CubeResult] = []

    for cube in combinations(public_bits, cube_size):
        s = cube_sum_tea3(
            cube_bits=cube,
            base_frame_number=base_frame_number,
            key_register=key_register,
            output_byte=output_byte,
            output_bit=output_bit,
        )
        if s == 0 and not keep_zero:
            continue

        results.append(
            CubeResult(
                cube=tuple(cube),
                cube_size=cube_size,
                cube_sum=s,
                output_byte=output_byte,
                output_bit=output_bit,
            )
        )

        if len(results) >= limit:
            break

    return results


def search_cubes_random_tea3(
    public_bits: Sequence[int],
    cube_size: int,
    base_frame_number: int,
    key_register: Sequence[int],
    samples: int = 1000,
    output_byte: int = 0,
    output_bit: int = 0,
    limit: int = 20,
    keep_zero: bool = False,
) -> list[CubeResult]:
    if cube_size < 0:
        raise ValueError("cube_size must be non-negative")
    if cube_size > len(public_bits):
        raise ValueError("cube_size cannot exceed the number of public bits")

    results: list[CubeResult] = []
    seen: set[tuple[int, ...]] = set()

    idxs = list(range(len(public_bits)))
    for _ in range(samples):
        cube_idx = tuple(sorted(sample(idxs, cube_size)))
        if cube_idx in seen:
            continue
        seen.add(cube_idx)

        cube = tuple(public_bits[i] for i in cube_idx)
        s = cube_sum_tea3(
            cube_bits=cube,
            base_frame_number=base_frame_number,
            key_register=key_register,
            output_byte=output_byte,
            output_bit=output_bit,
        )
        if s == 0 and not keep_zero:
            continue

        results.append(
            CubeResult(
                cube=cube,
                cube_size=cube_size,
                cube_sum=s,
                output_byte=output_byte,
                output_bit=output_bit,
            )
        )

        if len(results) >= limit:
            break

    return results


def run_cube_attack_offline_tea3(
    key_register: Sequence[int],
    base_frame_number: int,
    public_bits: Sequence[int] = tuple(range(32)),
    cube_size: int = 4,
    mode: str = "random",
    samples: int = 2000,
    limit: int = 20,
    output_byte: int = 0,
    output_bit: int = 0,
    keep_zero: bool = False,
) -> list[CubeResult]:
    """
    Blackbox offline cube search against TEA3.

    This version does not use symbolic polynomials. It evaluates the cipher on all cube assignments and returns cubes whose cube sum is nonzero by default.

    Note:
        This is an empirical offline phase. It can identify candidate cubes, but it does not compute the exact superpoly degree.
    """
    print("=" * 50)
    print("TEA3 black-box cube search")
    print(f"Base frame number: 0x{base_frame_number:08x}")
    print(f"Public bits: {len(public_bits)}")
    print(f"Cube size: {cube_size}")
    print(f"Mode: {mode}")
    print(f"Output byte/bit: {output_byte}/{output_bit}")
    print("=" * 50)

    if mode == "exhaustive":
        results = search_cubes_exhaustive_tea3(
            public_bits=public_bits,
            cube_size=cube_size,
            base_frame_number=base_frame_number,
            key_register=key_register,
            output_byte=output_byte,
            output_bit=output_bit,
            limit=limit,
            keep_zero=keep_zero,
        )
    elif mode == "random":
        results = search_cubes_random_tea3(
            public_bits=public_bits,
            cube_size=cube_size,
            base_frame_number=base_frame_number,
            key_register=key_register,
            samples=samples,
            output_byte=output_byte,
            output_bit=output_bit,
            limit=limit,
            keep_zero=keep_zero,
        )
    else:
        raise ValueError("mode must be 'random' or 'exhaustive'")

    if not results:
        print("No candidate cubes found.")
        return []

    print(f"Found {len(results)} candidate cube(s):")
    for i, res in enumerate(results, 1):
        cube_str = " ".join(f"b{b}" for b in res.cube)
        print("-" * 50)
        print(f"[{i}] cube = {cube_str}")
        print(f"    cube_sum = {res.cube_sum}")

    return results