import numpy as np
import numpy.typing as npt
import pytest
from src.atmos.lfe_check import (
    LFE_WARN_400_HZ_DBFS,
    LFE_WARN_2000_HZ_DBFS,
    LfeFftAccumulator,
    LfeMeasurement,
    apply_lfe_measurement,
)
from src.atmos.models import (
    SILENCE_FLOOR_DBFS,
    AtmosValidationBuilder,
    AtmosValidationFindingCode,
    AtmosValidationMetadata,
)

_LFE_400_KEY = AtmosValidationFindingCode.LFE_LEVEL_ABOVE_400HZ_TOO_HIGH.metadata_key
_LFE_2000_KEY = AtmosValidationFindingCode.LFE_LEVEL_ABOVE_2000HZ_TOO_HIGH.metadata_key

# One FFT bin at the analysis parameters (4096-point at 48 kHz). Tones placed on an exact bin
# avoid Hann scalloping loss, so the measured bin magnitude equals the tone amplitude and tests
# can assert dBFS values tightly.
_BIN_HZ = 48_000 / 4096  # 11.71875
_EXACT_BIN_500_HZ = 43 * _BIN_HZ  # 503.90625 — first-threshold band (>= 400 Hz, < 2000 Hz)
_EXACT_BIN_2500_HZ = 214 * _BIN_HZ  # 2507.8125 — second-threshold band (>= 2000 Hz)
_EXACT_BIN_47_HZ = 4 * _BIN_HZ  # 46.875 — legitimate LFE band, below both cutoffs


def _tone(
    frequency_hz: float, amplitude: float, sample_count: int, sample_rate_hz: int = 48_000
) -> npt.NDArray[np.float64]:
    t = np.arange(sample_count) / sample_rate_hz
    return amplitude * np.sin(2 * np.pi * frequency_hz * t)


