import json
import subprocess
import threading
from collections.abc import Iterable
from typing import Final

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

from src.atmos.models import AtmosValidationBuilder, AtmosValidationFindingCode
from src.atmos.render_errors import FfmpegError, FfmpegUnavailable, LoudnormOutputError

# Apple's published Atmos limits, measured per ITU-R BS.1770-4 (the verbatim "should not exceed"
# quotes are on the AtmosValidationFindingCode members). "Should not exceed" → warn strictly above.
LOUDNESS_WARN_LKFS: Final = -18.0
TRUE_PEAK_WARN_DBTP: Final = -1.0


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

    integrated_loudness_lkfs: float
    true_peak_dbtp: float


# Render-agnostic: takes any iterable of PCM blocks plus the channel layout, so the measurement is
# decoupled from how the blocks are produced — any PCM block source can stream into its own
# consumer. The layout must be declared because the EAR render is channel_layout=unknown; without it
# ffmpeg can't identify the LFE/surrounds and BS.1770 weighting/exclusion would be wrong.
def _measure_loudnorm(
    blocks: Iterable[npt.NDArray[np.float64]],
    *,
    sample_rate_hz: int,
    channel_count: int,
    ffmpeg_channel_layout: str,
) -> LoudnessMeasurement:
    command = [
        "ffmpeg",
        "-hide_banner",
        "-nostats",
        "-f",
        "f32le",
        "-ar",
        str(sample_rate_hz),
        "-ac",
        str(channel_count),
        "-i",
        "pipe:0",
        "-af",
        f"aformat=channel_layouts={ffmpeg_channel_layout},loudnorm=print_format=json",
        "-f",
        "null",
        "-",
    ]
    try:
        process = subprocess.Popen(command, stdin=subprocess.PIPE, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
    except OSError as e:
        raise FfmpegUnavailable(f"could not execute ffmpeg: {e}") from e

    # Stream blocks to stdin on a thread while the main thread drains stderr: loudnorm emits its
    # JSON only after end-of-input, and writing everything before reading stderr would deadlock if
    # ffmpeg's stderr pipe filled mid-stream.
    ffmpeg_stdin = process.stdin
    ffmpeg_stderr = process.stderr
    assert ffmpeg_stdin is not None  # guaranteed by stdin=PIPE; narrows for the closure below
    assert ffmpeg_stderr is not None
    feed_error: list[Exception] = []

    def _feed() -> None:
        try:
            for block in blocks:
                ffmpeg_stdin.write(np.ascontiguousarray(block, dtype=np.float32).tobytes())
        except Exception as e:  # re-raised on the main thread after join
            feed_error.append(e)
        finally:
            ffmpeg_stdin.close()

    feeder = threading.Thread(target=_feed)
    feeder.start()
    try:
        stderr_output = ffmpeg_stderr.read().decode("utf-8", errors="replace")
        return_code = process.wait()
    finally:
        # If an error escaped before ffmpeg finished, it may still be running with the feeder still
        # writing — kill it FIRST so the broken stdin pipe unblocks the feeder's next write, THEN
        # join. Joining first would wait for the feeder to stream the whole render (a long hang on a
        # big master). On the success path ffmpeg has already exited, so the kill is skipped.
        if process.poll() is None:
            process.kill()
            process.wait()
        feeder.join()

    if feed_error:
        # The feeder both drives the block generator (which renders) and writes to ffmpeg's stdin, so
        # its error is either a broken pipe (ffmpeg died) or a render/data failure from the generator.
        # Only the pipe case is an ffmpeg streaming failure; anything else must surface as itself so it
        # maps to the right failure category instead of being mislabeled a render-to-ffmpeg error.
        error = feed_error[0]
        if isinstance(error, OSError):
            raise FfmpegError(f"streaming blocks to ffmpeg failed: {error}") from error
        raise error
    if return_code != 0:
        raise FfmpegError(f"ffmpeg exited with code {return_code}: {stderr_output.strip()[-500:]}")

    return _parse_loudnorm_output(stderr_output)


def _parse_loudnorm_output(stderr_output: str) -> LoudnessMeasurement:
    # loudnorm prints one flat JSON object as the last brace-block in stderr; take the outermost
    # braces from the end (rindex) so a "{" in an earlier ffmpeg log line can't mis-anchor the slice.
    try:
        start = stderr_output.rindex("{")
        end = stderr_output.rindex("}") + 1
        payload = json.loads(stderr_output[start:end])
    except ValueError as e:
        raise LoudnormOutputError(f"no loudnorm JSON in ffmpeg output: {stderr_output.strip()[-500:]}") from e

    try:
        return LoudnessMeasurement(
            integrated_loudness_lkfs=float(payload["input_i"]),
            true_peak_dbtp=float(payload["input_tp"]),
        )
    except (KeyError, ValueError) as e:
        raise LoudnormOutputError(f"loudnorm JSON missing or non-numeric input_i/input_tp: {payload}") from e


def apply_loudness_measurement(
    measurement: LoudnessMeasurement,
    atmos_validation_builder: AtmosValidationBuilder,
) -> None:
    atmos_validation_builder.update_metadata(
        integrated_loudness_lkfs=measurement.integrated_loudness_lkfs,
        true_peak_dbtp=measurement.true_peak_dbtp,
    )

    _check_loudness(measurement, atmos_validation_builder)
    _check_true_peak(measurement, atmos_validation_builder)


def _check_loudness(measurement: LoudnessMeasurement, atmos_validation_builder: AtmosValidationBuilder) -> None:
    if measurement.integrated_loudness_lkfs > LOUDNESS_WARN_LKFS:
        atmos_validation_builder.warning(
            AtmosValidationFindingCode.LOUDNESS_TOO_HIGH,
            f"Integrated loudness is {measurement.integrated_loudness_lkfs} LKFS, "
            f"should not exceed {LOUDNESS_WARN_LKFS} LKFS",
        )


def _check_true_peak(measurement: LoudnessMeasurement, atmos_validation_builder: AtmosValidationBuilder) -> None:
    if measurement.true_peak_dbtp > TRUE_PEAK_WARN_DBTP:
        atmos_validation_builder.warning(
            AtmosValidationFindingCode.TRUE_PEAK_TOO_HIGH,
            f"True peak is {measurement.true_peak_dbtp} dBTP, should not exceed {TRUE_PEAK_WARN_DBTP} dBTP",
        )