#!/usr/bin/env python3
"""Fail if a committed audio vector is a coding of a recording in a reference set.

Every audio vector in this tree must come from a seeded generator. This check finds
the ones that do not, by content: it decodes or reads each committed audio vector to
PCM and computes its normalised cross-correlation, at every lag, against every
recording in a reference set the host has. A vector that correlates above the
threshold (0.5 by default) with any reference fails the check. A generated signal
correlates with an unrelated recording at a few hundredths; a lossy coding of a
recording correlates with it near 1.

    python3 tools/wai_audio_provenance_check.py [--wai DIR] [--threshold T]
                                               [--ref DIR]... [--pcm PATH:RATE[:CH]]...

Run from `wai/` (or pass `--wai`). Python standard library only.

**The reference set** is every audio file under each `--ref` directory, under each
directory in `WAI_AUDIO_REFERENCE_DIRS` (separated by the platform's path separator),
and under the host's system sound directory (`/System/Library/Sounds`) when it exists.
Each reference is converted to mono PCM at the analysis rate by the host's audio
converter (`afconvert`, else `ffmpeg`); raw `.s16` references are read as 16 kHz mono.
When the set is empty, or no converter can read any of it, the check prints why and
exits 0: it is skipped, not passed.

**The committed audio vectors** are found by rule, not listed, so a new one is
covered when it lands:

- each base of a staged-delivery case (`staged-conformance/cases/*/case.json`) that
  decodes through a `wai.audio.*` capability: its stage 0 (`stage0.bin`, the decoded
  PCM, i16 interleaved), at the sample rate its payload states (a `WIA1` header's
  `sr`, a FLAC stream's STREAMINFO);
- each staged-measurement title over such a base
  (`staged-measure-conformance/fixtures/*/title.json`): its source (`source.bin`), at
  the base's rate scaled by the source's length over the base's;
- each spatial-audio source (`audio-conformance/vectors/*.sources.bin`), at the rate
  its scene (`.wscn`) states;
- every file with an audio extension (`.wav`, `.aif`, `.aiff`, `.flac`, `.s16`, …)
  anywhere under `wai/`;
- and every `--pcm PATH:RATE[:CH]` given (raw s16le, CH channels interleaved).

A signal shorter than a quarter of a second at the analysis rate, or silent, is listed
and not judged: over so few samples a generated tone can match a recorded one by
chance (a 50 ms pair of triangle waves reaches 0.5 against a low recorded tone).

A learned-codec vector (`wai-rs/tests/vectors/*/audio_model.bin` beside a
`sha256.json`) is covered through its decode: the check requires that the decoded
PCM's SHA-256 its sidecar states is the stage 0 of a staged-delivery case above, so
the PCM it scans is that vector's output. A codec vector with no such case fails.

Exit status: 0 when every vector is under the threshold or the check is skipped;
1 when a vector correlates above it, or a codec vector's decode is not covered.
"""
import argparse
import cmath
import glob
import hashlib
import json
import math
import os
import re
import shutil
import struct
import subprocess
import sys
import tempfile
import wave

ANALYSIS_RATE = 8000           # Hz: every signal is compared here, mono
SYSTEM_SOUNDS = '/System/Library/Sounds'
AUDIO_EXT = {'.wav', '.wave', '.aif', '.aiff', '.aifc', '.flac', '.mp3', '.ogg', '.opus',
             '.m4a', '.caf', '.s16', '.pcm'}
SKIP_DIRS = {'.git', 'target', 'node_modules', 'jpeg-ai-ref', '__pycache__'}
MIN_SAMPLES = ANALYSIS_RATE // 4   # a quarter second: a shorter signal is listed, not judged


# ---- PCM helpers --------------------------------------------------------------------
def i16le(b):
    n = len(b) // 2
    return list(struct.unpack(f'<{n}h', b[:2 * n]))


def channels(pcm, ch):
    return [pcm[c::ch] for c in range(ch)]


def lowpass_taps(cut, n=63):
    """A Hann-windowed sinc low-pass, cut-off `cut` in cycles per input sample."""
    m = n // 2
    taps = []
    for k in range(n):
        x = k - m
        s = 2 * cut if x == 0 else math.sin(2 * math.pi * cut * x) / (math.pi * x)
        taps.append(s * (0.5 - 0.5 * math.cos(2 * math.pi * k / (n - 1))))
    g = sum(taps)
    return [t / g for t in taps]


