import math
import subprocess
from typing import Final

import numpy as np
import numpy.typing as npt
from pydantic import BaseModel, ConfigDict
from scipy.signal import correlate, correlation_lags, resample_poly

from src.atmos.models import AtmosValidationBuilder, AtmosValidationFindingCode
from src.atmos.render_errors import FfmpegError, FfmpegUnavailable

# Conform & sync: the Atmos is rendered to stereo (BS.2051 0+2+0) and cross-correlated
# against the stereo reference, both reduced to mono at 8 kHz. Sony QC parity — a single global lag
# over the first 60 s from sample 0 (leading silence included, no content-onset detection). The mono
# downmix checks program-level conform/sync only; per-channel defects (L/R swap, one-channel dropout)
# are invisible to it and out of scope for this conform/sync check.
_ANALYSIS_SAMPLE_RATE_HZ: Final = 8_000
_ANALYSIS_DURATION_S: Final = 60
_ANALYSIS_SAMPLES: Final = _ANALYSIS_DURATION_S * _ANALYSIS_SAMPLE_RATE_HZ

# Sony QC parity: warn when the global lag exceeds 50 ms.
SYNC_MISMATCH_WARN_MS: Final = 50.0
# Content-match guard: below this |r| the content doesn't match well enough to trust the lag, so warn
# CONTENT_MISMATCH and suppress the (meaningless) offset.
CONTENT_MISMATCH_WARN_R: Final = 0.50


class AlignmentMeasurement(BaseModel):
    model_config = ConfigDict(frozen=True)

    # Global cross-correlation lag in ms; None when |content_match_r| < CONTENT_MISMATCH_WARN_R (the
    # lag is meaningless when the content doesn't match). Positive = Atmos lags the stereo reference.
    alignment_offset_ms: float | None
    # Signed correlation at the best lag (see _cross_correlate for the normalization): |r| is the
    # content-match strength, the sign is polarity. Equals Pearson r near lag 0.
    content_match_r: float


def _to_alignment_measurement(lag_samples: int, content_match_r: float) -> AlignmentMeasurement:
    # Below the content-match floor the lag is meaningless (the reference may not even be the same
    # program), so suppress the offset and let apply_alignment_measurement warn CONTENT_MISMATCH.
    if abs(content_match_r) < CONTENT_MISMATCH_WARN_R:
        return AlignmentMeasurement(alignment_offset_ms=None, content_match_r=content_match_r)

    lag_ms = lag_samples / (_ANALYSIS_SAMPLE_RATE_HZ / 1000.0)
    return AlignmentMeasurement(alignment_offset_ms=lag_ms, content_match_r=content_match_r)


def _resample_to_8k(mono: npt.NDArray[np.float64], source_sample_rate_hz: int) -> npt.NDArray[np.float64]:
    # Atmos masters are always 48 or 96 kHz, so there's no 8 kHz passthrough branch — always resample
    # down to the analysis rate. gcd keeps the polyphase up/down factors small.
    divisor = math.gcd(source_sample_rate_hz, _ANALYSIS_SAMPLE_RATE_HZ)
    resampled: npt.NDArray[np.float64] = resample_poly(
        mono, _ANALYSIS_SAMPLE_RATE_HZ // divisor, source_sample_rate_hz // divisor
    )
    return resampled


