import warnings
from collections.abc import Callable
from pathlib import Path
from typing import BinaryIO
from unittest.mock import patch

import numpy as np
import numpy.typing as npt
import pytest
from src.atmos.render import (
    _LFE_SPEAKER_LABELS,
    RENDER_BLOCK_SIZE_FRAMES,
    AtmosRender,
    AtmosRenderLayout,
    MissingAdmMetadataError,
)

_FIXTURES_DIR = Path(__file__).parent.parent.parent / "fixtures"
_TEST_BWF = _FIXTURES_DIR / "test_bwf.wav"
_SILENT_STEREO_WAV = _FIXTURES_DIR / "silent_stereo_2s.wav"
# A real Dolby ADM Profile master that fails BS.2127 validation until the
# fix_dolby pre-fixes run (see test_fix_dolby.py); the end-to-end test below
# proves AtmosRender repairs and renders it.
_DOLBY_RENDERER_WAV = _FIXTURES_DIR / "dolby_renderer_silent_2s.wav"


def _concatenated_render(atmos_render: AtmosRender, layout: AtmosRenderLayout) -> npt.NDArray[np.float64]:
    blocks = [atmos_render.render_block(layout, input_samples) for input_samples in atmos_render.input_blocks()]
    blocks.append(atmos_render.tail(layout))
    return np.concatenate(blocks)


def _render_single_layout(file_like: BinaryIO, layout: AtmosRenderLayout) -> npt.NDArray[np.float64]:
    return _concatenated_render(AtmosRender(file_like, [layout]), layout)


def test_render_through_seekable_handle_produces_signal() -> None:
    with _TEST_BWF.open("rb") as file_handle:
        render_samples = _render_single_layout(file_handle, "0+2+0")

    assert render_samples.shape[1] == 2
    assert np.abs(render_samples).max() > 0


@pytest.mark.parametrize(
    ("layout", "expected_output_channel_count"),
    [
        pytest.param("0+2+0", 2, id="stereo"),
        pytest.param("0+5+0", 6, id="five_one"),
    ],
)
def test_blocks_cover_all_input_frames_including_tail_flush(
    layout: AtmosRenderLayout, expected_output_channel_count: int
) -> None:
    from ear.fileio.bw64 import Bw64Reader

    with _TEST_BWF.open("rb") as file_handle:
        bw64_reader = Bw64Reader(file_handle)
        input_frame_count = bw64_reader._chunks[b"data"].size // bw64_reader.formatInfo.blockAlignment

    with _TEST_BWF.open("rb") as file_handle:
        atmos_render = AtmosRender(file_handle, [layout])
        rendered_blocks = [
            atmos_render.render_block(layout, input_samples) for input_samples in atmos_render.input_blocks()
        ]
        rendered_blocks.append(atmos_render.tail(layout))

    assert atmos_render.sample_rate_hz == 48000
    assert atmos_render.output_channel_count(layout) == expected_output_channel_count
    assert all(block.shape[1] == expected_output_channel_count for block in rendered_blocks)
    assert all(len(block) <= RENDER_BLOCK_SIZE_FRAMES for block in rendered_blocks)
    assert sum(len(block) for block in rendered_blocks) == input_frame_count


def test_multi_layout_shares_parse_and_renders_both_layouts() -> None:
    # Both layouts share one ADM parse but get FRESH per-layout rendering items (sharing the same
    # items renders silence — each item's metadata_source is a single-use iterator). The max>0
    # assertions prove both layouts render non-silent content: the guard against that silence bug.
    with _TEST_BWF.open("rb") as file_handle:
        atmos_render = AtmosRender(file_handle, ["0+5+0", "0+2+0"])
        assert atmos_render.output_channel_count("0+5+0") == 6
        assert atmos_render.output_channel_count("0+2+0") == 2

        five_one_max = 0.0
        stereo_max = 0.0
        for input_samples in atmos_render.input_blocks():
            five_one_block = atmos_render.render_block("0+5+0", input_samples)
            stereo_block = atmos_render.render_block("0+2+0", input_samples)
            assert five_one_block.shape[1] == 6
            assert stereo_block.shape[1] == 2
            assert np.all(np.isfinite(five_one_block))
            assert np.all(np.isfinite(stereo_block))
            five_one_max = max(five_one_max, float(np.abs(five_one_block).max()))
            stereo_max = max(stereo_max, float(np.abs(stereo_block).max()))

    assert five_one_max > 0
    assert stereo_max > 0


