#!/usr/bin/env python3
"""Was this video made by ClonifyLab? The public check of ClonifyLab's invisible watermarks.

Every video ClonifyLab makes carries two invisible watermarks with one fixed, public key:
  picture  Meta Video Seal 1.0 (the TorchScript model y_256b_img.jit), 256 bits;
  sound    Meta AudioSeal 0.2 (its 16-bit detector, detector_base.pth), 16 bits.
The key is the SHA-256 digest of the text "ClonifyLab AI-generated v1": all 256 bits for the picture and the first 16
for the sound, in this order: byte 0 first, and in each byte the most significant bit first.

The check gives one p-value: the chance that a video WITHOUT our watermark matches the key at least this well. It says
"yes" when p <= 1e-6 (one in a million).
  picture  every 15th frame (2 per second at 30 fps), scaled to 256x256. Video Seal's 256 outputs are averaged over
           these frames; a bit is 1 when its average is above 0. With k of the 256 bits matching the key:
           p = P(Binomial(256, 1/2) >= k).
  sound    the first audio stream as 16 kHz mono. AudioSeal gives a probability (the share of samples it finds marked)
           and 16 bits. With a probability of at least 0.5: p = P(Binomial(16, 1/2) >= matching bits), else p = 1.
           (16 bits alone can never reach one in a million: the sound backs up the picture.)
  both     Fisher's method: X = -2 (ln p_picture + ln p_sound), p = P(chi-square with 4 degrees of freedom >= X).
The models are open (MIT) and the key is public, so the watermark is a label, not a signature: anyone with the same
models could remove or copy it. The C2PA Content Credentials in our files are the signed proof.

Needs Python 3.10 or newer, FFmpeg 5.1 or newer (ffmpeg and ffprobe on PATH), and (CPU only is fine):
    pip install torch --index-url https://download.pytorch.org/whl/cpu
    pip install numpy omegaconf einops
    pip install --no-deps audioseal==0.2.0
and the two model files (MIT licence; this script checks their SHA-256 before it loads them):
    https://dl.fbaipublicfiles.com/videoseal/y_256b_img.jit
    https://huggingface.co/facebook/audioseal/resolve/main/detector_base.pth

Usage:
    python verify.py VIDEO [VIDEO ...] [--video-model y_256b_img.jit] [--audio-model detector_base.pth] [--json]
Exit code: 0 = every file is "yes", 1 = at least one "no", 2 = a file or model could not be read. With --json: 0
whenever the report was printed (each file's verdict, or its error, is in it), 2 when the models could not be loaded.
The model paths default to where the ClonifyLab render image keeps them; the render service runs this same script
(with --json) for its QA, and scripts/mark.py, which embeds the watermarks, uses its key and detection code.
"""
import argparse
import hashlib
import json
import math
import os
import subprocess
import sys
import tempfile
import time
import traceback
import warnings

# AudioSeal 0.2 would otherwise wrap its networks in torch.compile, which needs a C++ compiler. Must be set before
# audioseal is imported.
os.environ.setdefault("NO_TORCH_COMPILE", "1")
# AudioSeal's loader still uses torch.nn.utils.weight_norm, which PyTorch reports as deprecated on every load.
warnings.filterwarnings("ignore", message=r".*weight_norm.* is deprecated", category=FutureWarning)

import numpy as np  # noqa: E402
import torch  # noqa: E402

KEY_TEXT = "ClonifyLab AI-generated v1"
VIDEO_BITS = 256
AUDIO_BITS = 16
# "yes" when the combined p-value is at most 1e-6.
YES_LOG10_P = -6.0
# Picture: every 15th frame (2 per second at 30 fps), scaled to the model's 256x256 input.
SAMPLE_EVERY_N_FRAMES = 15
DETECT_SIZE = 256
DETECT_BATCH = 16
# Sound: AudioSeal's models work on 16 kHz mono.
AUDIO_RATE = 16000
AUDIO_FOUND_PROBABILITY = 0.5
FFMPEG_THREADS = 4

