# Single-pass combined render pipeline: every render-based check shares one ADM parse (the
# expensive step on a real master) and one pass over the input audio. The loudness (0+5+0) and
# alignment (0+2+0) checks consume the renders; the LFE and silent-object checks consume the raw
# source channels of the same blocks. Sharing the parse holds the peak near a single render
# instead of the ~2x of two separate AtmosRender instances (which would each re-parse the master).
from collections.abc import Iterator
from typing import BinaryIO, Final

import numpy as np
import numpy.typing as npt
from pydantic import BaseModel, ConfigDict

from src.atmos.alignment_check import (
    _ANALYSIS_DURATION_S,
    _ANALYSIS_SAMPLES,
    AlignmentMeasurement,
    _cross_correlate,
    _decode_head_mono_8k,
    _resample_to_8k,
    _to_alignment_measurement,
)
from src.atmos.lfe_check import LfeFftAccumulator, LfeMeasurement
from src.atmos.loudness_check import LoudnessMeasurement, _measure_loudnorm
from src.atmos.render import AtmosRender, AtmosRenderLayout
from src.atmos.render_errors import RenderOutputError
from src.atmos.silent_object_check import SilentObjectMeasurement, to_silent_object_measurement
from src.clients import s3
from src.config import presigned_url_ttl_seconds

# Loudness/true peak are measured on a 5.1 (BS.2051 0+5+0) render per Apple/Netflix/Dolby
# convention; alignment cross-correlates a stereo (0+2+0) render against the stereo reference.
_LOUDNESS_LAYOUT: Final[AtmosRenderLayout] = "0+5+0"
_ALIGNMENT_LAYOUT: Final[AtmosRenderLayout] = "0+2+0"
# The EAR render carries no layout tag (channel_layout=unknown), so the layout is forced when handing
# PCM to ffmpeg — without it ffmpeg can't identify LFE/surrounds and BS.1770 weighting/exclusion
# would be wrong. EAR's 0+5+0 order (L R C LFE Ls Rs — the BS.2051 channel order) lines up with
# ffmpeg's 5.1 (FL FR FC LFE BL BR) position-for-position, so the tag only re-labels and never
# reorders samples: LFE stays at position 4 (excluded from loudness) and the surrounds at 5-6.
_FFMPEG_CHANNEL_LAYOUT: Final = "5.1"


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

    loudness: LoudnessMeasurement
    alignment: AlignmentMeasurement
    # None when the ADM has no LFE channel — nothing to measure (Sony: no-bed files pass-skip).
    lfe: LfeMeasurement | None
    silent_object: SilentObjectMeasurement


def run_render_checks(
    atmos_bucket: str,
    atmos_key: str,
    stereo_reference_bucket: str,
    stereo_reference_key: str,
) -> RenderCheckMeasurements:
    # Top-level and picklable so a spawn-based ProcessPoolExecutor can dispatch it to a worker
    # process. The stereo reference is fetched by presigned URL (ffmpeg decodes it directly); the
    # atmos is opened seekable for the renderer. Both S3 accesses happen in the worker (boto3 clients
    # aren't fork-safe). One worker produces every render-based measurement from a single ADM parse.
    stereo_reference_url = str(
        s3.create_presigned_url(
            stereo_reference_bucket, stereo_reference_key, expires_in_seconds=presigned_url_ttl_seconds()
        )
    )
    with s3.open_s3_seekable(atmos_bucket, atmos_key) as atmos_handle:
        return measure_render_checks(atmos_handle, stereo_reference_url)


