#!/usr/bin/env python3
"""Measure a recording's A/V sync without listening to it.

Prints the three numbers that matter, from the file itself:
  1. how long each stream actually is (a length mismatch is a clock problem),
  2. how far apart they drift per second (a rate problem),
  3. the residual offset between mouth motion and speech (a constant problem).

Usage:  avsync.py <recording.mp4> [--ffmpeg <path>] [--crop WxH+X+Y]

Exit code is 0 when the file is measurable, 1 when it is not.
"""
from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys

import numpy as np

DEFAULT_FFMPEG = r"C:\Users\gramp\AppData\Roaming\ytLlive\tools\ffmpeg.exe"


def host_path(path: str) -> str:
    """Accept a Windows path under WSL (and a native path everywhere else)."""
    if os.name != "nt" and len(path) > 2 and path[1] == ":":
        return f"/mnt/{path[0].lower()}/{path[2:].replace(chr(92), '/').lstrip('/')}"
    return path


def run(ffmpeg: str, args: list[str]) -> bytes:
    proc = subprocess.run([host_path(ffmpeg), "-v", "error", *args], capture_output=True)
    if proc.returncode != 0:
        sys.stderr.write(proc.stderr.decode("utf-8", "replace"))
        raise SystemExit(1)
    return proc.stdout


def probe(ffprobe: str, path: str) -> dict:
    out = run(ffprobe, ["-show_entries",
                        "stream=index,codec_type,duration",
                        "-show_entries", "format=duration",
                        "-of", "json", path])
    return json.loads(out)


def stream_durations(meta: dict) -> tuple[float, float]:
    video = audio = 0.0
    for s in meta.get("streams", []):
        d = float(s.get("duration") or 0)
        if s.get("codec_type") == "video" and d:
            video = max(video, d)
        elif s.get("codec_type") == "audio" and d:
            audio = max(audio, d)
    return video, audio


def envelopes(ffmpeg: str, path: str, crop: str | None):
    audio_bytes = run(ffmpeg, ["-i", path, "-map", "0:a", "-ac", "1", "-ar", "48000",
                               "-f", "s16le", "-"])
    audio = np.frombuffer(audio_bytes, dtype=np.int16).astype(np.float64) / 32768.0

    filters = []
    if crop:
        # Accept the everyday "WxH+X+Y" geometry and translate it to the filter's w:h:x:y.
        spec = crop.replace("x", ":", 1).replace("+", ":")
        if ":" not in spec and crop.count(":") == 3:
            spec = crop
        filters.append(f"crop={spec}")
    filters += ["scale=64:64", "format=gray"]
    video_bytes = run(ffmpeg, ["-i", path, "-map", "0:v", "-vf", ",".join(filters),
                               "-f", "rawvideo", "-"])
    frames = np.frombuffer(video_bytes, dtype=np.uint8).reshape(-1, 64, 64).astype(np.float64)
    motion = np.abs(np.diff(frames, axis=0)).mean(axis=(1, 2))
    motion = np.concatenate([[motion[0] if len(motion) else 0.0], motion])
    return audio, motion


