from dataclasses import dataclass
from pathlib import Path
from typing import Optional, 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
START_LOOKBACK_MS = 300
START_LOOKAHEAD_MS = 400


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 refine_sentence_start_ms(
    samples: np.ndarray,
    sample_rate: int = SAMPLE_RATE,
    *,
    raw_start_ms: int,
    previous_end_ms: int,
) -> int:
    """Pull a sentence start forward to the actual speech onset.

    Whisper's segment start can be earlier than the real speech onset (the tail
    of the previous sentence is still audible at the beginning of the next
    sentence). When the raw start lands inside continuous speech, this looks for
    the last pause of at least MIN_PAUSE_FRAMES before the next speech run and
    moves the start to that onset, so adjacent sentences share no audio.
    """
    if raw_start_ms <= previous_end_ms:
        return max(previous_end_ms, raw_start_ms)
    frame_ms = 1000 * max(1, int(round(sample_rate * 0.03))) / sample_rate
    total_ms = samples.size / sample_rate * 1000
    window_start_ms = max(0, raw_start_ms - START_LOOKBACK_MS)
    window_end_ms = min(int(total_ms), raw_start_ms + START_LOOKAHEAD_MS)
    start_sample = int(window_start_ms / 1000 * sample_rate)
    end_sample = min(samples.size, int(window_end_ms / 1000 * sample_rate))
    if end_sample - start_sample < max(1, int(round(sample_rate * 0.03))):
        return raw_start_ms
    try:
        speech = speech_frame_mask(samples[start_sample:end_sample], sample_rate)
    except AudioAnalysisError:
        return raw_start_ms
    onset_frame = int(round((raw_start_ms - window_start_ms) / frame_ms))
    if (
        0 <= onset_frame < len(speech)
        and speech[onset_frame]
        and (onset_frame == 0 or not speech[onset_frame - 1])
    ):
        return raw_start_ms  # already at a speech onset
    boundary_ms: Optional[float] = None
    silence_frames = 0
    for index, is_speech in enumerate(speech):
        if not is_speech:
            silence_frames += 1
            continue
        if silence_frames >= MIN_PAUSE_FRAMES:
            boundary_ms = window_start_ms + index * frame_ms
        silence_frames = 0
    if boundary_ms is None:
        return raw_start_ms
    boundary_ms = min(boundary_ms, raw_start_ms + START_LOOKAHEAD_MS)
    return max(previous_end_ms, int(boundary_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