def measure_render_checks(
    atmos_file_like: BinaryIO,
    stereo_reference_url: str,
) -> RenderCheckMeasurements:
    render = AtmosRender(atmos_file_like, [_LOUDNESS_LAYOUT, _ALIGNMENT_LAYOUT])
    alignment_head_frames = _ANALYSIS_DURATION_S * render.sample_rate_hz
    atmos_mono_chunks: list[npt.NDArray[np.float64]] = []
    lfe_fft_accumulators = {
        track_index: LfeFftAccumulator(render.sample_rate_hz) for track_index in render.lfe_track_indices
    }
    channel_peaks = np.zeros(render.input_channel_count)

    def loudness_blocks() -> Iterator[npt.NDArray[np.float64]]:
        # Single pass over the input audio: the block stream is single-use (the ADM block iterator
        # exhausts after one pass). A second pass is possible in principle — Bw64Reader.seek() could
        # rewind the handle — but that means re-reading the multi-GB master from S3 and re-running the
        # ~24 s ADM preprocess, so every consumer feeds from the same stream here instead. The 5.1
        # blocks stream straight to loudnorm; the source-channel scans (per-channel peaks, LFE FFT)
        # take every block of the whole track; only the stereo-render capture stops at the 60 s head,
        # buffered for the cross-correlation that runs after the pass. This generator is driven by
        # loudnorm's feeder thread, so the captures are side effects of that thread — they are fully
        # populated only once _measure_loudnorm returns (the feeder joins).
        collected = 0
        for input_samples in render.input_blocks():
            # max(|x|) == max(max(x), -min(x)): two channel-sized reductions instead of a full
            # (block_size, channels) np.abs temp allocated for every block of a multi-GB track.
            np.maximum(channel_peaks, input_samples.max(axis=0), out=channel_peaks)
            np.maximum(channel_peaks, -input_samples.min(axis=0), out=channel_peaks)
            for track_index, lfe_fft_accumulator in lfe_fft_accumulators.items():
                lfe_fft_accumulator.add_block(input_samples[:, track_index])
            # The boundary block that crosses 60 s is captured whole (pushing slightly past
            # alignment_head_frames); the excess is trimmed by the [:alignment_head_frames] slice below.
            if collected < alignment_head_frames:
                stereo_block = render.render_block(_ALIGNMENT_LAYOUT, input_samples)
                atmos_mono_chunks.append(stereo_block.mean(axis=1))
                collected += len(input_samples)
            yield render.render_block(_LOUDNESS_LAYOUT, input_samples)
        yield render.tail(_LOUDNESS_LAYOUT)

    loudness_measurement = _measure_loudnorm(
        loudness_blocks(),
        sample_rate_hz=render.sample_rate_hz,
        channel_count=render.output_channel_count(_LOUDNESS_LAYOUT),
        ffmpeg_channel_layout=_FFMPEG_CHANNEL_LAYOUT,
    )

    # Both cross-correlation inputs must be finite and non-empty — a NaN or a zero-length side silently
    # defeats the comparison gates in _to_alignment_measurement and _check_*. The stereo side guards
    # both in _decode_head_mono_8k; these two atmos-side guards mirror it (empty master, then non-finite
    # samples).
    if not atmos_mono_chunks:
        raise RenderOutputError("atmos master produced no audio blocks (empty or zero-length master)")

    atmos_mono = _resample_to_8k(np.concatenate(atmos_mono_chunks)[:alignment_head_frames], render.sample_rate_hz)[
        :_ANALYSIS_SAMPLES
    ]
    if not np.all(np.isfinite(atmos_mono)):
        raise RenderOutputError("atmos stereo render produced non-finite samples")

    stereo_mono = _decode_head_mono_8k(stereo_reference_url)
    lag_samples, content_match_r = _cross_correlate(stereo_mono, atmos_mono)
    alignment_measurement = _to_alignment_measurement(lag_samples, content_match_r)

    lfe_measurement = None
    if lfe_fft_accumulators:
        # Worst maxima across LFE channels, matching Sony's worst-across-beds roll-up.
        lfe_channel_measurements = tuple(accumulator.finish() for accumulator in lfe_fft_accumulators.values())
        lfe_measurement = LfeMeasurement(
            above_400hz_dbfs=max(m.above_400hz_dbfs for m in lfe_channel_measurements),
            above_2000hz_dbfs=max(m.above_2000hz_dbfs for m in lfe_channel_measurements),
        )

    silent_object_measurement = to_silent_object_measurement(
        channel_peaks, render.object_track_indices, render.bed_height_track_indices
    )

    return RenderCheckMeasurements(
        loudness=loudness_measurement,
        alignment=alignment_measurement,
        lfe=lfe_measurement,
        silent_object=silent_object_measurement,
    )