#!/usr/bin/env python3
"""Align two recordings of the same take by their audio, and say how sure it is.

Filming yourself while a screen recorder captures the screen gives two files that
start at different moments and drift apart if the clocks differ. Aligning them by
eye costs an hour and is never exact; aligning them by audio is a cross-correlation
and takes seconds — both files heard the same voice.

What this adds beyond a plain correlation:

  * a **confidence** figure. The correlation peak means nothing on its own; what
    matters is how far it stands above the next-best peak. A ratio near 1 means the
    two files probably do not contain the same audio, and the report says so
    instead of returning a confident wrong number.

  * a **drift** check. The offset is measured independently in the first and last
    third of the overlap. If they disagree, the recorders ran at different sample
    clocks and a single offset cannot fix the whole take — a very common problem
    with phone-plus-OBS setups, and invisible until the end of a long video.

  * an **audio recommendation**, from measured loudness, noise floor and speech
    bandwidth, because the screen recorder's audio is usually the worse of the two
    and picking it wastes the good microphone.
"""
from __future__ import annotations

import argparse
import json
import math
import subprocess
import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent))

from media_probe import loudness, probe  # noqa: E402
from skill_config import emit, require_binary, utf8_stdout  # noqa: E402

RATE = 8_000  # plenty for envelope correlation, and four times faster than 16 kHz
HOP = 0.005


def decode(media: Path, rate: int = RATE) -> np.ndarray:
    completed = subprocess.run(
        ["ffmpeg", "-v", "error", "-nostdin", "-i", str(media),
         "-map", "0:a:0", "-ac", "1", "-ar", str(rate), "-f", "s16le", "-"],
        capture_output=True,
    )
    if completed.returncode or not completed.stdout:
        detail = (completed.stderr or b"").decode("utf-8", errors="replace")[-300:]
        raise SystemExit(f"No usable audio in {media}: {detail}")
    return np.frombuffer(completed.stdout, dtype="<i2").astype(np.float32) / 32768.0


