import math
from collections.abc import Callable, Iterable
from pathlib import Path
from unittest.mock import MagicMock, patch

import numpy as np
import numpy.typing as npt
import pytest
from src.atmos.alignment_check import AlignmentMeasurement
from src.atmos.lfe_check import LFE_WARN_400_HZ_DBFS, LFE_WARN_2000_HZ_DBFS
from src.atmos.loudness_check import LoudnessMeasurement
from src.atmos.render_checks import RenderCheckMeasurements, measure_render_checks, run_render_checks
from src.atmos.render_errors import RenderOutputError
from src.atmos.silent_object_check import SilentObjectMeasurement

_FIXTURES_DIR = Path(__file__).parent.parent.parent / "fixtures"
_TEST_BWF = _FIXTURES_DIR / "test_bwf.wav"

_TTL_SECONDS = "src.atmos.render_checks.presigned_url_ttl_seconds"
_PRESIGN = "src.atmos.render_checks.s3.create_presigned_url"
_OPEN_S3 = "src.atmos.render_checks.s3.open_s3_seekable"
_MEASURE = "src.atmos.render_checks.measure_render_checks"

_ATMOS_RENDER = "src.atmos.render_checks.AtmosRender"
_MEASURE_LOUDNORM = "src.atmos.render_checks._measure_loudnorm"
_DECODE_MONO = "src.atmos.render_checks._decode_head_mono_8k"
_XCORR = "src.atmos.render_checks._cross_correlate"
_RESAMPLE = "src.atmos.render_checks._resample_to_8k"

_STEREO_URL = "https://example.com/stereo.flac"


