#!/usr/bin/env python3
"""Local-only verifier for Stormtape synthetic-voice AudioSeal marks.

The verifier does not upload media and does not download models. It extracts a
bounded audio window with a caller-supplied FFmpeg binary, verifies the pinned
AudioSeal detector and policy, and prints a machine-readable JSON result.

Dependencies:
  Python 3.10+, torch, numpy, soundfile, omegaconf, audioseal==0.2.0, FFmpeg.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import os
import subprocess
import sys
import tempfile
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

os.environ.setdefault("HF_HUB_OFFLINE", "1")
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
os.environ.setdefault("NO_TORCH_COMPILE", "1")

EXPECTED_POLICY_ID = "stormtape-audioseal-v1"
EXPECTED_POLICY_SHA256 = "7A1770FAC86BEAE4E696FC2508461AB85CE7C9F2FE017760432EA888861463DD"
EXPECTED_DETECTOR_SHA256 = "8A78E8A83584113523E161FC599FCAB10FD0E94C04D2EB9D2FA1E9EC91AB69D9"
EXPECTED_PAYLOAD_HEX = "5354"
EXPECTED_THRESHOLD = 0.99
MAX_MEDIA_BYTES = 64 * 1024 * 1024 * 1024
MAX_POLICY_BYTES = 64 * 1024
MAX_WINDOW_SECONDS = 60
VERIFIER_ID = "stormtape-public-audioseal-verifier-v1"


def sha256_bytes(value: bytes) -> str:
    return hashlib.sha256(value).hexdigest().upper()


def resolve_regular_file(value: str | Path, *, maximum: int) -> Path:
    path = Path(value).expanduser().resolve(strict=True)
    stat = path.lstat()
    if not path.is_file() or path.is_symlink():
        raise RuntimeError(f"regular file required: {path}")
    if stat.st_size < 1 or stat.st_size > maximum:
        raise RuntimeError(f"unsafe file size: {path} ({stat.st_size})")
    return path


def hash_regular_file(path: Path, *, maximum: int) -> tuple[str, int]:
    checked = resolve_regular_file(path, maximum=maximum)
    before = checked.stat()
    digest = hashlib.sha256()
    total = 0
    with checked.open("rb") as handle:
        while chunk := handle.read(4 * 1024 * 1024):
            total += len(chunk)
            if total > maximum:
                raise RuntimeError(f"file grew beyond limit while hashing: {checked}")
            digest.update(chunk)
    after = checked.stat()
    identity_before = (
        before.st_dev,
        before.st_ino,
        before.st_size,
        before.st_mtime_ns,
        before.st_ctime_ns,
    )
    identity_after = (
        after.st_dev,
        after.st_ino,
        after.st_size,
        after.st_mtime_ns,
        after.st_ctime_ns,
    )
    if total != before.st_size or identity_before != identity_after:
        raise RuntimeError(f"file changed while hashing: {checked}")
    return digest.hexdigest().upper(), total


def read_policy(path: Path) -> tuple[dict[str, Any], str]:
    raw = resolve_regular_file(path, maximum=MAX_POLICY_BYTES).read_bytes()
    policy_sha256 = sha256_bytes(raw)
    if policy_sha256 != EXPECTED_POLICY_SHA256:
        raise RuntimeError(f"policy SHA-256 mismatch: {policy_sha256}")
    policy = json.loads(raw.decode("utf-8"))
    expected = {
        "algorithm": "AudioSeal",
        "implementationVersion": "0.2.0",
        "policyId": EXPECTED_POLICY_ID,
        "payloadHex": EXPECTED_PAYLOAD_HEX,
        "detectorThreshold": EXPECTED_THRESHOLD,
        "failMode": "closed",
    }
    for key, value in expected.items():
        if policy.get(key) != value:
            raise RuntimeError(f"policy field mismatch: {key}")
    if policy.get("detectorModel", {}).get("sha256") != EXPECTED_DETECTOR_SHA256:
        raise RuntimeError("policy detector SHA-256 mismatch")
    return policy, policy_sha256


def parse_arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Verify a Stormtape AudioSeal mark without uploading the media.",
    )
    parser.add_argument("--input", required=True, type=Path, help="media file to inspect")
    parser.add_argument("--detector", required=True, type=Path, help="pinned detector_base.pth")
    parser.add_argument("--policy", required=True, type=Path, help="Stormtape public policy JSON")
    parser.add_argument("--ffmpeg", default="ffmpeg", help="FFmpeg executable")
    parser.add_argument(
        "--vendor",
        type=Path,
        help="optional directory containing the audioseal Python package",
    )
    parser.add_argument("--start-seconds", type=float, default=0.0)
    parser.add_argument("--duration-seconds", type=float, default=30.0)
    parser.add_argument(
        "--evidence",
        type=Path,
        help="optional new JSON file; the verifier refuses to overwrite it",
    )
    return parser.parse_args()


def import_runtime(vendor: Path | None):
    if vendor is not None:
        resolved = vendor.expanduser().resolve(strict=True)
        if not resolved.is_dir() or resolved.is_symlink():
            raise RuntimeError(f"regular vendor directory required: {resolved}")
        sys.path.insert(0, str(resolved))
    try:
        import numpy as np  # type: ignore
        import soundfile as sf  # type: ignore
        import torch  # type: ignore
        from audioseal import AudioSeal  # type: ignore
    except ImportError as error:
        raise RuntimeError(
            "missing local dependency; install audioseal==0.2.0, torch, numpy, "
            "soundfile and omegaconf before running the verifier"
        ) from error
    return np, sf, torch, AudioSeal


def extract_window(
    *,
    ffmpeg: str,
    source: Path,
    output: Path,
    start_seconds: float,
    duration_seconds: float,
) -> None:
    command = [
        ffmpeg,
        "-nostdin",
        "-hide_banner",
        "-loglevel",
        "error",
        "-ss",
        str(start_seconds),
        "-i",
        str(source),
        "-t",
        str(duration_seconds),
        "-map",
        "0:a:0",
        "-ac",
        "1",
        "-ar",
        "48000",
        "-c:a",
        "pcm_s16le",
        "-f",
        "wav",
        "-y",
        str(output),
    ]
    result = subprocess.run(
        command,
        check=False,
        capture_output=True,
        text=True,
        timeout=120,
    )
    if result.returncode != 0:
        detail = (result.stderr or result.stdout or str(result.returncode)).strip()
        raise RuntimeError(f"FFmpeg extraction failed: {detail[:500]}")


def stable_json(value: dict[str, Any]) -> bytes:
    return (
        json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
        + "\n"
    ).encode("utf-8")


def write_exclusive(path: Path, value: bytes) -> None:
    resolved = path.expanduser().resolve()
    resolved.parent.mkdir(parents=True, exist_ok=True)
    with resolved.open("xb") as handle:
        handle.write(value)
        handle.flush()
        os.fsync(handle.fileno())


def main() -> int:
    args = parse_arguments()
    if not math.isfinite(args.start_seconds) or args.start_seconds < 0:
        raise RuntimeError("--start-seconds must be a finite value >= 0")
    if (
        not math.isfinite(args.duration_seconds)
        or args.duration_seconds < 1
        or args.duration_seconds > MAX_WINDOW_SECONDS
    ):
        raise RuntimeError(
            f"--duration-seconds must be between 1 and {MAX_WINDOW_SECONDS}"
        )

    media = resolve_regular_file(args.input, maximum=MAX_MEDIA_BYTES)
    detector = resolve_regular_file(args.detector, maximum=2 * 1024 * 1024 * 1024)
    policy_path = resolve_regular_file(args.policy, maximum=MAX_POLICY_BYTES)
    policy, policy_sha256 = read_policy(policy_path)
    detector_sha256, _ = hash_regular_file(
        detector,
        maximum=2 * 1024 * 1024 * 1024,
    )
    if detector_sha256 != EXPECTED_DETECTOR_SHA256:
        raise RuntimeError(f"detector SHA-256 mismatch: {detector_sha256}")
    input_sha256, input_bytes = hash_regular_file(media, maximum=MAX_MEDIA_BYTES)
    verifier_path = resolve_regular_file(__file__, maximum=2 * 1024 * 1024)
    verifier_sha256, _ = hash_regular_file(
        verifier_path,
        maximum=2 * 1024 * 1024,
    )

    np, sf, torch, AudioSeal = import_runtime(args.vendor)
    with tempfile.TemporaryDirectory(prefix="stormtape-audioseal-verify-") as directory:
        extracted = Path(directory) / "audio-window.wav"
        extract_window(
            ffmpeg=args.ffmpeg,
            source=media,
            output=extracted,
            start_seconds=args.start_seconds,
            duration_seconds=args.duration_seconds,
        )
        extracted_sha256, extracted_bytes = hash_regular_file(
            extracted,
            maximum=512 * 1024 * 1024,
        )
        audio, sample_rate = sf.read(
            extracted,
            dtype="float32",
            always_2d=True,
        )
        if sample_rate != 48000 or audio.shape[1] != 1 or audio.shape[0] < 12000:
            raise RuntimeError(
                f"unexpected extracted audio: rate={sample_rate} shape={audio.shape}"
            )
        if not np.isfinite(audio).all():
            raise RuntimeError("extracted audio contains non-finite samples")
        tensor = torch.from_numpy(audio.T.copy()).unsqueeze(0)
        model = AudioSeal.load_detector(str(detector), nbits=16, device="cpu")
        model.eval()
        with torch.inference_mode():
            result, detected_message = model.detect_watermark(tensor)
        score = float(result.item() if hasattr(result, "item") else result)
        detected_bits = "".join(
            "1" if float(bit) >= 0.5 else "0"
            for bit in detected_message.reshape(-1).tolist()
        )

    expected_bits = f"{int(EXPECTED_PAYLOAD_HEX, 16):016b}"
    payload_matches = detected_bits == expected_bits
    verified = (
        math.isfinite(score)
        and score >= EXPECTED_THRESHOLD
        and payload_matches
    )
    evidence = {
        "algorithm": policy["algorithm"],
        "checkedAt": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
        "detectedBits": detected_bits,
        "detectorModelSha256": detector_sha256,
        "detectorScore": round(score, 9) if math.isfinite(score) else None,
        "detectorThreshold": EXPECTED_THRESHOLD,
        "durationSeconds": args.duration_seconds,
        "extractedAudioBytes": extracted_bytes,
        "extractedAudioSha256": extracted_sha256,
        "inputBytes": input_bytes,
        "inputPath": str(media),
        "inputSha256": input_sha256,
        "messageBits": expected_bits,
        "payloadHex": EXPECTED_PAYLOAD_HEX,
        "payloadMatches": payload_matches,
        "policyId": EXPECTED_POLICY_ID,
        "policySha256": policy_sha256,
        "schemaVersion": 1,
        "startSeconds": args.start_seconds,
        "verified": verified,
        "verifierId": VERIFIER_ID,
        "verifierSha256": verifier_sha256,
    }
    encoded = stable_json(evidence)
    if args.evidence is not None:
        write_exclusive(args.evidence, encoded)
    sys.stdout.buffer.write(encoded)
    return 0 if verified else 2


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except Exception as error:  # pylint: disable=broad-exception-caught
        failure = {
            "error": str(error),
            "schemaVersion": 1,
            "verified": False,
            "verifierId": VERIFIER_ID,
        }
        sys.stderr.buffer.write(stable_json(failure))
        raise SystemExit(2)