"""Statistical health checks for the raw entropy source.

These run on the frame hash (the input to the DRBG, not the DRBG output)
to detect a frozen or degraded camera before its entropy is mixed in.
A failing check means the frame is rejected and not used to reseed.

The thresholds are the standard FIPS 140-2 (Annex C) continuous-test
bounds for a 20000-bit sample. They are intentionally conservative: a
healthy camera frame hashed with SHA-256 passes easily; a frozen frame
(all identical bytes) fails hard.
"""

# FIPS 140-2 (Annex C) thresholds for 20000 bits, doubled because we
# count runs of 0s and runs of 1s together (the spec tracks them
# separately with these per-value bounds).
MONOBIT_MIN = 9725
MONOBIT_MAX = 10275
RUNS_MIN = {1: 4630, 2: 2228, 3: 1054, 4: 480, 5: 206, 6: 206}
RUNS_MAX = {1: 5370, 2: 2772, 3: 1446, 4: 768, 5: 418, 6: 418}
LONG_RUN_LIMIT = 26  # max run length of identical bits


class HealthCheckError(Exception):
    """Raised when a frame fails a statistical health check."""


def _bits(data):
    """Yield bits (as 0/1 ints) from a byte string, MSB first."""
    for byte in data:
        for shift in range(7, -1, -1):
            yield (byte >> shift) & 1


def monobit(data):
    """Count of 1-bits. Must be within [MONOBIT_MIN, MONOBIT_MAX]."""
    ones = sum(_bits(data))
    if not (MONOBIT_MIN <= ones <= MONOBIT_MAX):
        raise HealthCheckError(f"monobit: {ones} ones outside [{MONOBIT_MIN}, {MONOBIT_MAX}]")
    return ones


def runs(data):
    """Distribution of runs (consecutive identical bits)."""
    counts = {1: 0, 2: 0, 3: 0, 4: 0, 5: 0, 6: 0}
    prev = None
    run_len = 0
    long_run = 0
    for bit in _bits(data):
        if bit == prev:
            run_len += 1
        else:
            if prev is not None:
                bucket = min(run_len, 6)
                counts[bucket] += 1
                long_run = max(long_run, run_len)
            prev = bit
            run_len = 1
    # final run
    bucket = min(run_len, 6)
    counts[bucket] += 1
    long_run = max(long_run, run_len)

    if long_run >= LONG_RUN_LIMIT:
        raise HealthCheckError(f"runs: long run of {long_run} >= {LONG_RUN_LIMIT}")
    for length, count in counts.items():
        if not (RUNS_MIN[length] <= count <= RUNS_MAX[length]):
            raise HealthCheckError(f"runs: length-{length} count {count} outside [{RUNS_MIN[length]}, {RUNS_MAX[length]}]")
    return counts


def check(data):
    """Run all health checks. Raises HealthCheckError on failure.

    Expects at least 2500 bytes (20000 bits) of input. Shorter inputs
    are rejected because the FIPS thresholds assume a 20000-bit sample.
    """
    if len(data) * 8 < 20000:
        raise HealthCheckError(f"sample too short: {len(data)} bytes (< 2500)")
    monobit(data)
    runs(data)
    return True