def _fake_render(
    input_blocks: list[npt.NDArray[np.float64]],
    *,
    sample_rate_hz: int = 48000,
    render_block: Callable[[str, npt.NDArray[np.float64]], npt.NDArray[np.float64]] | None = None,
    lfe_track_indices: tuple[int, ...] = (),
    object_track_indices: tuple[int, ...] = (),
    bed_height_track_indices: tuple[int, ...] = (),
    input_channel_count: int = 10,
) -> MagicMock:
    # A mocked AtmosRender: input_blocks() yields the given raw blocks, and render_block() maps a
    # layout + input to output PCM (default: 6ch of ones for 0+5+0, 2ch of ones for 0+2+0).
    def default_render_block(layout: str, samples: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
        channels = 6 if layout == "0+5+0" else 2
        return np.ones((len(samples), channels), dtype=np.float64)

    render = MagicMock()
    render.sample_rate_hz = sample_rate_hz
    render.output_channel_count.return_value = 6
    render.input_blocks.return_value = iter(input_blocks)
    render.render_block.side_effect = render_block or default_render_block
    render.tail.return_value = np.zeros((0, 6), dtype=np.float64)
    render.lfe_track_indices = lfe_track_indices
    render.object_track_indices = object_track_indices
    render.bed_height_track_indices = bed_height_track_indices
    render.input_channel_count = input_channel_count
    return render


def _draining_loudnorm(loudness: LoudnessMeasurement) -> Callable[..., LoudnessMeasurement]:
    # _measure_loudnorm drives the loudness_blocks generator (in production, on the feeder thread) —
    # draining it here is what triggers the source-channel capture side effects the checks read.
    def drain(blocks: Iterable[npt.NDArray[np.float64]], **_: object) -> LoudnessMeasurement:
        for _block in blocks:
            pass
        return loudness

    return drain


def _loudness() -> LoudnessMeasurement:
    return LoudnessMeasurement(integrated_loudness_lkfs=-20.0, true_peak_dbtp=-3.0)


class TestRunRenderChecks:
    def test_presigns_stereo_opens_atmos_and_measures(self) -> None:
        # The picklable callable a spawn-based ProcessPoolExecutor dispatches to a worker process.
        render_check_measurements = RenderCheckMeasurements(
            loudness=LoudnessMeasurement(integrated_loudness_lkfs=-18.0, true_peak_dbtp=-2.0),
            alignment=AlignmentMeasurement(alignment_offset_ms=0.0, content_match_r=0.95),
            lfe=None,
            silent_object=SilentObjectMeasurement(
                object_track_count=0,
                silent_object_track_indices=(),
                height_track_count=0,
                silent_height_track_indices=(),
            ),
        )
        with (
            patch(_TTL_SECONDS, return_value=1800) as ttl_seconds,
            patch(_PRESIGN, return_value="https://example.com/stereo") as presign,
            patch(_OPEN_S3) as open_s3,
            patch(_MEASURE, return_value=render_check_measurements) as measure,
        ):
            result = run_render_checks("atmos-bucket", "atmos/key.wav", "stereo-bucket", "stereo/key.flac")

        # The presign TTL is derived from the message-processing budget, matching the validator.
        ttl_seconds.assert_called_once_with()
        presign.assert_called_once_with("stereo-bucket", "stereo/key.flac", expires_in_seconds=1800)
        open_s3.assert_called_once_with("atmos-bucket", "atmos/key.wav")
        measure.assert_called_once()
        assert result is render_check_measurements


class TestMeasureRenderChecks:
    def test_drives_both_layouts_and_returns_all_measurements(self) -> None:
        # One shared parse feeds every check: the 5.1 blocks stream to loudnorm, the stereo render
        # is captured as mono for the cross-correlation, and the raw source channels feed the LFE
        # and silent-object scans — all from a single pass over the input blocks.
        render = MagicMock()
        render.sample_rate_hz = 48000
        render.output_channel_count.return_value = 6
        input_samples = np.ones((8192, 10), dtype=np.float64)
        render.input_blocks.return_value = iter([input_samples])
        render.render_block.side_effect = lambda layout, samples: (
            np.ones((len(samples), 6), dtype=np.float64)
            if layout == "0+5+0"
            else np.ones((len(samples), 2), dtype=np.float64)
        )
        render.tail.return_value = np.zeros((0, 6), dtype=np.float64)
        render.lfe_track_indices = ()
        render.object_track_indices = (0, 1)
        render.bed_height_track_indices = ()
        render.input_channel_count = 10

        def drain_loudnorm(blocks: object, **_: object) -> LoudnessMeasurement:
            # _measure_loudnorm drives the generator (in production, on the feeder thread) — draining
            # it here is what triggers the source-channel capture side effects.
            for _block in blocks:  # type: ignore[attr-defined]
                pass
            return _loudness()

        with (
            patch(_ATMOS_RENDER, return_value=render),
            patch(_MEASURE_LOUDNORM, side_effect=drain_loudnorm) as measure_loudnorm,
            patch(_DECODE_MONO, return_value=np.zeros(80, dtype=np.float64)) as decode_mono,
            patch(_XCORR, return_value=(80, 0.95)) as cross_correlate,
        ):
            render_check_measurements = measure_render_checks(MagicMock(), "https://example.com/stereo.flac")

        # Both layouts were rendered from the shared parse.
        rendered_layouts = {call.args[0] for call in render.render_block.call_args_list}
        assert rendered_layouts == {"0+5+0", "0+2+0"}

        # The 5.1 blocks were fed to loudnorm with the 5.1 channel layout.
        measure_loudnorm.assert_called_once()
        assert measure_loudnorm.call_args.kwargs["channel_count"] == 6
        assert measure_loudnorm.call_args.kwargs["ffmpeg_channel_layout"] == "5.1"

        # The stereo reference was decoded and cross-correlated against the captured atmos mono.
        decode_mono.assert_called_once_with("https://example.com/stereo.flac")
        cross_correlate.assert_called_once()

        assert render_check_measurements.loudness.integrated_loudness_lkfs == -20.0
        assert render_check_measurements.alignment.content_match_r == 0.95
        assert render_check_measurements.alignment.alignment_offset_ms == 10.0
        # No LFE channel in the (mocked) ADM; the all-ones object channels are clearly active.
        assert render_check_measurements.lfe is None
        assert render_check_measurements.silent_object.object_track_count == 2
        assert render_check_measurements.silent_object.silent_object_track_indices == ()
        # The mocked render declares no bed height channels.
        assert render_check_measurements.silent_object.height_track_count == 0

    def test_captured_atmos_mono_reaches_cross_correlate_as_real_content(self) -> None:
        # The stereo render captured in the single pass must reach _cross_correlate as the REAL mono
        # content — not silence, not the wrong buffer. Feed a known non-silent 0+2+0 render and assert
        # the atmos_mono argument to _cross_correlate is that content, downmixed and resampled.
        frame_count = 8000  # 1 s at the tiny 8 kHz sample rate below -> resamples 1:1 to 8 kHz

        def render_block(layout: str, samples: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
            if layout == "0+2+0":
                # Both stereo channels carry the same ramp, so the mono downmix equals the ramp.
                ramp = np.linspace(0.1, 0.9, len(samples), dtype=np.float64)
                return np.stack([ramp, ramp], axis=1)
            return np.ones((len(samples), 6), dtype=np.float64)

        render = _fake_render(
            [np.ones((frame_count, 10), dtype=np.float64)],
            sample_rate_hz=8000,
            render_block=render_block,
        )

        with (
            patch(_ATMOS_RENDER, return_value=render),
            patch(_MEASURE_LOUDNORM, side_effect=_draining_loudnorm(_loudness())),
            patch(_DECODE_MONO, return_value=np.zeros(frame_count, dtype=np.float64)),
            patch(_XCORR, return_value=(0, 0.95)) as cross_correlate,
        ):
            measure_render_checks(MagicMock(), _STEREO_URL)

        cross_correlate.assert_called_once()
        atmos_mono_arg = cross_correlate.call_args.args[1]
        expected_mono = np.linspace(0.1, 0.9, frame_count, dtype=np.float64)
        assert len(atmos_mono_arg) == frame_count
        assert np.any(atmos_mono_arg != 0.0)
        np.testing.assert_allclose(atmos_mono_arg, expected_mono, atol=1e-9)

    def test_stereo_rendered_only_for_head_while_loudness_continues(self) -> None:
        # 0+2+0 is rendered only until the 60 s head cap trips; 0+5+0 keeps rendering the whole track.
        # A tiny sample rate shrinks alignment_head_frames so the cap trips after the first block.
        block_count = 4
        frames_per_block = 100
        sample_rate_hz = 1  # alignment_head_frames = 60 * 1 = 60 frames -> tripped by the first block

        render = _fake_render(
            [np.ones((frames_per_block, 10), dtype=np.float64) for _ in range(block_count)],
            sample_rate_hz=sample_rate_hz,
        )

        with (
            patch(_ATMOS_RENDER, return_value=render),
            patch(_MEASURE_LOUDNORM, side_effect=_draining_loudnorm(_loudness())),
            patch(_DECODE_MONO, return_value=np.zeros(60, dtype=np.float64)),
            patch(_XCORR, return_value=(0, 0.95)),
        ):
            measure_render_checks(MagicMock(), _STEREO_URL)

        layout_calls = [call.args[0] for call in render.render_block.call_args_list]
        stereo_calls = layout_calls.count("0+2+0")
        loudness_calls = layout_calls.count("0+5+0")
        assert stereo_calls < loudness_calls
        assert stereo_calls == 1  # only the first block is within the head cap
        assert loudness_calls == block_count

    def test_captured_atmos_mono_is_trimmed_to_the_head_cap(self) -> None:
        # The block that crosses the 60 s cap is captured WHOLE, so the concatenated mono overshoots
        # alignment_head_frames; the [:alignment_head_frames] trim must drop the excess before
        # cross-correlation. Two blocks that together overshoot the cap, each with a distinct mono
        # value, let us assert the tail beyond the cap is gone (resample patched to identity so the
        # assertion isolates the trim from the 8 kHz conversion).
        sample_rate_hz = 1  # alignment_head_frames = 60 * 1 = 60 frames

        def render_block(layout: str, samples: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
            if layout == "0+2+0":
                mono = samples[:, 0]  # carry the input's marker value through to the mono downmix
                return np.stack([mono, mono], axis=1)
            return np.ones((len(samples), 6), dtype=np.float64)

        render = _fake_render(
            [np.full((40, 10), 1.0, dtype=np.float64), np.full((40, 10), 2.0, dtype=np.float64)],
            sample_rate_hz=sample_rate_hz,
            render_block=render_block,
        )

        with (
            patch(_ATMOS_RENDER, return_value=render),
            patch(_MEASURE_LOUDNORM, side_effect=_draining_loudnorm(_loudness())),
            patch(_RESAMPLE, side_effect=lambda mono, *_: mono),
            patch(_DECODE_MONO, return_value=np.zeros(60, dtype=np.float64)),
            patch(_XCORR, return_value=(0, 0.95)) as cross_correlate,
        ):
            measure_render_checks(MagicMock(), _STEREO_URL)

        atmos_mono_arg = cross_correlate.call_args.args[1]
        assert len(atmos_mono_arg) == 60  # trimmed to the head cap, not the 80 frames captured
        assert np.count_nonzero(atmos_mono_arg == 1.0) == 40  # block 0 is wholly within the cap
        assert np.count_nonzero(atmos_mono_arg == 2.0) == 20  # block 1's tail beyond the cap is dropped

    def test_source_channel_scans_consume_the_whole_track_not_just_the_head(self) -> None:
        # The stereo capture stops at the 60 s head cap, but the LFE FFT and the per-channel peak
        # scan must see every block: an object that only sounds after the cap is NOT silent, and
        # LFE content after the cap still counts.
        sample_rate_hz = 1  # head cap = 60 frames -> the second block is entirely past it
        quiet_block = np.zeros((60, 10), dtype=np.float64)
        loud_late_block = np.zeros((60, 10), dtype=np.float64)
        loud_late_block[:, 5] = 0.5  # the object channel only sounds in the post-cap block

        render = _fake_render(
            [quiet_block, loud_late_block],
            sample_rate_hz=sample_rate_hz,
            lfe_track_indices=(3,),
            object_track_indices=(5,),
        )

        with (
            patch(_ATMOS_RENDER, return_value=render),
            patch(_MEASURE_LOUDNORM, side_effect=_draining_loudnorm(_loudness())),
            patch(_DECODE_MONO, return_value=np.zeros(60, dtype=np.float64)),
            patch(_XCORR, return_value=(0, 0.95)),
            patch("src.atmos.render_checks.LfeFftAccumulator") as lfe_accumulator_class,
        ):
            render_check_measurements = measure_render_checks(MagicMock(), _STEREO_URL)

        assert render_check_measurements.silent_object.silent_object_track_indices == ()
        # The single LFE accumulator received both blocks' LFE column, not just the pre-cap head.
        lfe_accumulator_class.assert_called_once_with(sample_rate_hz)
        added_blocks = lfe_accumulator_class.return_value.add_block.call_args_list
        assert len(added_blocks) == 2
        assert sum(len(call.args[0]) for call in added_blocks) == 120

    def test_lfe_measurement_aggregates_worst_per_band_across_lfe_channels(self) -> None:
        # Two LFE channels (a two-bed master), each carrying its band's worst content on a DIFFERENT
        # channel: the roll-up is the per-band maximum across channels (Sony's worst-across-beds),
        # not one worst channel's pair. Tones sit on exact FFT bins and fade to silence so the
        # zero-padded flush block adds no broadband leakage.
        frame_count = 8192

        def faded_tone(frequency_hz: float, amplitude: float) -> npt.NDArray[np.float64]:
            tone = amplitude * np.sin(2 * np.pi * frequency_hz * np.arange(frame_count) / 48000)
            fade_sample_count = 2048
            tone[-fade_sample_count:] *= 0.5 * (1 + np.cos(np.pi * np.arange(fade_sample_count) / fade_sample_count))
            return tone

        input_block = np.zeros((frame_count, 10), dtype=np.float64)
        input_block[:, 7] = faded_tone(503.90625, 10 ** (-30 / 20))  # 400 Hz band's worst on channel 7
        input_block[:, 3] = faded_tone(2507.8125, 10 ** (-90 / 20))  # 2000 Hz band's worst on channel 3

        render = _fake_render([input_block, np.zeros((6144, 10), dtype=np.float64)], lfe_track_indices=(3, 7))

        with (
            patch(_ATMOS_RENDER, return_value=render),
            patch(_MEASURE_LOUDNORM, side_effect=_draining_loudnorm(_loudness())),
            patch(_DECODE_MONO, return_value=np.zeros(8000, dtype=np.float64)),
            patch(_XCORR, return_value=(0, 0.95)),
        ):
            render_check_measurements = measure_render_checks(MagicMock(), _STEREO_URL)

        assert render_check_measurements.lfe is not None
        assert render_check_measurements.lfe.above_400hz_dbfs == pytest.approx(-30.0, abs=0.5)
        assert render_check_measurements.lfe.above_2000hz_dbfs == pytest.approx(-90.0, abs=0.5)

    def test_silent_object_and_height_channels_are_reported(self) -> None:
        input_block = np.zeros((8192, 10), dtype=np.float64)
        input_block[:, 1] = 0.5  # object 1 is active; object 2 stays digitally silent

        render = _fake_render([input_block], object_track_indices=(1, 2), bed_height_track_indices=(8, 9))

        with (
            patch(_ATMOS_RENDER, return_value=render),
            patch(_MEASURE_LOUDNORM, side_effect=_draining_loudnorm(_loudness())),
            patch(_DECODE_MONO, return_value=np.zeros(8000, dtype=np.float64)),
            patch(_XCORR, return_value=(0, 0.95)),
        ):
            render_check_measurements = measure_render_checks(MagicMock(), _STEREO_URL)

        assert render_check_measurements.silent_object.object_track_count == 2
        assert render_check_measurements.silent_object.silent_object_track_indices == (2,)
        assert render_check_measurements.silent_object.height_track_count == 2
        assert render_check_measurements.silent_object.silent_height_track_indices == (8, 9)

    def test_negative_only_channel_is_not_reported_silent(self) -> None:
        # The per-channel peak is max |sample|, so a channel whose loudest sample is negative (e.g. a
        # DC-offset or purely-negative-going object) reads its true magnitude and is NOT silent.
        # Guards the two-sided reduction: a max-only peak would read the negative peak as <= 0 and
        # fall to the silence floor, mislabeling a loud channel as blank.
        input_block = np.zeros((8192, 10), dtype=np.float64)
        input_block[:, 1] = -0.5  # object 1 is loud but negative-going; object 2 stays digitally silent

        render = _fake_render([input_block], object_track_indices=(1, 2))

        with (
            patch(_ATMOS_RENDER, return_value=render),
            patch(_MEASURE_LOUDNORM, side_effect=_draining_loudnorm(_loudness())),
            patch(_DECODE_MONO, return_value=np.zeros(8000, dtype=np.float64)),
            patch(_XCORR, return_value=(0, 0.95)),
        ):
            render_check_measurements = measure_render_checks(MagicMock(), _STEREO_URL)

        assert render_check_measurements.silent_object.silent_object_track_indices == (2,)

    def test_empty_master_raises_render_output_error(self) -> None:
        render = _fake_render([])  # input_blocks() yields nothing

        with (
            patch(_ATMOS_RENDER, return_value=render),
            patch(_MEASURE_LOUDNORM, side_effect=_draining_loudnorm(_loudness())),
            patch(_DECODE_MONO, return_value=np.zeros(80, dtype=np.float64)),
            patch(_XCORR, return_value=(0, 0.95)),
            pytest.raises(RenderOutputError, match="no audio blocks"),
        ):
            measure_render_checks(MagicMock(), _STEREO_URL)

    def test_non_finite_atmos_render_raises_render_output_error(self) -> None:
        def render_block(layout: str, samples: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
            channels = 6 if layout == "0+5+0" else 2
            block = np.ones((len(samples), channels), dtype=np.float64)
            if layout == "0+2+0":
                block[0, 0] = np.nan
            return block

        render = _fake_render(
            [np.ones((8000, 10), dtype=np.float64)],
            sample_rate_hz=8000,
            render_block=render_block,
        )

        with (
            patch(_ATMOS_RENDER, return_value=render),
            patch(_MEASURE_LOUDNORM, side_effect=_draining_loudnorm(_loudness())),
            patch(_DECODE_MONO, return_value=np.zeros(8000, dtype=np.float64)),
            patch(_XCORR, return_value=(0, 0.95)),
            pytest.raises(RenderOutputError, match="non-finite"),
        ):
            measure_render_checks(MagicMock(), _STEREO_URL)

    def test_real_render_produces_finite_measurements(self) -> None:
        # End-to-end through the real EAR parse + both renders + real loudnorm, cross-correlated
        # against a (mocked) silent stereo decode. test_bwf has an LFE (52 Hz tone, source channel
        # 3) and two active objects (channels 0-1), so the source-channel checks measure real
        # content: in-band LFE stays under both gates and neither object is silent.
        silence = np.zeros(60 * 8000, dtype=np.float32)
        decoded = MagicMock(returncode=0, stdout=silence.tobytes(), stderr=b"")
        with (
            _TEST_BWF.open("rb") as file_handle,
            patch("src.atmos.alignment_check.subprocess.run", return_value=decoded),
        ):
            render_check_measurements = measure_render_checks(file_handle, "https://example.com/stereo.flac")

        assert math.isfinite(render_check_measurements.loudness.true_peak_dbtp)
        assert not math.isnan(render_check_measurements.loudness.integrated_loudness_lkfs)
        assert isinstance(render_check_measurements.alignment, AlignmentMeasurement)
        assert math.isfinite(render_check_measurements.alignment.content_match_r)
        assert render_check_measurements.alignment.alignment_offset_ms is None or math.isfinite(
            render_check_measurements.alignment.alignment_offset_ms
        )
        assert render_check_measurements.lfe is not None
        assert render_check_measurements.lfe.above_400hz_dbfs < LFE_WARN_400_HZ_DBFS
        assert render_check_measurements.lfe.above_2000hz_dbfs < LFE_WARN_2000_HZ_DBFS
        assert render_check_measurements.silent_object.object_track_count == 2
        assert render_check_measurements.silent_object.silent_object_track_indices == ()
        # test_bwf's U-030 height format shares source channel 2, which carries an active tone.
        assert render_check_measurements.silent_object.height_track_count == 1
        assert render_check_measurements.silent_object.silent_height_track_indices == ()