def test_shared_multi_layout_render_matches_single_layout_render() -> None:
    # Rendering 0+5+0 from the shared ADM parse (with fresh per-layout items, alongside 0+2+0) must be
    # identical to rendering 0+5+0 from a single-layout parse — one shared parse doesn't perturb output.
    with _TEST_BWF.open("rb") as file_handle:
        single_layout_render = _render_single_layout(file_handle, "0+5+0")

    with _TEST_BWF.open("rb") as file_handle:
        shared_render = _concatenated_render(AtmosRender(file_handle, ["0+5+0", "0+2+0"]), "0+5+0")

    np.testing.assert_array_equal(single_layout_render, shared_render)


def test_prep_applies_dolby_pre_fixes_to_parsed_adm() -> None:
    with (
        _TEST_BWF.open("rb") as file_handle,
        patch("src.atmos.render.fix_stream_pack_refs", return_value=()) as fix_stream_pack_refs_mock,
        patch("src.atmos.render.fix_ds_frequency", return_value=()) as fix_ds_frequency_mock,
    ):
        AtmosRender(file_handle, ["0+2+0"])

    fix_stream_pack_refs_mock.assert_called_once()
    fix_ds_frequency_mock.assert_called_once()


def test_prep_warnings_are_aggregated_into_counts_not_emitted() -> None:
    def warn_three_times(*_: object, **__: object) -> None:
        warnings.warn("fixed blockFormat 1", UserWarning, stacklevel=1)
        warnings.warn("fixed blockFormat 2", UserWarning, stacklevel=1)
        warnings.warn("clock drift", RuntimeWarning, stacklevel=1)

    with (
        _TEST_BWF.open("rb") as file_handle,
        patch("src.atmos.render.timing_fixes.check_blockFormat_timings", side_effect=warn_three_times),
        warnings.catch_warnings(record=True) as escaped_warnings,
    ):
        warnings.simplefilter("always")
        atmos_render = AtmosRender(file_handle, ["0+2+0"])

    assert atmos_render.prep_warning_counts == {"UserWarning": 2, "RuntimeWarning": 1}
    assert escaped_warnings == []


def test_render_of_wav_without_adm_metadata_raises_clear_error() -> None:
    with _SILENT_STEREO_WAV.open("rb") as plain_wav_handle, pytest.raises(MissingAdmMetadataError, match="chna"):
        AtmosRender(plain_wav_handle, ["0+2+0"])


@pytest.mark.parametrize(
    ("layout", "expected_output_channel_count"),
    [
        pytest.param("0+2+0", 2, id="stereo"),
        pytest.param("0+5+0", 6, id="five_one"),
    ],
)
def test_renders_real_dolby_master_after_pre_fixes(
    layout: AtmosRenderLayout, expected_output_channel_count: int
) -> None:
    with _DOLBY_RENDERER_WAV.open("rb") as file_handle:
        rendered = _render_single_layout(file_handle, layout)

    assert rendered.shape[0] > 0
    assert rendered.shape[1] == expected_output_channel_count


def test_classifies_source_channels_of_the_real_dolby_master() -> None:
    # One 7.1.2 bed, no objects; the LFE (RoomCentricLFE, chna track 4) is source channel 3.
    # Its speakerLabel is "RC_LFE" — not a BS.2051 LFE label — and the file carries no frequency
    # element, so classification rests on the lowPass that fix_ds_frequency stamps during prep.
    with _DOLBY_RENDERER_WAV.open("rb") as file_handle:
        render = AtmosRender(file_handle, ["0+2+0"])

    assert render.lfe_track_indices == (3,)
    assert render.object_track_indices == ()
    # The 7.1.2 bed's top-surround pair carries Cartesian positions with Z=1.0.
    assert render.bed_height_track_indices == (8, 9)