def _decode_head_mono_8k(stereo_reference_url: str) -> npt.NDArray[np.float64]:
    # The stereo reference is a compressed file (e.g. FLAC), so decode it with ffmpeg rather than the
    # raw-PCM render path. First 60 s from sample 0 (leading silence included), downmixed to mono.
    command = [
        "ffmpeg",
        "-nostdin",
        "-hide_banner",
        "-v",
        "error",
        "-t",
        str(_ANALYSIS_DURATION_S),
        "-i",
        stereo_reference_url,
        "-af",
        "pan=mono|c0=0.5*c0+0.5*c1",
        "-ar",
        str(_ANALYSIS_SAMPLE_RATE_HZ),
        "-f",
        "f32le",
        "-",
    ]
    try:
        process = subprocess.run(command, capture_output=True, timeout=300)
    except OSError as e:
        raise FfmpegUnavailable(f"could not execute ffmpeg: {e}") from e
    except subprocess.TimeoutExpired as e:
        raise FfmpegError(f"ffmpeg timed out decoding the stereo reference: {e}") from e

    if process.returncode != 0:
        stderr = process.stderr.decode("utf-8", errors="replace").strip()[-500:]
        raise FfmpegError(f"ffmpeg exited with code {process.returncode}: {stderr}")

    try:
        samples = np.frombuffer(process.stdout, dtype=np.float32).astype(np.float64)
    except ValueError as e:
        raise FfmpegError(f"ffmpeg produced a malformed f32le stream: {e}") from e

    # An empty decode (ffmpeg exits 0 but the reference is header-only or unreadable in the first 60 s)
    # would otherwise reach _cross_correlate and surface as a bare scipy IndexError — categorize it
    # here, mirroring the atmos side's empty-master guard in run_render_checks.
    if samples.size == 0:
        raise FfmpegError("decoded stereo reference is empty (0 samples)")

    # A non-finite sample (e.g. a float source carrying NaN/inf) would poison the correlation into a
    # NaN r, which silently defeats every `<`/`>` gate downstream — reject it here instead.
    if not np.all(np.isfinite(samples)):
        raise FfmpegError("decoded stereo reference contains non-finite samples")

    return samples


def _cross_correlate(
    reference: npt.NDArray[np.float64],
    test: npt.NDArray[np.float64],
) -> tuple[int, float]:
    # Single global lag by cross-correlation over the full analysis window. Both inputs are
    # z-normalized, so dividing the peak by the full window length gives Pearson r at lag 0, decaying
    # by (length-|lag|)/length as the lag grows. That decay is deliberate: it down-weights spurious
    # large-lag peaks (dividing by the shrinking per-lag overlap would blow up their variance instead)
    # and is negligible (<0.5%) at the sub-second lags we expect. The signed extremum by |magnitude|
    # gives the lag (offset) and its sign (polarity). No windowing — matching the Sony QC approach.
    length = min(len(reference), len(test))
    reference = _z_normalize(reference[:length])
    test = _z_normalize(test[:length])

    correlation = correlate(test, reference, mode="full", method="fft")
    lags = correlation_lags(len(test), len(reference), mode="full")
    peak = int(np.argmax(np.abs(correlation)))

    return int(lags[peak]), float(correlation[peak] / length)


def _z_normalize(signal: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
    normalized: npt.NDArray[np.float64] = (signal - signal.mean()) / (signal.std() + 1e-12)
    return normalized


def apply_alignment_measurement(
    alignment_measurement: AlignmentMeasurement,
    atmos_validation_builder: AtmosValidationBuilder,
) -> None:
    atmos_validation_builder.update_metadata(
        stereo_reference_alignment_offset_ms=alignment_measurement.alignment_offset_ms,
        stereo_reference_content_match_r=alignment_measurement.content_match_r,
    )

    _check_content_match(alignment_measurement, atmos_validation_builder)
    _check_sync(alignment_measurement, atmos_validation_builder)


def _check_content_match(
    alignment_measurement: AlignmentMeasurement,
    atmos_validation_builder: AtmosValidationBuilder,
) -> None:
    if abs(alignment_measurement.content_match_r) < CONTENT_MISMATCH_WARN_R:
        atmos_validation_builder.warning(
            AtmosValidationFindingCode.CONTENT_MISMATCH,
            f"Stereo reference correlation is {alignment_measurement.content_match_r:.2f}, "
            f"below {CONTENT_MISMATCH_WARN_R} — cannot confirm the stereo reference is the same "
            f"program as the Atmos master",
        )


def _check_sync(
    alignment_measurement: AlignmentMeasurement,
    atmos_validation_builder: AtmosValidationBuilder,
) -> None:
    if alignment_measurement.alignment_offset_ms is None:
        return

    if abs(alignment_measurement.alignment_offset_ms) > SYNC_MISMATCH_WARN_MS:
        atmos_validation_builder.warning(
            AtmosValidationFindingCode.SYNC_MISMATCH,
            f"Atmos and stereo reference are misaligned by {alignment_measurement.alignment_offset_ms:.0f} ms, "
            f"should not exceed {SYNC_MISMATCH_WARN_MS} ms",
        )