from __future__ import annotations

from dataclasses import dataclass
from random import Random

from tea3.tea3 import Tea3


def rand_key(rng: Random) -> list[int]:
    return [rng.getrandbits(8) for _ in range(10)]


def rand_frame_number(rng: Random) -> int:
    return rng.getrandbits(32)


@dataclass
class BiasResult:
    samples: int
    matches: int
    p_match: float
    bias: float
    estimated_needed_samples: float


def estimate_black_box_output_bias(
    *,
    key: list[int],
    output_byte_index: int = 0,
    output_bit: int = 0,
    samples: int = 10000,
    seed: int = 1,
) -> BiasResult:
    rng = Random(seed)
    matches = 0

    for _ in range(samples):
        frame_number = rand_frame_number(rng)
        tea = Tea3(frame_number, key)

        for _ in range(output_byte_index + 1):
            tea.next_byte()

        out_byte = tea.iv_view() >> 56

        bit = (out_byte >> output_bit) & 1
        if bit == 0:
            matches += 1

    p_match = matches / samples
    bias = abs(p_match - 0.5)
    needed = float("inf") if bias == 0 else 1.0 / (bias * bias)

    return BiasResult(
        samples=samples,
        matches=matches,
        p_match=p_match,
        bias=bias,
        estimated_needed_samples=needed,
    )


def main():
    rng = Random(123)
    key = rand_key(rng)

    res = estimate_black_box_output_bias(
        key=key,
        output_byte_index=0,
        output_bit=0,
        samples=20000,
        seed=123,
    )
    print("black-box output bias")
    print("expected non-biased p_match: 1/2")
    print(res)


if __name__ == "__main__":
    main()