# LFE full-frequency check, ported verbatim from Sony QC's fft_maxima_stream (ventura_1_1.py) for
# parity with their operational gate: the LFE channel is resampled to 48 kHz and streamed through a
# Hann-windowed 4096-point FFT at 50% hop; the running maximum bin magnitude at/above each cutoff
# is compared against that cutoff's own limit, and each exceedance warns under its own finding (the
# 400 Hz and 2000 Hz gates are independent — see AtmosValidationFindingCode). Magnitudes are
# normalized by 2/sum(window), so a full-scale sine on an exact bin reads 0 dBFS. The port keeps
# Sony's float32 arithmetic, per-chunk resampling, 1e-12 log guard, and zero-padded final block so
# results are bit-identical on the same samples.
import math
from typing import Final

import numpy as np
import numpy.typing as npt
from pydantic import BaseModel, ConfigDict
from scipy.fft import rfft, rfftfreq
from scipy.signal import resample_poly

from src.atmos.models import (
    SILENCE_FLOOR_DBFS,
    AtmosValidationBuilder,
    AtmosValidationFindingCode,
)

LFE_WARN_400_HZ_DBFS: Final = -60.0
LFE_WARN_2000_HZ_DBFS: Final = -100.0
_FFT_BLOCK_SIZE: Final = 4096
_FFT_HOP_SIZE: Final = _FFT_BLOCK_SIZE // 2
_FFT_SAMPLE_RATE_HZ: Final = 48_000

_WINDOW: Final[npt.NDArray[np.float32]] = np.hanning(_FFT_BLOCK_SIZE).astype(np.float32)
_WINDOW_SUM: Final = float(np.sum(_WINDOW))
_BIN_FREQUENCIES_HZ: Final[npt.NDArray[np.float64]] = rfftfreq(_FFT_BLOCK_SIZE, 1 / _FFT_SAMPLE_RATE_HZ)
_FIRST_BIN_400_HZ: Final = int(np.searchsorted(_BIN_FREQUENCIES_HZ, 400, side="left"))
_FIRST_BIN_2000_HZ: Final = int(np.searchsorted(_BIN_FREQUENCIES_HZ, 2000, side="left"))


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

    above_400hz_dbfs: float
    above_2000hz_dbfs: float


class LfeFftAccumulator:
    # One per LFE channel: consumes that channel's samples block-by-block during the single render
    # pass, holding only the sub-FFT-block remainder between calls.
    def __init__(self, source_sample_rate_hz: int) -> None:
        self._source_sample_rate_hz = source_sample_rate_hz
        self._buffer: npt.NDArray[np.float32] = np.empty(0, dtype=np.float32)
        self._max_400_hz_dbfs = -math.inf
        self._max_2000_hz_dbfs = -math.inf

    def add_block(self, mono_samples: npt.NDArray[np.float64]) -> None:
        chunk = mono_samples.astype(np.float32)
        if self._source_sample_rate_hz != _FFT_SAMPLE_RATE_HZ:
            # Sony resamples each read chunk independently (their reader chunks frames the same way
            # our render blocks do), so per-chunk polyphase edges are part of the parity contract.
            divisor = math.gcd(self._source_sample_rate_hz, _FFT_SAMPLE_RATE_HZ)
            chunk = resample_poly(chunk, _FFT_SAMPLE_RATE_HZ // divisor, self._source_sample_rate_hz // divisor).astype(
                np.float32
            )

        self._buffer = chunk if self._buffer.size == 0 else np.concatenate((self._buffer, chunk))
        while self._buffer.size >= _FFT_BLOCK_SIZE:
            self._measure_fft_block(self._buffer[:_FFT_BLOCK_SIZE])
            self._buffer = self._buffer[_FFT_HOP_SIZE:]

    def finish(self) -> LfeMeasurement:
        if self._buffer.size > 0:
            padding = np.zeros(_FFT_BLOCK_SIZE - self._buffer.size, dtype=np.float32)
            self._measure_fft_block(np.concatenate((self._buffer, padding)))
            self._buffer = np.empty(0, dtype=np.float32)

        return LfeMeasurement(
            above_400hz_dbfs=max(self._max_400_hz_dbfs, SILENCE_FLOOR_DBFS),
            above_2000hz_dbfs=max(self._max_2000_hz_dbfs, SILENCE_FLOOR_DBFS),
        )

    def _measure_fft_block(self, fft_block: npt.NDArray[np.float32]) -> None:
        magnitudes = (np.abs(rfft(fft_block * _WINDOW)) * 2.0) / _WINDOW_SUM

        peak_400 = float(np.max(magnitudes[_FIRST_BIN_400_HZ:]))
        if peak_400 > 1e-12:
            self._max_400_hz_dbfs = max(self._max_400_hz_dbfs, 20.0 * math.log10(peak_400))

        peak_2000 = float(np.max(magnitudes[_FIRST_BIN_2000_HZ:]))
        if peak_2000 > 1e-12:
            self._max_2000_hz_dbfs = max(self._max_2000_hz_dbfs, 20.0 * math.log10(peak_2000))


def apply_lfe_measurement(
    lfe_measurement: LfeMeasurement | None,
    atmos_validation_builder: AtmosValidationBuilder,
) -> None:
    # None means the ADM has no LFE channel — nothing to measure (Sony: no-bed files pass-skip).
    if lfe_measurement is None:
        return

    atmos_validation_builder.update_metadata(
        lfe_above_400hz_dbfs=lfe_measurement.above_400hz_dbfs,
        lfe_above_2000hz_dbfs=lfe_measurement.above_2000hz_dbfs,
    )

    _check_lfe_above_400hz(lfe_measurement, atmos_validation_builder)
    _check_lfe_above_2000hz(lfe_measurement, atmos_validation_builder)


def _check_lfe_above_400hz(lfe_measurement: LfeMeasurement, atmos_validation_builder: AtmosValidationBuilder) -> None:
    if lfe_measurement.above_400hz_dbfs > LFE_WARN_400_HZ_DBFS:
        atmos_validation_builder.warning(
            AtmosValidationFindingCode.LFE_LEVEL_ABOVE_400HZ_TOO_HIGH,
            f"LFE peak level at/above 400 Hz is {lfe_measurement.above_400hz_dbfs:.1f} dBFS, "
            f"should not exceed {LFE_WARN_400_HZ_DBFS} dBFS",
        )


def _check_lfe_above_2000hz(lfe_measurement: LfeMeasurement, atmos_validation_builder: AtmosValidationBuilder) -> None:
    if lfe_measurement.above_2000hz_dbfs > LFE_WARN_2000_HZ_DBFS:
        atmos_validation_builder.warning(
            AtmosValidationFindingCode.LFE_LEVEL_ABOVE_2000HZ_TOO_HIGH,
            f"LFE peak level at/above 2000 Hz is {lfe_measurement.above_2000hz_dbfs:.1f} dBFS, "
            f"should not exceed {LFE_WARN_2000_HZ_DBFS} dBFS",
        )