from dataclasses import dataclass
from pathlib import Path
from typing import Tuple

import numpy as np


SAMPLE_RATE = 16_000
FRAME_SAMPLES = 480
# A silence run at least this long (in 30 ms frames) separates two sentences.
# speech_frame_mask bridges gaps of up to 5 frames, so use 6 frames (~180 ms).
MIN_PAUSE_FRAMES = 6
ONSET_MARGIN_MS = 50


class AudioAnalysisError(RuntimeError):
    pass


@dataclass(frozen=True)
class AudioMetrics:
    recording_duration_ms: int
    speech_duration_ms: int
    internal_silence_ms: int
    internal_pause_ratio: float


def decode_audio_mono(path: Path) -> Tuple[np.ndarray, int]:
    try:
        import av
    except ImportError as exc:
        raise AudioAnalysisError("PyAV is required for audio analysis.") from exc

    container = av.open(str(path))
    chunks = []
    try:
        audio_stream = next(
            (stream for stream in container.streams if stream.type == "audio"),
            None,
        )
        if audio_stream is None:
            raise AudioAnalysisError("The uploaded file does not contain an audio stream.")
        resampler = av.AudioResampler(format="fltp", layout="mono", rate=SAMPLE_RATE)
        for packet in container.demux(audio_stream):
            for frame in packet.decode():
                for output in resampler.resample(frame):
                    chunks.append(output.to_ndarray()[0].astype(np.float32, copy=False))
        for output in resampler.resample(None):
            chunks.append(output.to_ndarray()[0].astype(np.float32, copy=False))
    finally:
        container.close()

    if not chunks:
        raise AudioAnalysisError("The uploaded audio is empty.")
    return np.concatenate(chunks), SAMPLE_RATE


def analyze_audio(path: Path) -> AudioMetrics:
    samples, sample_rate = decode_audio_mono(path)
    return analyze_samples(samples, sample_rate)


def analyze_samples(samples: np.ndarray, sample_rate: int = SAMPLE_RATE) -> AudioMetrics:
    if samples.ndim != 1:
        samples = samples.reshape(-1)
    speech = speech_frame_mask(samples, sample_rate)
    speech_indexes = np.flatnonzero(speech)
    if speech_indexes.size == 0:
        raise AudioAnalysisError("No usable speech was detected in the recording.")

    recording_duration_ms = max(1, int(round(samples.size / sample_rate * 1000)))
    window = max(1, int(round(sample_rate * 0.03)))
    frame_ms = window / sample_rate * 1000
    speech_duration_ms = max(1, int(round(speech.sum() * frame_ms)))
    first = int(speech_indexes[0])
    last = int(speech_indexes[-1])
    internal_frames = max(1, last - first + 1)
    internal_silence_frames = int((~speech[first : last + 1]).sum())
    internal_silence_ms = int(round(internal_silence_frames * frame_ms))
    internal_pause_ratio = internal_silence_frames / internal_frames
    return AudioMetrics(
        recording_duration_ms=recording_duration_ms,
        speech_duration_ms=speech_duration_ms,
        internal_silence_ms=internal_silence_ms,
        internal_pause_ratio=round(internal_pause_ratio, 4),
    )


def speech_frame_mask(samples: np.ndarray, sample_rate: int = SAMPLE_RATE) -> np.ndarray:
    """Per-frame speech mask (30 ms frames) using the same RMS threshold as analyze_samples."""
    if samples.ndim != 1:
        samples = samples.reshape(-1)
    if samples.size == 0 or sample_rate <= 0:
        raise AudioAnalysisError("The uploaded audio is empty.")

    window = max(1, int(round(sample_rate * 0.03)))
    complete_frames = int(np.ceil(samples.size / window))
    padded = np.pad(samples, (0, complete_frames * window - samples.size))
    frames = padded.reshape(complete_frames, window).astype(np.float64, copy=False)
    rms = np.sqrt(np.mean(frames * frames, axis=1))

    signal_level = float(np.percentile(rms, 95))
    if signal_level < 0.002:
        raise AudioAnalysisError("No usable speech was detected in the recording.")
    noise_floor = float(np.percentile(rms, 10))
    threshold = min(signal_level * 0.45, max(0.006, noise_floor * 2.2))
    speech = rms >= threshold

    # Treat very short gaps inside a word as speech, then reject short clicks.
    _bridge_false_runs(speech, max_frames=5)
    _remove_true_runs(speech, max_frames=2)
    return speech


def refine_sentence_end_ms(
    samples: np.ndarray,
    sample_rate: int = SAMPLE_RATE,
    *,
    raw_end_ms: int,
    padded_end_ms: int,
) -> int:
    """Pull a padded sentence end back to just before the next sentence's speech.

    Whisper's timestamps are not always aligned with the real audio: the next
    segment's start can be later than the actual speech onset, so a fixed
    end-padding may occasionally run into the next sentence's beginning. This
    finds the first silence run of at least MIN_PAUSE_FRAMES inside the padded
    region and stops the sentence just before the speech that follows it.
    """
    if padded_end_ms <= raw_end_ms:
        return padded_end_ms
    frame_ms = 1000 * max(1, int(round(sample_rate * 0.03))) / sample_rate
    start_sample = max(0, int(raw_end_ms / 1000 * sample_rate))
    end_sample = min(samples.size, int(padded_end_ms / 1000 * sample_rate))
    if end_sample <= start_sample:
        return padded_end_ms
    try:
        speech = speech_frame_mask(samples[start_sample:end_sample], sample_rate)
    except AudioAnalysisError:
        return padded_end_ms
    silence_frames = 0
    for index, is_speech in enumerate(speech):
        if not is_speech:
            silence_frames += 1
            continue
        if silence_frames >= MIN_PAUSE_FRAMES:
            onset_ms = int(raw_end_ms + index * frame_ms)
            return max(raw_end_ms, min(padded_end_ms, onset_ms - ONSET_MARGIN_MS))
        silence_frames = 0
    return padded_end_ms


def _bridge_false_runs(values: np.ndarray, max_frames: int) -> None:
    start = None
    for index, value in enumerate(values):
        if not value and start is None:
            start = index
        elif value and start is not None:
            if start > 0 and index - start <= max_frames:
                values[start:index] = True
            start = None


def _remove_true_runs(values: np.ndarray, max_frames: int) -> None:
    start = None
    for index, value in enumerate(values):
        if value and start is None:
            start = index
        elif not value and start is not None:
            if index - start <= max_frames:
                values[start:index] = False
            start = None
    if start is not None and len(values) - start <= max_frames:
        values[start:] = False