def resample(x, rate):
    """Resample `x` from `rate` to the analysis rate: low-pass, then linear interpolation."""
    if rate == ANALYSIS_RATE or not x:
        return [float(v) for v in x]
    if rate > ANALYSIS_RATE:
        taps = lowpass_taps(0.45 * ANALYSIS_RATE / rate)
        m = len(taps) // 2
        n = len(x)
        y = []
        for i in range(n):
            acc = 0.0
            lo, hi = max(0, i - m), min(n, i + m + 1)
            for j in range(lo, hi):
                acc += taps[j - i + m] * x[j]
            y.append(acc)
    else:
        y = [float(v) for v in x]
    step = rate / ANALYSIS_RATE
    out, t = [], 0.0
    while t <= len(y) - 1:
        i = int(t)
        f = t - i
        out.append(y[i] if f == 0 or i + 1 >= len(y) else y[i] * (1 - f) + y[i + 1] * f)
        t += step
    return out


# ---- correlation ----------------------------------------------------------------------
_TW = {}


def fft(a, inverse=False):
    """In-place iterative radix-2 FFT of a list of complex numbers (len a power of 2)."""
    n = len(a)
    j = 0
    for i in range(1, n):
        bit = n >> 1
        while j & bit:
            j ^= bit
            bit >>= 1
        j |= bit
        if i < j:
            a[i], a[j] = a[j], a[i]
    key = (n, inverse)
    if key not in _TW:
        sgn = 1 if inverse else -1
        _TW[key] = [cmath.exp(sgn * 2j * math.pi * k / n) for k in range(n // 2)]
    tw = _TW[key]
    size = 2
    while size <= n:
        half = size // 2
        stride = n // size
        for start in range(0, n, size):
            for k in range(half):
                u = a[start + k]
                v = a[start + k + half] * tw[k * stride]
                a[start + k] = u + v
                a[start + k + half] = u - v
        size *= 2
    if inverse:
        for i in range(n):
            a[i] /= n
    return a


def spectrum(x, n):
    return fft([complex(v) for v in x] + [0j] * (n - len(x)))


def prefix_sq(x):
    p = [0.0]
    for v in x:
        p.append(p[-1] + v * v)
    return p


def best_ncc(xa, fa, pa, xb, fb, pb, n):
    """max over lags of |Σ a[i]·b[i+lag]| / sqrt(Σ a² · Σ b²), both sums over the
    overlap, among lags whose overlap is at least half the shorter signal."""
    la, lb = len(xa), len(xb)
    prod = [fa[k].conjugate() * fb[k] for k in range(n)]
    c = fft(prod, inverse=True)                    # c[lag mod n] = Σ a[i] b[i+lag]
    need = max(1, min(la, lb) // 2)
    best, at = 0.0, 0
    for lag in range(-(la - 1), lb):
        i0, i1 = max(0, -lag), min(la, lb - lag)   # a's overlapping range
        if i1 - i0 < need:
            continue
        ea = pa[i1] - pa[i0]
        eb = pb[i1 + lag] - pb[i0 + lag]
        if ea <= 0 or eb <= 0:
            continue
        r = abs(c[lag % n].real) / math.sqrt(ea * eb)
        if r > best:
            best, at = r, lag
    return best, at


# ---- the reference set --------------------------------------------------------------
def converter():
    if shutil.which('afconvert'):
        return 'afconvert'
    if shutil.which('ffmpeg'):
        return 'ffmpeg'
    return None


def read_wav(path):
    """(samples per channel list, rate) for PCM 8/16/24/32-bit or float32 WAV; raises on
    anything else."""
    with open(path, 'rb') as f:
        b = f.read()
    if b[:4] != b'RIFF' or b[8:12] != b'WAVE':
        raise ValueError('not RIFF/WAVE')
    pos, fmt, data = 12, None, None
    while pos + 8 <= len(b):
        cid, size = b[pos:pos + 4], struct.unpack('<I', b[pos + 4:pos + 8])[0]
        body = b[pos + 8:pos + 8 + size]
        if cid == b'fmt ' and len(body) >= 16:
            fmt = struct.unpack('<HHIIHH', body[:16])
        elif cid == b'data':
            data = body
        pos += 8 + size + (size & 1)
    if fmt is None or data is None:
        raise ValueError('no fmt or data chunk')
    tag, ch, rate, _, _, bits = fmt
    if ch < 1 or ch > 64 or rate < 1:
        raise ValueError(f'{ch} channels at {rate} Hz')
    if tag == 1 and bits == 16:
        x = i16le(data)
    elif tag == 1 and bits == 8:
        x = [(v - 128) * 256 for v in data]
    elif tag == 1 and bits == 24:
        x = [int.from_bytes(data[i:i + 3], 'little', signed=True) >> 8 for i in range(0, len(data) - 2, 3)]
    elif tag == 1 and bits == 32:
        x = [v >> 16 for v in struct.unpack(f'<{len(data) // 4}i', data[:len(data) // 4 * 4])]
    elif tag == 3 and bits == 32:
        x = [v * 32767 for v in struct.unpack(f'<{len(data) // 4}f', data[:len(data) // 4 * 4])]
    else:
        raise ValueError(f'format tag {tag}, {bits} bits')
    x = x[:len(x) // ch * ch]
    return channels(x, ch), rate


def convert(conv, src, tmp):
    """Mono s16 WAV at the analysis rate, via the host's converter."""
    out = os.path.join(tmp, hashlib.sha256(src.encode()).hexdigest()[:16] + '.wav')
    if conv == 'afconvert':
        cmd = ['afconvert', '-f', 'WAVE', '-d', f'LEI16@{ANALYSIS_RATE}', '-c', '1', src, out]
    else:
        cmd = ['ffmpeg', '-v', 'error', '-y', '-i', src, '-ac', '1', '-ar', str(ANALYSIS_RATE),
               '-c:a', 'pcm_s16le', out]
    subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
    chans, rate = read_wav(out)
    return resample(chans[0], rate)


def references(dirs, conv, tmp):
    refs, notes = [], []
    for d in dirs:
        for root, _, files in os.walk(d):
            for f in sorted(files):
                p = os.path.join(root, f)
                ext = os.path.splitext(f)[1].lower()
                if ext not in AUDIO_EXT:
                    continue
                try:
                    if ext == '.s16':
                        with open(p, 'rb') as fh:
                            x = resample(i16le(fh.read()), 16000)
                    elif conv:
                        x = convert(conv, p, tmp)
                    else:
                        chans, rate = read_wav(p)
                        x = resample(chans[0], rate)
                except Exception as e:  # noqa: BLE001 — an unreadable reference is reported
                    notes.append(f'reference unreadable, skipped: {p} ({e})')
                    continue
                refs.append((p, x))
    return refs, notes


# ---- the committed audio vectors ----------------------------------------------------
def envelope_payload(b):
    """The payload of a WAI1 envelope: magic, u32 manifest length, manifest, then the
    payload (a u32 length, or to the end)."""
    if b[:4] != b'WAI1':
        return None
    mlen = struct.unpack('<I', b[4:8])[0]
    rest = b[8 + mlen:]
    if len(rest) >= 4:
        plen = struct.unpack('<I', rest[:4])[0]
        if plen <= len(rest) - 4:
            return rest[4:4 + plen]
    return rest


def payload_rate(p):
    """The sample rate a payload states: a WIA1 header's `sr`, or a FLAC STREAMINFO's."""
    if p is None:
        return None
    if p[:4] == b'WIA1' and len(p) >= 12:
        return struct.unpack('<I', p[8:12])[0]
    i = p.find(b'fLaC')
    if i >= 0 and len(p) >= i + 8 + 13:
        si = p[i + 8:i + 8 + 34]                   # STREAMINFO body
        return (si[10] << 12) | (si[11] << 4) | (si[12] >> 4)
    return None


def staged_bases(wai):
    """(case dir, base index) → (pcm, channels, rate, stage0 sha256) for audio bases."""
    out = {}
    for cj in sorted(glob.glob(os.path.join(wai, 'staged-conformance/cases/*/case.json'))):
        case = json.load(open(cj))
        d = os.path.dirname(cj)
        for bi, base in enumerate(case.get('bases', [])):
            if not str(base.get('decoded_via', '')).startswith('wai.audio.'):
                continue
            units = base.get('units') or [{}]
            shape = units[0].get('shape') or [0, 1, 1]
            ch = int(shape[2]) if len(shape) > 2 else 1
            with open(os.path.join(d, base['stage0']), 'rb') as f:
                s0 = f.read()
            with open(os.path.join(d, base['file']), 'rb') as f:
                rate = payload_rate(envelope_payload(f.read()))
            out[(os.path.relpath(d, wai), bi)] = (i16le(s0), ch, rate, hashlib.sha256(s0).hexdigest())
    return out


def vectors(wai, extra):
    """[(name, [channel samples], rate)], notes, failures."""
    vs, notes, fails = [], [], []
    bases = staged_bases(wai)
    for (case, bi), (pcm, ch, rate, _) in bases.items():
        name = f'{case}/stage0 (base {bi})'
        if not rate:
            fails.append(f'{name}: no sample rate in its payload')
            continue
        vs.append((name, channels(pcm, ch), rate))

    for tj in sorted(glob.glob(os.path.join(wai, 'staged-measure-conformance/fixtures/*/title.json'))):
        t = json.load(open(tj))
        for b in t.get('bases', []):
            key = (b.get('case'), b.get('base', 0))
            if key not in bases:
                continue
            pcm, ch, rate, _ = bases[key]
            src = os.path.join(os.path.dirname(tj), t['source']['file'])
            with open(src, 'rb') as f:
                x = i16le(f.read())
            if not rate:
                continue
            r = round(rate * len(x) / len(pcm))
            vs.append((os.path.relpath(src, wai), channels(x, ch), r))

    for sb in sorted(glob.glob(os.path.join(wai, 'audio-conformance/vectors/*.sources.bin'))):
        with open(sb, 'rb') as f:
            b = f.read()
        scene = sb[:-len('.sources.bin')] + '.wscn'
        m = re.search(rb'"sample_rate":(\d+)', open(scene, 'rb').read()) if os.path.exists(scene) else None
        rate = int(m.group(1)) if m else 48000
        n, pos = struct.unpack('<I', b[:4])[0], 4
        for k in range(n):
            h, ln = b[pos:pos + 32].hex()[:12], struct.unpack('<I', b[pos + 32:pos + 36])[0]
            pos += 36
            vs.append((f'{os.path.relpath(sb, wai)}#{k} ({h}…)', [i16le(b[pos:pos + 2 * ln])], rate))
            pos += 2 * ln

    for root, dirs, files in os.walk(wai):
        dirs[:] = [d for d in dirs if d not in SKIP_DIRS]
        for f in sorted(files):
            ext = os.path.splitext(f)[1].lower()
            if ext not in AUDIO_EXT:
                continue
            p = os.path.join(root, f)
            try:
                if ext in ('.s16', '.pcm'):
                    with open(p, 'rb') as fh:
                        vs.append((os.path.relpath(p, wai), [i16le(fh.read())], 16000))
                else:
                    chans, rate = read_wav(p)
                    vs.append((os.path.relpath(p, wai), chans, rate))
            except Exception as e:  # noqa: BLE001
                notes.append(f'not readable as PCM, listed only: {os.path.relpath(p, wai)} ({e})')

    for spec in extra:
        parts = spec.split(':')
        p, rate = parts[0], int(parts[1])
        ch = int(parts[2]) if len(parts) > 2 else 1
        with open(p, 'rb') as f:
            vs.append((p, channels(i16le(f.read()), ch), rate))

    # A learned codec vector is covered by its decode, which must be a case's stage 0.
    stage0 = {sha for (_, _, _, sha) in bases.values()}
    for mb in sorted(glob.glob(os.path.join(wai, 'wai-rs/tests/vectors/*/audio_model.bin'))):
        side = os.path.join(os.path.dirname(mb), 'sha256.json')
        if not os.path.exists(side):
            fails.append(f'{os.path.relpath(mb, wai)}: no sha256.json sidecar to tie its decode to a scanned PCM')
            continue
        sha = json.load(open(side)).get('audio', {}).get('sha256')
        if sha not in stage0:
            fails.append(f'{os.path.relpath(os.path.dirname(mb), wai)}: its decoded PCM (sha256 {sha}) is no '
                         'staged-delivery case\'s stage 0, so it is not scanned; add a case that decodes it')
        else:
            notes.append(f'{os.path.relpath(os.path.dirname(mb), wai)}: decode covered as a staged case\'s stage 0')
    return vs, notes, fails


def main():
    ap = argparse.ArgumentParser(description=__doc__.split('\n')[0])
    ap.add_argument('--wai', default='.', help='the wai/ directory (default: .)')
    ap.add_argument('--threshold', type=float, default=0.5)
    ap.add_argument('--ref', action='append', default=[], help='a reference directory (repeatable)')
    ap.add_argument('--pcm', action='append', default=[], help='PATH:RATE[:CH], raw s16le (repeatable)')
    a = ap.parse_args()
    wai = os.path.abspath(a.wai)

    dirs = list(a.ref)
    env = os.environ.get('WAI_AUDIO_REFERENCE_DIRS', '')
    dirs += [d for d in env.split(os.pathsep) if d]
    if os.path.isdir(SYSTEM_SOUNDS):
        dirs.append(SYSTEM_SOUNDS)
    dirs = [d for d in dirs if os.path.isdir(d)]
    if not dirs:
        print(f'audio provenance: SKIPPED — no reference set on this host ({SYSTEM_SOUNDS} is absent and '
              'no --ref or WAI_AUDIO_REFERENCE_DIRS directory was given)')
        return 0
    conv = converter()
    with tempfile.TemporaryDirectory() as tmp:
        refs, rnotes = references(dirs, conv, tmp)
    for n in rnotes:
        print(f'  note: {n}')
    if not refs:
        print(f'audio provenance: SKIPPED — the reference set ({", ".join(dirs)}) has no file this host can '
              f'read (converter: {conv or "none; install afconvert or ffmpeg"})')
        return 0

    vs, notes, fails = vectors(wai, a.pcm)
    print(f'audio provenance: {len(vs)} committed audio signals against {len(refs)} references '
          f'({", ".join(dirs)}), at {ANALYSIS_RATE} Hz, threshold {a.threshold}')
    for n in notes:
        print(f'  note: {n}')

    rs = [(path, x, prefix_sq(x)) for path, x in refs]
    ref_spec = {}                                   # (reference, n) → its spectrum
    over = []
    for name, chans, rate in vs:
        for ci, c in enumerate(chans):
            x = resample(c, rate)
            label = name + (f' ch{ci}' if len(chans) > 1 else '')
            if len(x) < MIN_SAMPLES or not any(x):
                print(f'  {"-":>6}  {label}  ({len(x)} samples at {ANALYSIS_RATE} Hz'
                      f'{", silent" if not any(x) else ""}: too short or silent to judge)')
                continue
            px = prefix_sq(x)
            best = (0.0, None, 0)
            cache = {}
            for path, y, py in rs:
                n = 1 << (len(x) + len(y) - 1).bit_length()
                if n not in cache:
                    cache[n] = spectrum(x, n)
                if (path, n) not in ref_spec:
                    ref_spec[(path, n)] = spectrum(y, n)
                r, lag = best_ncc(x, cache[n], px, y, ref_spec[(path, n)], py, n)
                if r > best[0]:
                    best = (r, path, lag)
            flag = 'FAIL' if best[0] > a.threshold else 'ok'
            print(f'  {best[0]:6.3f}  {label}  (rate {rate}; nearest {os.path.basename(best[1] or "-")} '
                  f'at lag {best[2]})  {flag}')
            if best[0] > a.threshold:
                over.append(f'{label}: {best[0]:.3f} with {best[1]}')
    for f in fails:
        print(f'  FAIL: {f}')
    for o in over:
        print(f'  FAIL: {o}')
    if over or fails:
        print(f'audio provenance: FAILED — {len(over)} signal(s) above {a.threshold}, {len(fails)} uncovered')
        return 1
    print('audio provenance: ok — no committed audio vector correlates with a reference above the threshold')
    return 0


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