#!/usr/bin/env python3
"""
Prüft, ob eine Audiodatei das unhörbare Wasserzeichen von lalabai® trägt.
Checks whether an audio file carries the inaudible lalabai® watermark.

lalabai® markiert jedes Hörspiel mit AudioSeal (Meta, MIT-Lizenz) und der
16-Bit-Kennung 0x4C42 ("LB"). Dieses Skript nutzt Metas offenen
AudioSeal-Detektor: Ob ein Wasserzeichen vorhanden ist, entscheidet der
Detektor selbst; die Kennung ordnet es lalabai® zu. Weil ein Schnitt oder
eine Neukodierung das Audio um Bruchteile einer Millisekunde verschiebt,
sucht das Skript die passende Ausrichtung (in Schritten von 4 Samples) und
bestätigt einen Fund an bis zu 15 Stellen. Bei Musik ohne Sprache ist die
Kennung manchmal nicht lesbar, das Wasserzeichen wird trotzdem erkannt.

  pip install audioseal torch numpy      (dazu ffmpeg im Suchpfad)
  python3 wasserzeichen-pruefen.py DATEI [--json]

Die Datei verlässt deinen Rechner nicht. Mehr: https://lalabai.com/de-de/ki-kennzeichnung
"""

from __future__ import annotations

import json
import subprocess
import sys

import numpy as np
import torch
from audioseal import AudioSeal

MESSAGE = 0x4C42  # lalabai® ("LB")
RATE = 16000      # AudioSeal works at 16 kHz
PROBE = 64000     # 4 s detector window
HOP = 320         # AudioSeal frame hop (20 ms)
STEP = 4          # alignment search step in samples
CONFIRM = 15      # windows that confirm a find
PRESENCE = 0.5    # share of watermarked samples that counts as detected (AudioSeal's threshold)


def decode(path: str) -> np.ndarray:
    r = subprocess.run(["ffmpeg", "-v", "error", "-i", path, "-map", "0:a:0", "-ac", "1", "-ar", str(RATE), "-f", "f32le", "-"],
                       capture_output=True)
    if r.returncode != 0:
        sys.exit(f"ffmpeg konnte die Datei nicht lesen / could not read the file: {r.stderr.decode()[-300:]}")
    return np.frombuffer(r.stdout, dtype=np.float32).copy()


def level_db(x: np.ndarray) -> float:
    return float(20 * np.log10(np.sqrt(np.mean(np.square(x, dtype=np.float64))) + 1e-12))


class Checker:
    def __init__(self):
        self.det = AudioSeal.load_detector("audioseal_detector_16bits").eval()

    def window(self, seg: np.ndarray):
        """One detector pass: share of watermarked samples, summed message logits, hits."""
        buf = np.zeros(PROBE, dtype=np.float32)
        buf[: len(seg)] = seg[:PROBE]
        with torch.no_grad():
            logits = self.det.detector(torch.from_numpy(buf)[None, None, :])[0].numpy()[:, : len(seg)]
        e = np.exp(logits[:2] - logits[:2].max(axis=0))
        hit = (e[1] / e.sum(axis=0)) > 0.5
        msg = logits[2:, hit].sum(axis=1) if hit.any() else np.zeros(16)
        return float(hit.mean()), msg, int(hit.sum())

    @staticmethod
    def value(msg: np.ndarray) -> int:
        return int("".join("1" if v > 0 else "0" for v in msg), 2)

    def confirm(self, x: np.ndarray, phase: int) -> dict:
        """Detector windows over the whole file at one alignment: presence and message."""
        span = max(0, (len(x) - PROBE - phase) // HOP)
        picks = sorted({phase + int(k) * HOP for k in np.linspace(0, span, min(CONFIRM, span + 1)).round()})
        total = hits = windows = marked = 0
        msg_sum = np.zeros(16)
        for p in picks:
            seg = x[p: p + PROBE]
            if len(seg) < PROBE // 2 or level_db(seg) < -50:
                continue
            share, msg, h = self.window(seg)
            windows += 1
            marked += int(share >= PRESENCE)
            total += len(seg)
            hits += h
            msg_sum += msg
        share = hits / max(1, total)
        found = self.value(msg_sum)
        return {"detected": share >= PRESENCE and 2 * marked > windows, "share": round(share, 3),
                "windows": f"{marked}/{windows}", "message": f"0x{found:04X}", "lalabai": found == MESSAGE,
                "seconds_checked": round(total / RATE, 1)}

    def check(self, x: np.ndarray) -> dict:
        """Presence tolerates small shifts, the message needs the exact alignment:
        scan every alignment of a 4 s probe, confirm those that best decode
        lalabai's identifier over the whole file."""
        n = len(x)
        if n < PROBE + HOP:
            return {"detected": False, "reason": "zu kurz / too short (min. 5 s)"}
        tried = 0
        best: dict = {}
        for frac in (0.5, 0.3, 0.7, 0.15, 0.85):
            start = min(max(0, int(n * frac) - PROBE // 2), n - PROBE - HOP)
            if level_db(x[start: start + PROBE]) < -50:
                continue
            tried += 1
            scan = []
            for off in range(0, HOP, STEP):
                share, msg, _ = self.window(x[start + off: start + off + PROBE])
                if share >= PRESENCE:
                    scan.append((16 - bin(self.value(msg) ^ MESSAGE).count("1"), share, off))
                    if scan[-1][0] == 16 and share >= 0.9:
                        break
            for _, _, off in sorted(scan, reverse=True)[:3]:
                result = self.confirm(x, (start + off) % HOP)
                if result["detected"] and result["lalabai"]:
                    return result
                if (result["detected"], result["share"]) > (best.get("detected", False), best.get("share", -1.0)):
                    best = result
            if best.get("detected") and tried >= 2:
                break
        if best.get("detected"):
            return best
        reason = "kein Wasserzeichen gefunden / no watermark found" if tried else "nur Stille / only silence"
        return {**best, "detected": False, "reason": reason}


def main() -> int:
    args = [a for a in sys.argv[1:] if not a.startswith("--")]
    if len(args) != 1:
        print(__doc__)
        return 2
    result = Checker().check(decode(args[0]))
    if "--json" in sys.argv:
        print(json.dumps(result, ensure_ascii=False))
    elif result.get("detected") and result.get("lalabai"):
        print(f"Wasserzeichen von lalabai® gefunden (Kennung {result['message']}, Anteil {result['share'] * 100:.0f} %, "
              f"{result['seconds_checked']:.0f} s geprüft).\nlalabai® watermark found (identifier 0x4C42).")
    elif result.get("detected"):
        print(f"Wasserzeichen gefunden: Die Datei ist als KI-erzeugt gekennzeichnet (Anteil {result['share'] * 100:.0f} %, "
              f"{result['seconds_checked']:.0f} s geprüft). Die Kennung von lalabai® ließ sich nicht sicher lesen; "
              f"das kommt bei Musik ohne Sprache vor.\n"
              f"Watermark found: the file is marked as AI-generated. The lalabai® identifier could not be read "
              f"reliably, which happens with music without speech.")
    else:
        print(f"Kein Wasserzeichen gefunden ({result.get('reason')}).\nNo watermark found.")
    return 0 if result.get("detected") else 1


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