"""Username <-> 256-bit videoseal payload codec.

Layout (32 bytes total = 256 bits):
    bytes [0:16]  ASCII username, NUL-padded right (max 16 chars).
    bytes [16:32] SHA-256(name_bytes)[:16] integrity tag.

The hash tail is the only thing protecting the puzzle from accidental matches:
a noisy extraction of an un-watermarked clip would otherwise have a non-trivial
probability of producing the admin name by chance.
"""

from __future__ import annotations

import hashlib
import string
from typing import Optional

import numpy as np

NAME_BYTES = 16
TAG_BYTES = 16
PAYLOAD_BYTES = NAME_BYTES + TAG_BYTES
PAYLOAD_BITS = PAYLOAD_BYTES * 8

_PRINTABLE = set(bytes(string.ascii_letters + string.digits + "_-", "ascii"))


def _name_to_bytes(username: str) -> bytes:
    if not username:
        raise ValueError("username must not be empty")
    if len(username) > NAME_BYTES:
        raise ValueError(f"username must be <= {NAME_BYTES} chars")
    raw = username.encode("ascii")
    if any(b not in _PRINTABLE for b in raw):
        raise ValueError("username must be ASCII alnum/_/-")
    return raw.ljust(NAME_BYTES, b"\x00")


def _tag(name_bytes: bytes) -> bytes:
    return hashlib.sha256(name_bytes).digest()[:TAG_BYTES]


def encode(username: str) -> np.ndarray:
    """Return a length-256 uint8 array of bits (0/1)."""
    name_bytes = _name_to_bytes(username)
    payload = name_bytes + _tag(name_bytes)
    return np.unpackbits(np.frombuffer(payload, dtype=np.uint8))


def _bits_to_bytes(bits: np.ndarray) -> bytes:
    bits = np.asarray(bits, dtype=np.uint8).reshape(-1)
    if bits.size != PAYLOAD_BITS:
        raise ValueError(f"expected {PAYLOAD_BITS} bits, got {bits.size}")
    return bytes(np.packbits(bits))


def _strip_padding(name_bytes: bytes) -> Optional[str]:
    stripped = name_bytes.rstrip(b"\x00")
    if not stripped or any(b not in _PRINTABLE for b in stripped):
        return None
    return stripped.decode("ascii")


def decode_strict(bits: np.ndarray) -> Optional[str]:
    """Return the username iff the hash tag matches exactly."""
    payload = _bits_to_bytes(bits)
    name_bytes, tag = payload[:NAME_BYTES], payload[NAME_BYTES:]
    if _tag(name_bytes) != tag:
        return None
    return _strip_padding(name_bytes)


def _hamming(a: bytes, b: bytes) -> int:
    return sum(bin(x ^ y).count("1") for x, y in zip(a, b))


def decode_fuzzy(bits: np.ndarray, hamming_tol: int = 12) -> Optional[str]:
    """Strict first; on failure, recompute the tag from a stripped candidate
    and accept if the recovered tag is within `hamming_tol` bits of the
    expected tag. The name half must already decode to printable ASCII --
    we never invent characters."""
    strict = decode_strict(bits)
    if strict is not None:
        return strict
    payload = _bits_to_bytes(bits)
    name_bytes, recovered_tag = payload[:NAME_BYTES], payload[NAME_BYTES:]
    candidate = _strip_padding(name_bytes)
    if candidate is None:
        return None
    expected_tag = _tag(name_bytes)
    if _hamming(expected_tag, recovered_tag) <= hamming_tol:
        return candidate
    return None


if __name__ == "__main__":
    names = ["Marcel", "Mike", "Marley", "Mira", "Maxine",
             "alice", "bob", "charlie", "diane", "eric"]
    for n in names:
        bits = encode(n)
        rt = decode_strict(bits)
        ok = "PASS" if rt == n else "FAIL"
        print(f"{ok} strict roundtrip: {n!r} -> {rt!r}")

    rng = np.random.default_rng(0)
    bits = encode("Marcel").copy()
    flip_idx = rng.choice(np.arange(NAME_BYTES * 8, PAYLOAD_BITS), size=8, replace=False)
    bits[flip_idx] ^= 1
    rt = decode_fuzzy(bits)
    print(f"{'PASS' if rt == 'Marcel' else 'FAIL'} fuzzy roundtrip with 8 tag bits flipped -> {rt!r}")

    bits = encode("Marcel").copy()
    bits[0:8] ^= 1
    rt = decode_fuzzy(bits)
    print(f"{'PASS' if rt is None or rt == 'Marcel' else 'FAIL'} name-half flip handled -> {rt!r}")
