#!/usr/bin/env python3
"""
Find the time offset between a local audio file and a YouTube video's audio.

Uses Chromaprint audio fingerprinting to find candidate offsets, then
generates stereo comparison WAV files (left ear = local, right ear = YouTube)
so you can verify alignment by ear.

Usage:
    python find_offset.py <local_audio_path> <youtube_video_id>
    python find_offset.py <local_audio_path> --yt=-GY8MhHMxLM   # IDs starting with -

Example:
    python find_offset.py ../../public/songs/get-lucky/audio.ogg 5NV6Rdv1a3I

Requirements:
    pip install -r requirements.txt
    ffmpeg and fpcalc must be on PATH

Output:
    - Prints top candidate offsets ranked by fingerprint similarity
    - Writes stereo WAV files for each candidate to --out-dir
    - Listen with headphones: left ear = local, right ear = YouTube
    - The file where both ears are in sync has the correct offset
"""

import argparse
import json
import os
import subprocess
import sys
import tempfile
import wave

import numpy as np
from scipy import signal


SAMPLE_RATE = 16000  # for comparison WAV generation
FPCALC_RATE = 0.1238  # seconds per Chromaprint hash


def audio_to_pcm(path: str, sr: int = SAMPLE_RATE) -> np.ndarray:
    """Decode any audio file to mono float32 PCM using ffmpeg."""
    cmd = [
        "ffmpeg", "-i", path,
        "-ac", "1", "-ar", str(sr),
        "-f", "f32le", "-v", "quiet", "pipe:1",
    ]
    result = subprocess.run(cmd, capture_output=True)
    if result.returncode != 0:
        raise RuntimeError(f"ffmpeg failed on {path}: {result.stderr.decode()}")
    return np.frombuffer(result.stdout, dtype=np.float32)


def download_youtube_audio(video_id: str, out_dir: str) -> str:
    """Download audio from YouTube video, return path to file."""
    out_path = os.path.join(out_dir, "yt_audio.%(ext)s")
    cmd = [
        "yt-dlp", "-x", "--audio-format", "wav", "--audio-quality", "0",
        "-o", out_path, f"https://www.youtube.com/watch?v={video_id}",
    ]
    print(f"Downloading YouTube audio for {video_id}...", file=sys.stderr)
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        raise RuntimeError(f"yt-dlp failed: {result.stderr}")
    for f in os.listdir(out_dir):
        if f.startswith("yt_audio"):
            return os.path.join(out_dir, f)
    raise RuntimeError("yt-dlp produced no output file")


def get_fingerprint(path: str) -> np.ndarray:
    """Get raw Chromaprint fingerprint as array of uint32 hashes."""
    r = subprocess.run(
        ["fpcalc", "-raw", "-length", "0", path],
        capture_output=True, text=True,
    )
    for line in r.stdout.splitlines():
        if line.startswith("FINGERPRINT="):
            return np.array(
                [int(x) for x in line.split("=", 1)[1].split(",")],
                dtype=np.uint32,
            )
    raise RuntimeError(f"fpcalc failed: {r.stderr}")


def _popcount_mean(fp_a: np.ndarray, fp_b: np.ndarray) -> float:
    """Average number of matching bits between aligned fingerprint arrays."""
    xor = np.bitwise_xor(fp_a.astype(np.int64), fp_b.astype(np.int64))
    differing = np.zeros(len(xor), dtype=np.int32)
    for bit in range(32):
        differing += ((xor >> bit) & 1).astype(np.int32)
    return float(np.mean(32 - differing))


def find_candidates(fp_local: np.ndarray, fp_yt: np.ndarray) -> list:
    """
    Find candidate offsets using Chromaprint fingerprint matching.

    1. Build histogram of offset differences from 16-bit sub-hash matches.
    2. Verify top peaks by computing average Hamming similarity of full 32-bit
       hashes over the entire aligned overlap.
    3. Return candidates sorted by similarity (best first).
    """
    # Step 1: sub-hash histogram
    mask = np.uint32(0xFFFF)
    local_hashes: dict[int, list[int]] = {}
    for i, h in enumerate(fp_local):
        key = int(h & mask)
        local_hashes.setdefault(key, []).append(i)

    offsets = []
    for j, h in enumerate(fp_yt):
        key = int(h & mask)
        if key in local_hashes:
            for i in local_hashes[key]:
                offsets.append(j - i)

    if not offsets:
        return []

    offsets_arr = np.array(offsets)
    lo, hi = int(offsets_arr.min()), int(offsets_arr.max())
    bins = np.arange(lo, hi + 2)
    hist, edges = np.histogram(offsets_arr, bins=bins)

    # Extract top peaks (at least 10s apart)
    min_gap = int(10 / FPCALC_RATE)
    peaks = []
    hist_work = hist.copy()
    for _ in range(8):
        idx = np.argmax(hist_work)
        count = int(hist_work[idx])
        if count == 0:
            break
        offset_h = int(edges[idx])
        offset_sec = offset_h * FPCALC_RATE
        if offset_sec >= 0:
            peaks.append((offset_h, offset_sec, count))
        hist_work[max(0, idx - min_gap) : min(len(hist_work), idx + min_gap)] = 0

    # Step 2: verify with full Hamming similarity
    results = []
    for offset_h, offset_sec, count in peaks:
        yt_start = max(0, offset_h)
        loc_start = max(0, -offset_h)
        overlap = min(len(fp_local) - loc_start, len(fp_yt) - yt_start)
        if overlap < 10:
            continue
        avg_bits = _popcount_mean(
            fp_local[loc_start : loc_start + overlap],
            fp_yt[yt_start : yt_start + overlap],
        )
        results.append({
            "offset": round(offset_sec, 3),
            "hist_matches": count,
            "avg_matching_bits": round(avg_bits, 1),
        })

    return sorted(results, key=lambda x: -x["avg_matching_bits"])