def envelope(samples: np.ndarray, rate: int = RATE) -> np.ndarray:
    """Onset-weighted energy envelope.

    Correlating raw waveforms fails as soon as the two microphones differ in
    frequency response or one is compressed. An envelope of *rises* in energy keys
    on the timing of syllables, which both recordings share regardless of tone.
    """
    step = max(1, int(HOP * rate))
    usable = samples[: (samples.size // step) * step]
    if usable.size == 0:
        return np.zeros(1)
    frames = np.sqrt((usable.reshape(-1, step) ** 2).mean(axis=1) + 1e-12)
    log_frames = np.log(frames + 1e-6)
    rises = np.diff(log_frames, prepend=log_frames[0])
    rises = np.maximum(rises, 0.0)
    if rises.std() > 0:
        rises = (rises - rises.mean()) / rises.std()
    return rises


def correlate(first: np.ndarray, second: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Full cross-correlation via FFT, plus the lag axis in seconds."""
    size = 1
    while size < first.size + second.size:
        size *= 2
    spectrum = np.fft.rfft(first, size) * np.conj(np.fft.rfft(second, size))
    correlation = np.fft.irfft(spectrum, size)
    correlation = np.concatenate([correlation[-(second.size - 1):], correlation[:first.size]])
    lags = (np.arange(correlation.size) - (second.size - 1)) * HOP
    return correlation, lags


def best_offset(first: np.ndarray, second: np.ndarray,
                limit_seconds: float | None = None) -> dict:
    correlation, lags = correlate(first, second)
    if limit_seconds is not None:
        mask = np.abs(lags) <= limit_seconds
        correlation = correlation[mask]
        lags = lags[mask]
    if correlation.size == 0:
        return {"offset": 0.0, "confidence": 0.0, "peak": 0.0}
    peak_index = int(np.argmax(correlation))
    peak = float(correlation[peak_index])
    offset = float(lags[peak_index])

    # Confidence = how far the winner stands above the best *unrelated* peak.
    # Everything within 250 ms of the winner belongs to the same peak.
    exclusion = max(1, int(0.25 / HOP))
    masked = correlation.copy()
    low = max(0, peak_index - exclusion)
    high = min(correlation.size, peak_index + exclusion + 1)
    masked[low:high] = -np.inf
    runner_up = float(np.max(masked)) if np.isfinite(masked).any() else 0.0
    confidence = float(peak / runner_up) if runner_up > 0 else float("inf")
    return {
        "offset": round(offset, 4),
        "confidence": round(confidence, 3) if math.isfinite(confidence) else 99.0,
        "peak": round(peak, 2),
        "runner_up": round(runner_up, 2),
    }


def audio_quality(media: Path) -> dict:
    """Enough measurement to choose which microphone to keep."""
    samples = decode(media, 16_000)
    if samples.size == 0:
        return {"usable": False}
    frame = 1600
    usable = samples[: (samples.size // frame) * frame].reshape(-1, frame)
    levels = 20 * np.log10(np.maximum(np.sqrt((usable ** 2).mean(axis=1)), 1e-9))
    loud = float(np.percentile(levels, 90))
    floor = float(np.percentile(levels, 10))
    span = min(samples.size, 16_000 * 30)
    spectrum = np.abs(np.fft.rfft(samples[:span] * np.hanning(span)))
    frequencies = np.fft.rfftfreq(span, 1 / 16_000)
    total = float(spectrum.sum()) or 1e-9
    speech_band = float(spectrum[(frequencies >= 300) & (frequencies <= 3400)].sum() / total)
    # Where 95% of the energy ends. A screen-capture or telephone-band mic rolls off
    # near 3.4 kHz; a usable microphone carries well past 6 kHz. This is the figure
    # that separates them — the speech-band *ratio* does the opposite, because
    # band-limiting a signal pushes that ratio up and would score the worse mic
    # higher.
    cumulative = np.cumsum(spectrum) / total
    rolloff_index = int(np.searchsorted(cumulative, 0.95))
    rolloff = float(frequencies[min(rolloff_index, frequencies.size - 1)])
    measured = loudness(media)
    clipping = float(np.mean(np.abs(samples) > 0.995))
    return {
        "usable": True,
        "integrated_lufs": measured.get("integrated_lufs"),
        "true_peak_dbtp": measured.get("true_peak_dbtp"),
        "loud_level_db": round(loud, 1),
        "noise_floor_db": round(floor, 1),
        "signal_to_noise_db": round(loud - floor, 1),
        "speech_band_ratio": round(speech_band, 3),
        "rolloff_95_hz": round(rolloff),
        "clipped_fraction": round(clipping, 5),
    }


def recommend_audio(first: dict, second: dict, names: tuple[str, str]) -> dict:
    """Pick the better microphone from measurements, and explain the choice."""
    def score(item: dict) -> float:
        """Signal-to-noise plus real bandwidth, minus clipping."""
        if not item.get("usable"):
            return -99.0
        value = item["signal_to_noise_db"]
        # Bandwidth counted in octaves above the telephone band, a few dB each.
        octaves = max(0.0, math.log2(max(item["rolloff_95_hz"], 1) / 3400.0))
        value += 6.0 * octaves
        value -= 200 * item["clipped_fraction"]
        return value

    scores = (score(first), score(second))
    winner = 0 if scores[0] >= scores[1] else 1
    return {
        "use": names[winner],
        "scores": {names[0]: round(scores[0], 2), names[1]: round(scores[1], 2)},
        "why": (
            f"Сигнал/шум {(first, second)[winner]['signal_to_noise_db']} дБ, "
            f"полоса до {(first, second)[winner]['rolloff_95_hz']} Гц, "
            f"клиппинг {(first, second)[winner]['clipped_fraction']}. "
            f"Проигравший: сигнал/шум "
            f"{(first, second)[1 - winner]['signal_to_noise_db']} дБ, полоса до "
            f"{(first, second)[1 - winner]['rolloff_95_hz']} Гц."
        ),
        "note": (
            "Вторую дорожку не выбрасывай: она нужна для звуков интерфейса и "
            "как страховка при выпадении основного микрофона."
        ),
    }


def main() -> None:
    utf8_stdout()
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("camera", help="Recording of you (usually the better microphone)")
    parser.add_argument("screen", help="Screen recording of the same take")
    parser.add_argument("--out", default="dual_sync.json")
    parser.add_argument("--max-offset", type=float, default=120.0,
                        help="Largest plausible start difference, in seconds")
    parser.add_argument("--analyse-seconds", type=float, default=240.0,
                        help="How much audio to correlate; the whole file for short takes")
    parser.add_argument("--min-confidence", type=float, default=1.6,
                        help="Peak-to-runner-up ratio below which the result is rejected")
    args = parser.parse_args()

    require_binary("ffmpeg")
    camera = Path(args.camera).resolve()
    screen = Path(args.screen).resolve()
    for path in (camera, screen):
        if not path.is_file():
            raise SystemExit(f"Not found: {path}")

    camera_info = probe(camera)
    screen_info = probe(screen)
    for name, info in (("camera", camera_info), ("screen", screen_info)):
        if not info.get("has_audio"):
            raise SystemExit(
                f"{name} has no audio track, so the takes cannot be aligned by sound. "
                f"Use a clap or a slate and align manually."
            )

    camera_audio = decode(camera)
    screen_audio = decode(screen)
    limit = int(args.analyse_seconds * RATE)
    first = envelope(camera_audio[:limit])
    second = envelope(screen_audio[:limit])
    overall = best_offset(first, second, args.max_offset)

    # Drift: apply the global offset first, then look for a *residual* offset in the
    # head and tail of the region the files actually share. Comparing the same index
    # ranges of unaligned envelopes compares different content and invents drift
    # that is not there — on a pair offset by 6.4 s it reported ten seconds of it.
    shift = int(round(overall["offset"] / HOP))
    if shift >= 0:
        aligned_first, aligned_second = first[shift:], second
    else:
        aligned_first, aligned_second = first, second[-shift:]
    common = min(aligned_first.size, aligned_second.size)
    aligned_first = aligned_first[:common]
    aligned_second = aligned_second[:common]
    third = common // 3
    drift = None
    if third > int(15 / HOP):
        # A narrow residual window: real clock drift over a take is small, and a
        # wide search would again lock onto unrelated speech.
        window = 4.0
        head = best_offset(aligned_first[:third], aligned_second[:third], window)
        tail = best_offset(aligned_first[-third:], aligned_second[-third:], window)
        drift = {
            "measured_on": "общая часть после применения глобального смещения",
            "residual_head": head["offset"],
            "residual_tail": tail["offset"],
            "difference": round(tail["offset"] - head["offset"], 4),
            "head_confidence": head["confidence"],
            "tail_confidence": tail["confidence"],
            "shared_seconds": round(common * HOP, 2),
        }

    camera_quality = audio_quality(camera)
    screen_quality = audio_quality(screen)
    recommendation = recommend_audio(camera_quality, screen_quality, ("camera", "screen"))

    offset = overall["offset"]
    confident = overall["confidence"] >= args.min_confidence
    warnings: list[str] = []
    if not confident:
        warnings.append(
            f"Уверенность {overall['confidence']} ниже порога {args.min_confidence}: "
            f"пик корреляции почти не выделяется. Скорее всего в файлах разный звук — "
            f"проверь синхронизацию вручную по хлопку."
        )
    if drift and abs(drift["difference"]) > 0.25 and \
            min(drift["head_confidence"], drift["tail_confidence"]) >= args.min_confidence:
        span = max(1.0, drift["shared_seconds"] * 2 / 3)
        warnings.append(
            f"Остаточное смещение в начале {drift['residual_head']:+.2f}s, в конце "
            f"{drift['residual_tail']:+.2f}s — расхождение {drift['difference']:+.2f}s "
            f"на {drift['shared_seconds']:.0f}s общего материала. Записи идут с разной "
            f"частотой сэмплирования: одним сдвигом весь дубль не выровнять. "
            f"Растяни экран примерно atempo={1 + drift['difference'] / span:.6f}."
        )
    if abs(offset) > args.max_offset * 0.95:
        warnings.append("Смещение упёрлось в предел поиска — увеличь --max-offset.")

    # Positive offset means the camera's audio happens later than the screen's, so
    # the screen recording started first and has to be trimmed by that much.
    if offset >= 0:
        trim = {"camera_start": 0.0, "screen_start": round(offset, 4)}
        explanation = (
            f"Экран начал запись на {offset:.3f}s раньше: обрежь экран с начала на "
            f"{offset:.3f}s, камеру оставь как есть."
        )
    else:
        trim = {"camera_start": round(-offset, 4), "screen_start": 0.0}
        explanation = (
            f"Камера начала запись на {-offset:.3f}s раньше: обрежь камеру с начала на "
            f"{-offset:.3f}s, экран оставь как есть."
        )

    camera_duration = float(camera_info.get("duration") or 0.0)
    screen_duration = float(screen_info.get("duration") or 0.0)
    overlap = min(camera_duration - trim["camera_start"], screen_duration - trim["screen_start"])

    report = {
        "status": "ok" if confident and not warnings else
                  "low_confidence" if not confident else "ok_with_warnings",
        "camera": {"path": str(camera), "duration": camera_duration,
                   "video": camera_info.get("video", {}).get("display_width"),
                   "audio": camera_quality},
        "screen": {"path": str(screen), "duration": screen_duration,
                   "video": screen_info.get("video", {}).get("display_width"),
                   "audio": screen_quality},
        "offset_seconds": offset,
        "confidence": overall["confidence"],
        "confidence_meaning": (
            "Отношение главного пика корреляции к следующему по величине. "
            "Меньше 1.6 — результату верить нельзя."
        ),
        "trim_to_align": trim,
        "how_to_align": explanation,
        "overlap_seconds": round(overlap, 3),
        "drift_check": drift,
        "audio_recommendation": recommendation,
        "warnings": warnings,
        "next": (
            "python scripts/render_dual_layout.py --sync dual_sync.json "
            "--layout-plan layout_plan.json OUT.mp4"
        ),
    }
    Path(args.out).parent.mkdir(parents=True, exist_ok=True)
    Path(args.out).write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n",
                              encoding="utf-8")
    report["saved"] = str(Path(args.out).resolve())
    emit(report)
    if not confident:
        raise SystemExit(2)


if __name__ == "__main__":
    main()