VIDEO_MODEL_SHA256 = "5c7a4581c36fc6090aafdcfb3999123bae5172a4847f22e2da4e7fd1a39d1e1b"
AUDIO_DETECTOR_SHA256 = "8a78e8a83584113523e161fc599fcab10fd0e94c04d2eb9d2fa1e9ec91ab69d9"
VIDEO_MODEL_URL = "https://dl.fbaipublicfiles.com/videoseal/y_256b_img.jit"
AUDIO_DETECTOR_URL = "https://huggingface.co/facebook/audioseal/resolve/main/detector_base.pth"
# Where the render image keeps them (render/Dockerfile).
DEFAULT_VIDEO_MODEL = "/opt/models/videoseal/y_256b_img.jit"
DEFAULT_AUDIO_MODEL = "/opt/models/audioseal/detector_base.pth"


class VerifyError(Exception):
    """A file or model that cannot be read; the message is one plain sentence."""


# ------------------------------------------------------------------------------------------------ key and statistics
def key_bits(n=VIDEO_BITS, text=KEY_TEXT):
    """The first n bits of SHA-256(text): byte 0 first, the most significant bit of each byte first."""
    digest = hashlib.sha256(text.encode("utf-8")).digest()
    bits = [(byte >> (7 - i)) & 1 for byte in digest for i in range(8)]
    if not 0 < n <= len(bits):
        raise ValueError(f"the key has {len(bits)} bits")
    return bits[:n]


def log10_binomial_tail(k, n):
    """log10 P(Binomial(n, 1/2) >= k), exact (integer sums)."""
    if k <= 0:
        return 0.0
    if k > n:
        return -math.inf
    return math.log10(sum(math.comb(n, i) for i in range(k, n + 1))) - n * math.log10(2)


def fisher_log10(log10_ps):
    """Fisher's method for independent p-values given as log10: log10 P(chi-square(2m) >= -2 sum ln p)."""
    m = len(log10_ps)
    if m == 0:
        return 0.0
    if m == 1:
        return log10_ps[0]
    half = -sum(log10_ps) * math.log(10)  # X / 2
    if half == 0:
        return 0.0
    # survival function of chi-square with 2m degrees of freedom: exp(-half) * sum_{i<m} half^i / i!
    terms = [i * math.log(half) - math.lgamma(i + 1) for i in range(m)]
    top = max(terms)
    log_sf = -half + top + math.log(sum(math.exp(t - top) for t in terms))
    return min(0.0, log_sf / math.log(10))


def format_p(log10_p):
    """A p-value for people, e.g. 4.1e-62 (works far below the smallest float)."""
    if log10_p == -math.inf:
        return "0"
    exponent = math.floor(log10_p)
    mantissa = 10 ** (log10_p - exponent)
    if round(mantissa, 1) >= 10:
        mantissa, exponent = mantissa / 10, exponent + 1
    return f"{mantissa:.1f}e{exponent:+03d}"


# ------------------------------------------------------------------------------------------------ models
def sha256_file(path):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def check_model_file(path, sha256, what, url):
    """Refuses a missing or different model file before anything loads it (loading runs pickled code)."""
    if not os.path.isfile(path):
        raise VerifyError(f"the {what} {path} was not found (download it from {url})")
    digest = sha256_file(path)
    if digest != sha256:
        raise VerifyError(f"the {what} {path} is not the expected file (sha256 {digest}, expected {sha256})")


def load_video_model(path=DEFAULT_VIDEO_MODEL):
    """Video Seal 1.0 (TorchScript), checked by SHA-256."""
    check_model_file(path, VIDEO_MODEL_SHA256, "Video Seal model", VIDEO_MODEL_URL)
    return torch.jit.load(path, map_location="cpu").eval()


def load_audio_detector(path=DEFAULT_AUDIO_MODEL):
    """The AudioSeal 0.2 16-bit detector from a local file, checked by SHA-256."""
    check_model_file(path, AUDIO_DETECTOR_SHA256, "AudioSeal detector", AUDIO_DETECTOR_URL)
    from audioseal import AudioSeal

    return AudioSeal.load_detector(path, nbits=AUDIO_BITS).eval()