def write_comparison_wav(
    path: str,
    local_pcm: np.ndarray,
    yt_pcm: np.ndarray,
    offset: float,
    sr: int = SAMPLE_RATE,
) -> None:
    """Write a stereo WAV: left = local audio, right = YouTube at offset."""
    start = max(0, int(offset * sr))
    end = start + len(local_pcm)
    if end > len(yt_pcm):
        end = len(yt_pcm)
    if start >= end:
        return
    yt_seg = yt_pcm[start:end]
    local_seg = local_pcm[: len(yt_seg)]
    if len(local_seg) == 0:
        return

    left = local_seg / (np.max(np.abs(local_seg)) + 1e-10) * 0.8
    right = yt_seg / (np.max(np.abs(yt_seg)) + 1e-10) * 0.8
    l16 = (left * 32767).astype(np.int16)
    r16 = (right * 32767).astype(np.int16)
    stereo = np.column_stack([l16, r16]).flatten()

    with wave.open(path, "w") as w:
        w.setnchannels(2)
        w.setsampwidth(2)
        w.setframerate(sr)
        w.writeframes(stereo.tobytes())


def main():
    parser = argparse.ArgumentParser(
        description="Find time offset between local audio and YouTube video audio."
    )
    parser.add_argument("local_audio", help="Path to the local audio file (e.g. audio.ogg)")
    parser.add_argument(
        "youtube_id", nargs="?", default=None,
        help="YouTube video ID (e.g. 5NV6Rdv1a3I)",
    )
    parser.add_argument(
        "--yt", dest="youtube_id_flag",
        help="YouTube video ID (use this for IDs starting with '-')",
    )
    parser.add_argument(
        "--out-dir", default=None,
        help="Directory to write comparison WAV files (default: ./compare_<id>/)",
    )
    parser.add_argument(
        "--json", action="store_true",
        help="Output result as JSON (for scripting)",
    )
    args = parser.parse_args()

    video_id = args.youtube_id_flag or args.youtube_id
    if not video_id:
        parser.error("youtube_id is required (use --yt=ID for IDs starting with '-')")

    if not os.path.exists(args.local_audio):
        print(f"Error: {args.local_audio} not found", file=sys.stderr)
        sys.exit(1)

    out_dir = args.out_dir or f"compare_{video_id}"
    os.makedirs(out_dir, exist_ok=True)

    with tempfile.TemporaryDirectory() as tmp:
        yt_path = download_youtube_audio(video_id, tmp)

        # Fingerprint matching
        print("Computing fingerprints...", file=sys.stderr)
        fp_local = get_fingerprint(args.local_audio)
        fp_yt = get_fingerprint(yt_path)
        print(
            f"Fingerprints: local={len(fp_local)} hashes "
            f"({len(fp_local)*FPCALC_RATE:.0f}s), "
            f"yt={len(fp_yt)} hashes ({len(fp_yt)*FPCALC_RATE:.0f}s)",
            file=sys.stderr,
        )

        candidates = find_candidates(fp_local, fp_yt)

        # Generate comparison WAVs
        print("Decoding audio for comparison files...", file=sys.stderr)
        local_pcm = audio_to_pcm(args.local_audio)
        yt_pcm = audio_to_pcm(yt_path)

        print(f"\nWriting comparison files to {out_dir}/", file=sys.stderr)
        print("  Left ear = local audio, Right ear = YouTube\n", file=sys.stderr)

        for c in candidates:
            label = f"{c['offset']:.1f}".replace(".", "_").replace("-", "neg")
            filename = f"offset_{label}s.wav"
            filepath = os.path.join(out_dir, filename)
            write_comparison_wav(filepath, local_pcm, yt_pcm, c["offset"])
            m, s = divmod(c["offset"], 60)
            print(
                f"  {filename:<28s}  offset={c['offset']:>8.3f}s ({int(m)}:{s:05.2f})  "
                f"bits={c['avg_matching_bits']:.1f}  hist={c['hist_matches']}",
                file=sys.stderr,
            )

        if args.json:
            print(json.dumps({
                "youtube_id": video_id,
                "candidates": candidates,
                "compare_dir": out_dir,
            }, indent=2))
        else:
            print(
                f"\nListen to each file with headphones. The one where both "
                f"ears are in sync is the correct offset.",
                file=sys.stderr,
            )
            if candidates:
                print(f"{candidates[0]['offset']:.3f}")


if __name__ == "__main__":
    main()
