import warnings
from collections import Counter
from collections.abc import Iterator, Sequence
from typing import Final, Literal, Protocol

import numpy as np
import numpy.typing as npt
from ear.core import Renderer, bs2051
from ear.core.direct_speakers.panner import DirectSpeakersPanner
from ear.core.metadata_input import DirectSpeakersRenderingItem, DirectTrackSpec, ObjectRenderingItem
from ear.core.metadata_processing import preprocess_rendering_items
from ear.core.renderer_common import is_lfe
from ear.core.select_items import select_rendering_items
from ear.fileio.adm import timing_fixes
from ear.fileio.adm.elements.geom import DirectSpeakerCartesianPosition, DirectSpeakerPolarPosition
from ear.fileio.bw64 import Bw64Reader
from ear.fileio.utils import Bw64AdmReader

from src.atmos.fix_dolby import fix_ds_frequency, fix_stream_pack_refs

RENDER_BLOCK_SIZE_FRAMES: Final = 8192

AtmosRenderLayout = Literal["0+2+0", "0+5+0"]

# The BS.2051 LFE speaker labels (LFE1/LFE2) plus the aliases EAR's DirectSpeakersPanner
# substitutes onto them (LFE/LFEL -> LFE1, LFER -> LFE2): every substitution lands inside the
# target set, so membership in this union equals the panner's substitute-then-match check.
# Frequency-based detection is the authoritative route on Dolby masters: their "RC_LFE" label
# matches no BS.2051 name, but fix_ds_frequency guarantees the LFE's lowPass element is present
# before classification runs.
_LFE_SPEAKER_LABELS: Final = frozenset({"LFE1", "LFE2", "LFE", "LFEL", "LFER"})


class MissingAdmMetadataError(Exception):
    pass


class _SeekableBinarySource(Protocol):
    def read(self, size: int = ..., /) -> bytes: ...

    def seek(self, offset: int, whence: int = ..., /) -> int: ...

    def tell(self) -> int: ...


class AtmosRender:
    # The canonical sequence from ear/cmdline/render_file.py, with the EAT
    # fix_dolby pre-fixes applied first — real Dolby ADM Profile masters fail
    # adm.validate() without them.
    #
    # Multi-layout by design: the ADM is parsed, fixed, and validated ONCE (the expensive step on a
    # real master), then each requested layout gets its own Renderer. Each layout's rendering items
    # are built with a fresh select+preprocess pass over that single parsed ADM — the items must NOT
    # be shared across renderers, because each item's metadata_source is a single-use iterator
    # (MetadataSourceIter): once one renderer's render() loop drains it, a second renderer sharing the
    # same item renders silence. Rebuilding the items per layout gives each Renderer independent
    # metadata sources while the heavy ADM parse stays shared.
    def __init__(self, file_like: _SeekableBinarySource, layout_names: Sequence[AtmosRenderLayout]) -> None:
        with warnings.catch_warnings(record=True) as prep_warnings:
            warnings.simplefilter("always")
            adm_reader = Bw64AdmReader(Bw64Reader(file_like))
            adm = adm_reader.adm
            if adm is None:
                raise MissingAdmMetadataError("file has no ADM metadata (missing chna chunk)")

            fix_stream_pack_refs(adm)
            fix_ds_frequency(adm)
            adm.validate()
            timing_fixes.check_blockFormat_timings(adm, fix=True)

            self._renderers: dict[AtmosRenderLayout, Renderer] = {}
            self._output_channel_counts: dict[AtmosRenderLayout, int] = {}
            for name in layout_names:
                layout = bs2051.get_layout(name)
                renderer = Renderer(layout)
                renderer.set_rendering_items(preprocess_rendering_items(select_rendering_items(adm)))
                self._renderers[name] = renderer
                self._output_channel_counts[name] = len(layout.channels)

            lfe_track_indices, object_track_indices, bed_height_track_indices = _classify_source_channels(adm)
            self.lfe_track_indices: tuple[int, ...] = lfe_track_indices
            self.object_track_indices: tuple[int, ...] = object_track_indices
            self.bed_height_track_indices: tuple[int, ...] = bed_height_track_indices

        # One real master emitted 202k individual prep warnings — aggregate to
        # bounded per-category counts and never surface them one by one.
        self.prep_warning_counts: dict[str, int] = dict(
            Counter(type(prep_warning.message).__name__ for prep_warning in prep_warnings)
        )
        self.sample_rate_hz: int = adm_reader.sampleRate
        self.input_channel_count: int = adm_reader.channels
        self._adm_reader = adm_reader

    def output_channel_count(self, layout_name: AtmosRenderLayout) -> int:
        return self._output_channel_counts[layout_name]

    def input_blocks(self) -> Iterator[npt.NDArray[np.float64]]:
        yield from self._adm_reader.iter_sample_blocks(RENDER_BLOCK_SIZE_FRAMES)

    def render_block(
        self, layout_name: AtmosRenderLayout, input_samples: npt.NDArray[np.float64]
    ) -> npt.NDArray[np.float64]:
        rendered: npt.NDArray[np.float64] = self._renderers[layout_name].render(self.sample_rate_hz, input_samples)
        return rendered

    def tail(self, layout_name: AtmosRenderLayout) -> npt.NDArray[np.float64]:
        tail_samples: npt.NDArray[np.float64] = self._renderers[layout_name].get_tail(
            self.sample_rate_hz, self.input_channel_count
        )
        return tail_samples