def test_classifies_source_channels_of_the_synthetic_bwf() -> None:
    # Two mono Objects on source channels 0-1; LFE1 (explicit lowPass 120) on source channel 3.
    # Channel 2 carries two non-LFE DirectSpeakers formats and must land in neither bucket.
    with _TEST_BWF.open("rb") as file_handle:
        render = AtmosRender(file_handle, ["0+5+0"])

    assert render.lfe_track_indices == (3,)
    assert render.object_track_indices == (0, 1)
    # Channel 2 carries a U-030 (upper-layer) DirectSpeakers format alongside a non-height M+045.
    assert render.bed_height_track_indices == (2,)


def _direct_speakers_item(track_index: int | None, speaker_label: str, elevation_degrees: float = 0.0) -> object:
    from ear.core.metadata_input import (
        ADMPath,
        DirectSpeakersRenderingItem,
        DirectTrackSpec,
        MetadataSourceIter,
        SilentTrackSpec,
    )
    from ear.fileio.adm.elements import AudioBlockFormatDirectSpeakers, AudioChannelFormat, TypeDefinition
    from ear.fileio.adm.elements.geom import BoundCoordinate, DirectSpeakerPolarPosition

    block_format = AudioBlockFormatDirectSpeakers(
        position=DirectSpeakerPolarPosition(
            bounded_azimuth=BoundCoordinate(0.0),
            bounded_elevation=BoundCoordinate(elevation_degrees),
            bounded_distance=BoundCoordinate(1.0),
        ),
        speakerLabel=[speaker_label],
    )
    audio_channel_format = AudioChannelFormat(
        audioChannelFormatName="test channel",
        type=TypeDefinition.DirectSpeakers,
        audioBlockFormats=[block_format],
    )
    return DirectSpeakersRenderingItem(
        track_spec=SilentTrackSpec() if track_index is None else DirectTrackSpec(track_index),
        metadata_source=MetadataSourceIter([]),
        adm_path=ADMPath(audioChannelFormat=audio_channel_format),
    )


def _object_item(track_index: int | None) -> object:
    from ear.core.metadata_input import DirectTrackSpec, MetadataSourceIter, ObjectRenderingItem, SilentTrackSpec

    return ObjectRenderingItem(
        track_spec=SilentTrackSpec() if track_index is None else DirectTrackSpec(track_index),
        metadata_source=MetadataSourceIter([]),
    )


def _classify(*rendering_items: object) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
    from src.atmos.render import _classify_source_channels

    with patch("src.atmos.render.select_rendering_items", return_value=list(rendering_items)):
        return _classify_source_channels(object())


@pytest.mark.parametrize(
    ("speaker_label", "expected_lfe"),
    [
        pytest.param("LFE1", True, id="LFE1"),
        pytest.param("LFE2", True, id="LFE2"),
        pytest.param("LFE", True, id="bare_LFE_alias"),
        pytest.param("LFEL", True, id="LFEL_alias"),
        pytest.param("LFER", True, id="LFER_alias"),
        pytest.param("urn:itu:bs:2051:0:speaker:LFE1", True, id="urn_LFE1"),
        pytest.param("urn:itu:bs:2051:1:speaker:LFEL", True, id="urn_LFEL_alias"),
        pytest.param("M+000", False, id="non_lfe_label"),
        pytest.param("urn:itu:bs:2051:0:speaker:M+000", False, id="urn_non_lfe_label"),
        pytest.param("RC_LFE", False, id="dolby_label_matches_only_via_frequency"),
    ],
)
def test_classifies_lfe_by_speaker_label_without_a_frequency_element(speaker_label: str, expected_lfe: bool) -> None:
    # The label route covers masters whose LFE carries a BS.2051 speakerLabel (bare, aliased, or
    # URN-wrapped) but no frequency element that fix_ds_frequency would have stamped.
    lfe_track_indices, _, _ = _classify(_direct_speakers_item(3, speaker_label))

    assert (lfe_track_indices == (3,)) is expected_lfe


