"""Extract the hidden username from a watermarked video.
python extract.py <video.mp4>

The 256-bit videoseal payload is 16 bytes of ASCII username
(NUL-padded) + 16 bytes of SHA-256(name)[:16] (hashed over the padded 16-byte name).
We recover the bits, split them, and only trust the name if the hash matches (so an unwatermarked clip
prints `None` instead of a made-up name).
"""

import hashlib
import shutil
import sys
from pathlib import Path

import numpy as np
import torch
import torchvision.io
import videoseal


def stage_videoseal_configs() -> None:
    """videoseal.load() reads `videoseal/cards/*.yaml` and
    `configs/attenuation.yaml` from the CWD; copy them in if missing so this
    runs anywhere, not just inside a videoseal checkout."""
    if not Path("videoseal/cards").exists():
        shutil.copytree(Path(videoseal.__file__).resolve().parent / "cards", "videoseal/cards")
    attn = Path("configs/attenuation.yaml")
    if not attn.exists():
        attn.parent.mkdir(parents=True, exist_ok=True)
        attn.write_text(
            "jnd_1_1: {in_channels: 1, out_channels: 1}\n"
            "jnd_3_3: {in_channels: 3, out_channels: 3}\n"
            "jnd_1_3: {in_channels: 1, out_channels: 3}\n"
            "jnd_3_1: {in_channels: 3, out_channels: 1}\n"
        )


def decode(bits: np.ndarray) -> str | None:
    payload = bytes(np.packbits(np.asarray(bits, np.uint8).reshape(-1)[:256]))
    print(f"message (256b): {payload.hex()}")
    name, tag = payload[:16], payload[16:]
    if hashlib.sha256(name).digest()[:16] != tag:
        return None
    print(f"hash: {hashlib.sha256(name).hexdigest()}")
    return name.rstrip(b"\x00").decode("ascii", "replace") or None


def main() -> None:
    if len(sys.argv) != 2:
        sys.exit("usage: extract.py <video.mp4>")
    stage_videoseal_configs()

    model = videoseal.load("videoseal").eval()
    video, _, _ = torchvision.io.read_video(sys.argv[1], pts_unit="sec")
    frames = video.permute(0, 3, 1, 2).float() / 255.0

    with torch.no_grad():
        msg = model.extract_message(frames, aggregation="avg")
    print(decode(msg.squeeze().cpu().numpy()))


if __name__ == "__main__":
    main()