def _classify_source_channels(adm: object) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
    # A dedicated selection pass: the per-layout rendering items are consumed by their renderers
    # (each item's metadata_source is single-use), so classification reads its own fresh items —
    # and only through plain attributes (adm_path, track_spec), never the metadata_source. Items
    # without a DirectTrackSpec (silent or matrix-derived tracks) map to no source channel, so
    # there is nothing to measure for them.
    lfe_track_indices: set[int] = set()
    object_track_indices: set[int] = set()
    bed_height_track_indices: set[int] = set()
    for rendering_item in select_rendering_items(adm):
        if isinstance(rendering_item, ObjectRenderingItem) and isinstance(rendering_item.track_spec, DirectTrackSpec):
            object_track_indices.add(rendering_item.track_spec.track_index)
        elif isinstance(rendering_item, DirectSpeakersRenderingItem) and isinstance(
            rendering_item.track_spec, DirectTrackSpec
        ):
            if _is_lfe_item(rendering_item):
                lfe_track_indices.add(rendering_item.track_spec.track_index)
            elif _is_height_item(rendering_item):
                bed_height_track_indices.add(rendering_item.track_spec.track_index)

    return (
        tuple(sorted(lfe_track_indices)),
        tuple(sorted(object_track_indices)),
        tuple(sorted(bed_height_track_indices)),
    )


def _is_lfe_item(rendering_item: DirectSpeakersRenderingItem) -> bool:
    # Mirrors EAR's DirectSpeakersPanner.is_lfe_channel (LFE by frequency element OR by speaker
    # label) without instantiating a panner, which needs a full layout.
    audio_channel_format = rendering_item.adm_path.audioChannelFormat
    if is_lfe(audio_channel_format.frequency):
        return True

    for block_format in audio_channel_format.audioBlockFormats:
        for speaker_label in block_format.speakerLabel:
            urn_match = DirectSpeakersPanner.SPEAKER_URN_REGEX.match(speaker_label)
            nominal_label = urn_match.group(1) if urn_match else speaker_label
            if nominal_label in _LFE_SPEAKER_LABELS:
                return True

    return False


def _is_height_item(rendering_item: DirectSpeakersRenderingItem) -> bool:
    # Height detection is position-based rather than label-based: real Dolby masters label their top
    # surrounds RC_Lts/RC_Rts (no BS.2051 match), but every DirectSpeakers block format carries a
    # position — polar elevation or Cartesian Z, either above zero means above the listener.
    for block_format in rendering_item.adm_path.audioChannelFormat.audioBlockFormats:
        position = block_format.position
        if isinstance(position, DirectSpeakerPolarPosition) and position.elevation > 0:
            return True
        if isinstance(position, DirectSpeakerCartesianPosition) and position.bounded_Z.value > 0:
            return True

    return False