def _ending_in_silence(tone: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
    # Fade out and end in silence: an abruptly-truncated tone leaks across the whole spectrum in
    # the final zero-padded FFT block — a block-edge artifact (present in Sony's method too), not
    # content the check measures. Real masters end in fades/room tone; synthetic signals must too.
    fade_sample_count = 2048
    faded = tone.copy()
    faded[-fade_sample_count:] *= 0.5 * (1 + np.cos(np.pi * np.arange(fade_sample_count) / fade_sample_count))
    return np.concatenate((faded, np.zeros(6144)))


def _measure(mono_samples: npt.NDArray[np.float64], source_sample_rate_hz: int = 48_000) -> LfeMeasurement:
    lfe_fft_accumulator = LfeFftAccumulator(source_sample_rate_hz)
    # Feed in render-sized blocks to exercise the streaming buffer across block boundaries.
    for start in range(0, len(mono_samples), 8192):
        lfe_fft_accumulator.add_block(mono_samples[start : start + 8192])
    return lfe_fft_accumulator.finish()


def _atmos_metadata() -> AtmosValidationMetadata:
    return AtmosValidationMetadata(
        container="Wave",
        codec="PCM",
        codec_id="1",
        channels=12,
        sample_rate_hz=48000,
        bits_per_sample=24,
        is_truncated=False,
        adm_profile="Dolby Atmos Master",
        duration_ms=180500,
    )


class TestThresholds:
    def test_sony_qc_parity_values(self) -> None:
        # Locks the gate values to Sony QC's (ventura_1_1.py THRESH_400_DB / THRESH_2000_DB) and the
        # silence floor to their -200 no-signal value.
        assert LFE_WARN_400_HZ_DBFS == -60.0
        assert LFE_WARN_2000_HZ_DBFS == -100.0
        assert SILENCE_FLOOR_DBFS == -200.0


class TestLfeFftAccumulator:
    def test_tone_above_400_hz_reads_its_amplitude(self) -> None:
        measurement = _measure(_ending_in_silence(_tone(_EXACT_BIN_500_HZ, 10 ** (-30 / 20), 48_000)))

        assert measurement.above_400hz_dbfs == pytest.approx(-30.0, abs=0.1)
        # 500 Hz leaks nothing measurable into the >= 2000 Hz band.
        assert measurement.above_2000hz_dbfs < LFE_WARN_2000_HZ_DBFS

    def test_tone_above_2000_hz_reads_in_both_bands(self) -> None:
        # A >= 2000 Hz bin is also >= 400 Hz, so the same peak appears in both maxima.
        measurement = _measure(_ending_in_silence(_tone(_EXACT_BIN_2500_HZ, 10 ** (-90 / 20), 48_000)))

        assert measurement.above_400hz_dbfs == pytest.approx(-90.0, abs=0.1)
        assert measurement.above_2000hz_dbfs == pytest.approx(-90.0, abs=0.1)

    def test_in_band_lfe_tone_stays_below_both_thresholds(self) -> None:
        # A legitimate full-scale LFE tone (sub-100 Hz): Hann leakage at 30+ bins is far below
        # both gates, so real LFE content never trips the check.
        measurement = _measure(_ending_in_silence(_tone(_EXACT_BIN_47_HZ, 1.0, 48_000)))

        assert measurement.above_400hz_dbfs < LFE_WARN_400_HZ_DBFS
        assert measurement.above_2000hz_dbfs < LFE_WARN_2000_HZ_DBFS

    @pytest.mark.parametrize(
        ("bin_index", "expected_400_hz_dbfs"),
        [
            # The >= 400 Hz band starts at bin 35 (410.15625 Hz, searchsorted-left of 400). A tone on
            # the last excluded bin reaches the band only via Hann leakage into its neighbor (exactly
            # half amplitude, -6.02 dB); a tone on the first included bin reads at full amplitude. An
            # off-by-one in the band start flips which reading appears.
            pytest.param(34, -36.0, id="bin_34_excluded_reads_leakage_only"),
            pytest.param(35, -30.0, id="bin_35_first_included_reads_amplitude"),
        ],
    )
    def test_400_hz_band_edge_bins(self, bin_index: int, expected_400_hz_dbfs: float) -> None:
        measurement = _measure(_ending_in_silence(_tone(bin_index * _BIN_HZ, 10 ** (-30 / 20), 48_000)))

        assert measurement.above_400hz_dbfs == pytest.approx(expected_400_hz_dbfs, abs=0.2)

    @pytest.mark.parametrize(
        ("bin_index", "expected_2000_hz_dbfs"),
        [
            # Same edge semantics for the >= 2000 Hz band, which starts at bin 171 (2003.90625 Hz).
            pytest.param(170, -96.0, id="bin_170_excluded_reads_leakage_only"),
            pytest.param(171, -90.0, id="bin_171_first_included_reads_amplitude"),
        ],
    )
    def test_2000_hz_band_edge_bins(self, bin_index: int, expected_2000_hz_dbfs: float) -> None:
        measurement = _measure(_ending_in_silence(_tone(bin_index * _BIN_HZ, 10 ** (-90 / 20), 48_000)))

        assert measurement.above_2000hz_dbfs == pytest.approx(expected_2000_hz_dbfs, abs=0.2)

    def test_silence_reads_the_silence_floor(self) -> None:
        measurement = _measure(np.zeros(48_000))

        assert measurement.above_400hz_dbfs == SILENCE_FLOOR_DBFS
        assert measurement.above_2000hz_dbfs == SILENCE_FLOOR_DBFS

    def test_signal_shorter_than_one_fft_block_is_still_measured(self) -> None:
        # The final partial block is zero-padded and analyzed (Sony's flush), so content in a
        # sub-4096-sample tail cannot escape the check.
        measurement = _measure(_tone(_EXACT_BIN_500_HZ, 1.0, 1_000))

        assert measurement.above_400hz_dbfs > LFE_WARN_400_HZ_DBFS

    def test_96_khz_source_is_resampled_to_the_analysis_rate(self) -> None:
        measurement = _measure(
            _ending_in_silence(_tone(_EXACT_BIN_500_HZ, 10 ** (-30 / 20), 96_000, sample_rate_hz=96_000)), 96_000
        )

        assert measurement.above_400hz_dbfs == pytest.approx(-30.0, abs=0.5)


class TestApplyLfeMeasurement:
    @pytest.mark.parametrize(
        ("above_400hz_dbfs", "above_2000hz_dbfs", "expect_400_warning", "expect_2000_warning"),
        [
            # Boundaries are strict >: exactly at a limit passes. The two gates are independent —
            # only_2000_over is the case the single-code design mishandled: a moderate high-frequency
            # peak breaches the 2 kHz gate (-100) without the 400 Hz gate (-60), since the >= 2000 Hz
            # band is a subset of the >= 400 Hz band and so never exceeds it.
            pytest.param(-60.0, -100.0, False, False, id="both_exactly_at_limits"),
            pytest.param(-59.9, -150.0, True, False, id="only_400_over"),
            pytest.param(-70.0, -99.9, False, True, id="only_2000_over"),
            pytest.param(-59.9, -99.9, True, True, id="both_over"),
            pytest.param(-150.0, -150.0, False, False, id="well_under_both"),
        ],
    )
    def test_warning_boundaries(
        self,
        above_400hz_dbfs: float,
        above_2000hz_dbfs: float,
        expect_400_warning: bool,
        expect_2000_warning: bool,
    ) -> None:
        atmos_validation_builder = AtmosValidationBuilder()

        apply_lfe_measurement(
            LfeMeasurement(above_400hz_dbfs=above_400hz_dbfs, above_2000hz_dbfs=above_2000hz_dbfs),
            atmos_validation_builder,
        )

        assert (_LFE_400_KEY in atmos_validation_builder.warnings) is expect_400_warning
        assert (_LFE_2000_KEY in atmos_validation_builder.warnings) is expect_2000_warning

    def test_stores_both_maxima_in_metadata(self) -> None:
        atmos_validation_builder = AtmosValidationBuilder()
        atmos_validation_builder.metadata = _atmos_metadata()

        apply_lfe_measurement(
            LfeMeasurement(above_400hz_dbfs=-52.1, above_2000hz_dbfs=-110.4),
            atmos_validation_builder,
        )

        assert atmos_validation_builder.metadata is not None
        assert atmos_validation_builder.metadata.lfe_above_400hz_dbfs == -52.1
        assert atmos_validation_builder.metadata.lfe_above_2000hz_dbfs == -110.4
        # -52.1 breaches the 400 Hz gate; -110.4 stays under the 2 kHz gate.
        assert _LFE_400_KEY in atmos_validation_builder.warnings
        assert _LFE_2000_KEY not in atmos_validation_builder.warnings

    def test_no_lfe_channel_skips_the_check(self) -> None:
        # No LFE channel in the ADM means nothing to measure (Sony: no-bed files PASS-skip).
        atmos_validation_builder = AtmosValidationBuilder()
        atmos_validation_builder.metadata = _atmos_metadata()

        apply_lfe_measurement(None, atmos_validation_builder)

        assert atmos_validation_builder.metadata is not None
        assert atmos_validation_builder.metadata.lfe_above_400hz_dbfs is None
        assert atmos_validation_builder.warnings == {}

    def test_no_metadata_still_emits_warnings(self) -> None:
        # metadata is None when the atmos mediainfo failed — the check still runs.
        atmos_validation_builder = AtmosValidationBuilder()

        apply_lfe_measurement(
            LfeMeasurement(above_400hz_dbfs=-10.0, above_2000hz_dbfs=-150.0),
            atmos_validation_builder,
        )

        assert _LFE_400_KEY in atmos_validation_builder.warnings