"""Forge a video that the board reads as <username>.
python forge.py <username> <input.mp4> <output.mp4>

Embeds the 256-bit payload the board expects (16 bytes of NUL-padded ASCII username +
16 bytes of SHA-256(name)[:16], hashed over the padded 16-byte name) and writes at CRF 10
so the watermark survives re-encoding.
"""

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 encode(name: str) -> np.ndarray:
    """16-byte NUL-padded ASCII name + 16-byte SHA-256(name)[:16] -> 256 bits."""
    raw = name.encode("ascii").ljust(16, b"\x00")
    payload = raw + hashlib.sha256(raw).digest()[:16]
    return np.unpackbits(np.frombuffer(payload, np.uint8))


def main() -> None:
    if len(sys.argv) != 4:
        sys.exit("usage: forge.py <username> <input.mp4> <output.mp4>")
    name, src, dst = sys.argv[1:]
    stage_videoseal_configs()

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

    msg = torch.from_numpy(encode(name)).float().unsqueeze(0)
    out = model.embed(frames, msgs=msg, is_video=True)
    wm = (out["imgs_w"].clamp(0, 1) * 255).byte().permute(0, 2, 3, 1).cpu()
    # CRF 10 (not the ~23 default): heavy compression flips the bit-perfect name half.
    torchvision.io.write_video(dst, wm, fps=24, video_codec="libx264", options={"crf": "10"})
    print(f"wrote {dst}")


if __name__ == "__main__":
    main()