def test_lfe_speaker_labels_stay_in_sync_with_ears_substitution_table() -> None:
    # _LFE_SPEAKER_LABELS hardcodes the labels EAR's DirectSpeakersPanner treats as LFE via its
    # substitute-then-match route (nominal_speaker_label -> substitutions -> "LFE1"/"LFE2"), so the
    # label detection in _is_lfe_item can skip instantiating a panner. Derive that set straight from
    # the (vendored) panner and pin the two together: a re-vendor that changes EAR's aliases fails
    # here instead of silently dropping newly-aliased LFE channels from classification.
    from ear.core import bs2051
    from ear.core.direct_speakers.panner import DirectSpeakersPanner

    substitutions = DirectSpeakersPanner(bs2051.get_layout("0+5+0")).substitutions
    ear_lfe_labels = {"LFE1", "LFE2"} | {alias for alias, target in substitutions.items() if target in {"LFE1", "LFE2"}}

    assert ear_lfe_labels == _LFE_SPEAKER_LABELS


def test_channels_without_a_source_track_are_not_classified() -> None:
    # SilentTrackSpec items reference no WAV channel, so there is nothing for either check to
    # measure — even an LFE-labelled one.
    lfe_track_indices, object_track_indices, bed_height_track_indices = _classify(
        _direct_speakers_item(None, "LFE1"), _object_item(None)
    )

    assert lfe_track_indices == ()
    assert object_track_indices == ()
    assert bed_height_track_indices == ()


def _cartesian_direct_speakers_item(track_index: int, z: float) -> object:
    from ear.core.metadata_input import ADMPath, DirectSpeakersRenderingItem, DirectTrackSpec, MetadataSourceIter
    from ear.fileio.adm.elements import AudioBlockFormatDirectSpeakers, AudioChannelFormat, TypeDefinition
    from ear.fileio.adm.elements.geom import BoundCoordinate, DirectSpeakerCartesianPosition

    block_format = AudioBlockFormatDirectSpeakers(
        position=DirectSpeakerCartesianPosition(
            bounded_X=BoundCoordinate(0.0),
            bounded_Y=BoundCoordinate(1.0),
            bounded_Z=BoundCoordinate(z),
        ),
        speakerLabel=["RC_Lts"],
    )
    audio_channel_format = AudioChannelFormat(
        audioChannelFormatName="test channel",
        type=TypeDefinition.DirectSpeakers,
        audioBlockFormats=[block_format],
    )
    return DirectSpeakersRenderingItem(
        track_spec=DirectTrackSpec(track_index),
        metadata_source=MetadataSourceIter([]),
        adm_path=ADMPath(audioChannelFormat=audio_channel_format),
    )


@pytest.mark.parametrize(
    ("rendering_item", "expected_heights"),
    [
        pytest.param(lambda: _direct_speakers_item(5, "U+030", elevation_degrees=30.0), (5,), id="polar_elevated"),
        pytest.param(lambda: _direct_speakers_item(5, "M+030", elevation_degrees=0.0), (), id="polar_ear_level"),
        pytest.param(lambda: _cartesian_direct_speakers_item(5, 1.0), (5,), id="cartesian_top"),
        pytest.param(lambda: _cartesian_direct_speakers_item(5, 0.0), (), id="cartesian_floor"),
        pytest.param(lambda: _direct_speakers_item(5, "LFE1", elevation_degrees=30.0), (), id="lfe_never_height"),
    ],
)
def test_classifies_bed_height_channels_by_position(
    rendering_item: Callable[[], object], expected_heights: tuple[int, ...]
) -> None:
    # Height detection is position-based (polar elevation or Cartesian Z above zero) because real
    # Dolby masters label their top surrounds RC_Lts/RC_Rts, which match no BS.2051 name.
    _, _, bed_height_track_indices = _classify(rendering_item())

    assert bed_height_track_indices == expected_heights


def test_parse_warnings_are_aggregated_into_counts_not_emitted() -> None:
    from ear.fileio.utils import Bw64AdmReader

    def parse_with_warning(bw64_reader: object) -> Bw64AdmReader:
        warnings.warn("added missing rtime to AB_00031001_00000001", UserWarning, stacklevel=1)
        return Bw64AdmReader(bw64_reader)

    with (
        _TEST_BWF.open("rb") as file_handle,
        patch("src.atmos.render.Bw64AdmReader", side_effect=parse_with_warning),
        warnings.catch_warnings(record=True) as escaped_warnings,
    ):
        warnings.simplefilter("always")
        atmos_render = AtmosRender(file_handle, ["0+2+0"])

    assert atmos_render.prep_warning_counts == {"UserWarning": 1}
    assert escaped_warnings == []