def versions():
    import audioseal

    return {"python": sys.version.split()[0], "torch": torch.__version__, "audioseal": audioseal.__version__}


# ------------------------------------------------------------------------------------------------ detection
def detect_frames(model, frames):
    """Sum of Video Seal's 256 raw outputs over frames (a float tensor N x 3 x H x W in [0, 1]), as float64 numpy."""
    with torch.inference_mode():
        out = model.detect(frames, True)
        return out.reshape(len(frames), -1)[:, 1:].double().sum(0).numpy().copy()


def video_result(sums, frames):
    """The picture's evidence from the summed outputs (numpy, 256 values) of `frames` frames."""
    if frames == 0:
        return {"frames": 0, "bits": VIDEO_BITS, "matching_bits": 0, "bit_accuracy": None, "log10_p": 0.0}
    bits = [int(v > 0) for v in (np.asarray(sums) / frames).tolist()]
    k = sum(int(b == kb) for b, kb in zip(bits, key_bits(VIDEO_BITS)))
    return {
        "frames": frames,
        "bits": VIDEO_BITS,
        "matching_bits": k,
        "bit_accuracy": round(k / VIDEO_BITS, 4),
        "log10_p": round(log10_binomial_tail(k, VIDEO_BITS), 2),
    }


def detect_audio_16k(detector, x16):
    """AudioSeal on 16 kHz mono samples (float32 numpy array): probability, matching key bits, p-value."""
    z = torch.from_numpy(np.array(x16, dtype=np.float32)).view(1, 1, -1)
    with torch.inference_mode():
        prob, message = detector.detect_watermark(z, AUDIO_RATE)
    # Rounded down, and "found" decided on that same number: QA reads the probability, people read "found", and a
    # value just under 0.5 can never show (or pass QA) as 0.5.
    prob = round_down(float(prob[0]), 4)
    bits = [int(b) for b in message[0].tolist()]
    k = sum(int(b == kb) for b, kb in zip(bits, key_bits(AUDIO_BITS)))
    found = prob >= AUDIO_FOUND_PROBABILITY
    return {
        "seconds": round(len(x16) / AUDIO_RATE, 2),
        "probability": prob,
        "bits": AUDIO_BITS,
        "matching_bits": k,
        "found": found,
        "log10_p": round(log10_binomial_tail(k, AUDIO_BITS), 2) if found else 0.0,
    }


def exact_log10_p(evidence):
    """One stream's log10 p from its counts, unrounded (its "log10_p" field is rounded for people): the binomial tail
    of its matching key bits, or 0 (p = 1) for a sound in which AudioSeal found no mark."""
    if evidence.get("found") is False:
        return 0.0
    return log10_binomial_tail(evidence["matching_bits"], evidence["bits"])


def combine(video, audio):
    """One verdict from the picture's and the sound's evidence (either may be None). The verdict uses the exact
    p-values, so rounding can never move a file across the threshold."""
    parts = [exact_log10_p(r) for r in (video, audio) if r is not None]
    log10_p = fisher_log10(parts) if parts else 0.0
    return {
        "ai_generated": log10_p <= YES_LOG10_P,
        "log10_p": round(log10_p, 2),
        "p": format_p(log10_p),
        "threshold_log10_p": YES_LOG10_P,
    }


def round_down(x, digits):
    """x rounded down to `digits` decimals. The 1e-9 keeps a value that already has that many decimals whole
    (0.57 * 100 = 56.99999999999999)."""
    scale = 10 ** digits
    return math.floor(x * scale + 1e-9) / scale


def short_error(e):
    """An unexpected exception in a few words: its type and the first line of its message."""
    text = str(e).strip().splitlines()
    return f"{type(e).__name__}: {text[0][:200]}" if text else type(e).__name__


# ------------------------------------------------------------------------------------------------ ffmpeg
def ffmpeg_error(stderr):
    lines = [line.strip() for line in stderr.decode("utf-8", "replace").splitlines() if line.strip()]
    return lines[-1] if lines else "no message"