def offset_ms(audio: np.ndarray, motion: np.ndarray, fps: float, max_ms: float = 400.0):
    """Cross-correlate mouth motion against speech energy on a 100 Hz grid.

    Returns (lag_ms, correlation) where a POSITIVE lag means the sound arrives later than the
    mouth movement (audio delayed / video ahead).
    """
    grid = 100.0
    duration = min(len(audio) / 48000.0, len(motion) / fps)
    t = np.arange(0, duration, 1.0 / grid)

    def resample(x: np.ndarray, rate: float) -> np.ndarray:
        return x[np.clip((t * rate).astype(int), 0, len(x) - 1)]

    win = max(1, int(48000 / grid))
    n = len(audio) // win
    audio_rms = np.sqrt((audio[: n * win].reshape(n, win) ** 2).mean(axis=1))
    audio_env = np.interp(t, np.arange(n) / grid, audio_rms)

    def smooth(x: np.ndarray) -> np.ndarray:
        k = max(3, int(grid * 0.05) | 1)
        return np.convolve(x, np.ones(k) / k, mode="same")

    def z(x: np.ndarray) -> np.ndarray:
        s = x.std()
        return (x - x.mean()) / (s if s > 1e-9 else 1.0)

    m = z(smooth(resample(motion, fps)))
    a = z(smooth(audio_env))

    limit = int(grid * max_ms / 1000)

    def best_for(m_seg: np.ndarray, a_seg: np.ndarray) -> tuple[float, float]:
        best_corr, best_lag = -9.0, 0
        for k in range(-limit, limit + 1):
            a_part = a_seg[max(0, k): len(a_seg) + min(0, k)]
            m_part = m_seg[max(0, -k): len(m_seg) - max(0, k)]
            corr = float((m_part * a_part).mean())
            if corr > best_corr:
                best_corr, best_lag = corr, k
        return best_lag / grid * 1000.0, best_corr

    # Whole-clip first; then in 3 s windows, because a periodic voice (a spoken alphabet, a
    # song) can lock the global search onto the wrong beat. Only windows that correlate
    # strongly are trusted, and their median is the answer.
    overall_lag, overall_corr = best_for(m, a)
    seg = int(3 * grid)
    windows = []
    for start in range(0, max(1, len(m) - seg), seg):
        lag, corr = best_for(m[start:start + seg], a[start:start + seg])
        windows.append((start / grid, lag, corr))
    strong = sorted(lag for _, lag, corr in windows if corr >= 0.35)
    if strong:
        mid = len(strong) // 2
        robust = strong[mid] if len(strong) % 2 else (strong[mid - 1] + strong[mid]) / 2
        return robust, max(corr for _, _, corr in windows), windows
    return overall_lag, overall_corr, windows


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("recording")
    ap.add_argument("--ffmpeg", default=os.environ.get("YTLLIVE_FFMPEG", DEFAULT_FFMPEG))
    ap.add_argument("--crop", default=None,
                    help="webcam region, e.g. 764x764+1032+159 (defaults to the whole frame)")
    ap.add_argument("--max-ms", type=float, default=400.0)
    args = ap.parse_args()

    ffmpeg = args.ffmpeg
    ffprobe = ffmpeg.replace("ffmpeg.exe", "ffprobe.exe")
    if not os.path.exists(host_path(ffmpeg)):
        sys.stderr.write(f"ffmpeg not found: {ffmpeg}\n")
        return 1

    video, audio = stream_durations(probe(ffprobe, args.recording))
    if not video or not audio:
        sys.stderr.write("could not read both stream durations\n")
        return 1

    diff = audio - video
    pct = diff / video * 100.0

    samples, motion = envelopes(ffmpeg, args.recording, args.crop)
    fps = len(motion) / (len(samples) / 48000.0) if len(motion) else 0.0
    lag, corr, windows = offset_ms(samples, motion, fps, args.max_ms) if fps > 0 else (0.0, 0.0, [])

    print(f"video   {video:8.3f} s")
    print(f"audio   {audio:8.3f} s")
    print(f"delta   {diff:+8.3f} s  ({pct:+.2f}% of the video length)")
    verdict = "lengths match" if abs(diff) < 0.05 else "LENGTH MISMATCH — a clock problem"
    print(f"        {verdict}")
    quality = "strong" if corr >= 0.35 else "weak — treat the offset as a hint only"
    for start, wlag, wcorr in windows:
        mark = "*" if wcorr >= 0.35 else " "
        print(f"        {start:5.1f}s {wlag:+7.0f} ms  corr {wcorr:.2f} {mark}")
    print(f"offset  {lag:+8.0f} ms  (correlation {corr:.2f}, {quality})")
    direction = "sound LATER than the lips" if lag > 0 else "sound EARLIER than the lips"
    if abs(lag) >= 20:
        print(f"        {direction} — adjust the mic delay by {-lag:+.0f} ms")
    else:
        print("        within a frame or two — leave it alone")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())