def probe_streams(path):
    """(has a picture, has sound) of a media file (cover art does not count as a picture)."""
    cmd = ["ffprobe", "-v", "error", "-protocol_whitelist", "file", "-show_entries",
           "stream=codec_type:stream_disposition=attached_pic", "-of", "json", path]
    r = subprocess.run(cmd, capture_output=True, check=False)
    if r.returncode != 0:
        raise VerifyError(f"{os.path.basename(path)} could not be read ({ffmpeg_error(r.stderr)})")
    streams = json.loads(r.stdout or b"{}").get("streams", [])
    video = any(s.get("codec_type") == "video" and not s.get("disposition", {}).get("attached_pic") for s in streams)
    audio = any(s.get("codec_type") == "audio" for s in streams)
    return video, audio


def check_video(model, path):
    """Decodes every 15th frame at 256x256 (RGB) and runs Video Seal's detector on them, a batch at a time."""
    vf = f"select=not(mod(n\\,{SAMPLE_EVERY_N_FRAMES})),scale={DETECT_SIZE}:{DETECT_SIZE}:flags=area,format=rgb24"
    cmd = ["ffmpeg", "-hide_banner", "-nostdin", "-v", "error", "-threads", str(FFMPEG_THREADS),
           "-protocol_whitelist", "file", "-i", path, "-map", "0:V:0", "-vf", vf, "-fps_mode", "passthrough",
           "-f", "rawvideo", "-pix_fmt", "rgb24", "pipe:1"]
    frame_bytes = DETECT_SIZE * DETECT_SIZE * 3
    started = time.perf_counter()
    # FFmpeg's messages go to a temporary file: a pipe nobody reads could fill up and stall it.
    with tempfile.TemporaryFile() as log:
        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=log)
        try:
            sums = np.zeros(VIDEO_BITS)
            frames = 0
            while True:
                buf = proc.stdout.read(frame_bytes * DETECT_BATCH)
                n = len(buf) // frame_bytes
                if n == 0:
                    break
                x = torch.frombuffer(bytearray(buf[: n * frame_bytes]), dtype=torch.uint8)
                x = x.view(n, DETECT_SIZE, DETECT_SIZE, 3).permute(0, 3, 1, 2).float().div_(255)
                sums += detect_frames(model, x)
                frames += n
        finally:
            proc.stdout.close()
            code = proc.wait()
        log.seek(0)
        stderr = log.read()
    if code != 0:
        if b"Unrecognized option" in stderr and b"fps_mode" in stderr:
            raise VerifyError("this FFmpeg is too old: the check needs FFmpeg 5.1 or newer")
        raise VerifyError(f"the picture of {os.path.basename(path)} could not be decoded ({ffmpeg_error(stderr)})")
    result = video_result(sums, frames)
    result["ms"] = round(1000 * (time.perf_counter() - started))
    return result


def read_audio_16k(path):
    """The first audio stream as 16 kHz mono float32 samples."""
    cmd = ["ffmpeg", "-hide_banner", "-nostdin", "-v", "error", "-protocol_whitelist", "file", "-i", path,
           "-map", "0:a:0", "-ac", "1", "-ar", str(AUDIO_RATE), "-f", "f32le", "pipe:1"]
    r = subprocess.run(cmd, capture_output=True, check=False)
    if r.returncode != 0:
        raise VerifyError(f"the sound of {os.path.basename(path)} could not be decoded ({ffmpeg_error(r.stderr)})")
    return np.frombuffer(r.stdout, dtype="<f4")


def check_audio(detector, path):
    started = time.perf_counter()
    x16 = read_audio_16k(path)
    if len(x16) < AUDIO_RATE // 10:
        result = {"seconds": round(len(x16) / AUDIO_RATE, 2), "probability": 0.0, "bits": AUDIO_BITS,
                  "matching_bits": 0, "found": False, "log10_p": 0.0}
    else:
        result = detect_audio_16k(detector, x16)
    result["ms"] = round(1000 * (time.perf_counter() - started))
    return result


def check_file(path, video_model, audio_detector):
    """The full check of one file: {file, ai_generated, log10_p, p, threshold_log10_p, video, audio}."""
    has_video, has_audio = probe_streams(path)
    if not has_video and not has_audio:
        raise VerifyError(f"{os.path.basename(path)} has no picture and no sound")
    video = check_video(video_model, path) if has_video else None
    audio = check_audio(audio_detector, path) if has_audio else None
    return {"file": path, **combine(video, audio), "video": video, "audio": audio}


# ------------------------------------------------------------------------------------------------ command line
def describe(result):
    """The human-readable lines for one file."""
    lines = [f"AI-generated by ClonifyLab: {'yes' if result['ai_generated'] else 'no'} (p = {result['p']})"]
    v, a = result["video"], result["audio"]
    if v is None:
        lines.append("  picture: none (no video stream)")
    elif v["frames"] == 0:
        lines.append("  picture: no frames could be checked")
    else:
        lines.append(f"  picture: {v['matching_bits']} of {v['bits']} key bits match "
                     f"(bit accuracy {v['bit_accuracy']:.3f}, {v['frames']} frames checked)")
    if a is None:
        lines.append("  sound: none (no audio stream)")
    elif a["found"]:
        lines.append(f"  sound: watermark found (probability {round_down(a['probability'], 2):.2f}), "
                     f"{a['matching_bits']} of {a['bits']} key bits match")
    else:
        lines.append(f"  sound: no watermark found (probability {round_down(a['probability'], 2):.2f})")
    return lines


def default_threads():
    try:
        return max(1, int(os.environ.get("OMP_NUM_THREADS", "")))
    except ValueError:
        return max(1, min(4, os.cpu_count() or 1))


def main(argv=None):
    ap = argparse.ArgumentParser(description="Was this video made by ClonifyLab? Checks its invisible watermarks.")
    ap.add_argument("files", nargs="+", metavar="VIDEO")
    ap.add_argument("--video-model", default=DEFAULT_VIDEO_MODEL, help="y_256b_img.jit (Video Seal 1.0)")
    ap.add_argument("--audio-model", default=DEFAULT_AUDIO_MODEL, help="detector_base.pth (AudioSeal 0.2, 16 bits)")
    ap.add_argument("--json", action="store_true", help="print one JSON object instead of text")
    ap.add_argument("--threads", type=int, default=default_threads(), help="CPU threads for the models")
    args = ap.parse_args(argv)
    torch.set_num_threads(max(1, args.threads))
    try:
        torch.set_num_interop_threads(1)
    except RuntimeError:
        pass
    out = sys.stdout
    sys.stdout = sys.stderr  # nothing but the result goes to stdout
    started = time.perf_counter()
    try:
        video_model = load_video_model(args.video_model)
        audio_detector = load_audio_detector(args.audio_model)
    except VerifyError as e:
        print(f"verify.py: error: {e}", file=sys.stderr)
        return 2
    except Exception as e:  # noqa: BLE001  (e.g. a missing Python package)
        traceback.print_exc()
        print(f"verify.py: error: the models could not be loaded ({short_error(e)})", file=sys.stderr)
        return 2
    load_ms = round(1000 * (time.perf_counter() - started))
    results, status = [], 0
    for path in args.files:
        try:
            result = check_file(path, video_model, audio_detector)
            if not result["ai_generated"]:
                status = max(status, 1)
        except VerifyError as e:
            result = {"file": path, "error": str(e)}
            status = 2
        except Exception as e:  # noqa: BLE001  (one file the models fail on; the others are still checked)
            traceback.print_exc()
            result = {"file": path, "error": f"{os.path.basename(path)} could not be checked ({short_error(e)})"}
            status = 2
        results.append(result)
        if not args.json:
            if len(args.files) > 1:
                out.write(f"{path}\n")
            out.write("\n".join(describe(result) if "error" not in result else [f"error: {result['error']}"]) + "\n")
            out.flush()
    if args.json:
        report = {"key": KEY_TEXT, "threshold_log10_p": YES_LOG10_P, "load_ms": load_ms, "versions": versions(),
                  "files": results}
        out.write(json.dumps(report) + "\n")
        out.flush()
        return 0
    return status


if __name__ == "__main__":
    sys.exit(main())