#!/usr/bin/env python3
"""Classify what happens between spoken words, from the audio itself.

A word-level transcript says when speech happens. It does not say whether the
0.4 s between two words is a breath, a lip smack, a thinking pause, or dead room —
and those want completely different treatment. Cutting all of them equally is what
makes a recut sound machine-gunned; cutting none of them is what makes it slow.

So each gap is measured, not guessed:

  loudness   relative to the speaker's own median speech level, because absolute
             dBFS says nothing about a quiet recording;
  spectral   centroid and high-frequency ratio — a breath is broadband with energy
             above 1 kHz, while room tone sits low;
  shape      crest factor and onset count, which separate a click or lip smack
             from a sustained inhale.

Filler *sounds* ("э-э", "мм") are treated separately from filler *words* ("ну",
"вот", "типа"): the first are safe to drop, the second carry rhythm and often
meaning, so they are only ever reported as candidates for a human decision.
"""
from __future__ import annotations

import math
import subprocess
from dataclasses import dataclass, field
from pathlib import Path

import numpy as np

ANALYSIS_RATE = 16_000

# Pure hesitation sounds. Safe to remove when they stand alone: they carry no
# lexical content in Russian and are what listeners hear as stumbling.
HESITATION_SOUNDS = {
    "э", "ээ", "эээ", "э-э", "э-э-э", "эм", "эмм", "ммм", "мм", "м-м", "м-м-м",
    "а-а", "аа", "ааа", "ну-у", "э-эм", "хм", "кхм", "мгм", "угу",
}

# Discourse markers. Frequently redundant, but they also set rhythm and can change
# meaning ("вот" as a pointer, "значит" as a connective). Reported, never cut
# automatically.
FILLER_WORD_CANDIDATES = {
    "ну", "вот", "типа", "значит", "короче", "собственно", "как-бы", "как бы",
    "в общем", "в-общем", "то есть", "то-есть", "скажем", "допустим",
    "получается", "реально", "просто", "буквально", "прям", "прямо",
}


@dataclass
class Gap:
    index: int
    start: float
    end: float
    after_word: str
    before_word: str
    after_sentence_end: bool
    kind: str = "unknown"
    rms_db: float = 0.0
    relative_db: float = 0.0
    centroid_hz: float = 0.0
    high_ratio: float = 0.0
    crest_db: float = 0.0
    onsets: int = 0
    keep: float = 0.0
    reason: str = ""

    @property
    def duration(self) -> float:
        return self.end - self.start

    def as_dict(self) -> dict:
        return {
            "index": self.index,
            "start": round(self.start, 3),
            "end": round(self.end, 3),
            "duration": round(self.duration, 3),
            "kind": self.kind,
            "after_word": self.after_word,
            "before_word": self.before_word,
            "after_sentence_end": self.after_sentence_end,
            "rms_db": round(self.rms_db, 1),
            "relative_db": round(self.relative_db, 1),
            "centroid_hz": round(self.centroid_hz),
            "high_ratio": round(self.high_ratio, 3),
            "crest_db": round(self.crest_db, 1),
            "onsets": self.onsets,
            "keep": round(self.keep, 3),
            "trimmed": round(max(0.0, self.duration - self.keep), 3),
            "reason": self.reason,
        }


@dataclass
class SpeechProfile:
    """What the speaker's own audio looks like, used as the reference level."""

    speech_rms_db: float
    noise_floor_db: float
    speech_centroid_hz: float
    samples: int
    rate: int = ANALYSIS_RATE
    warnings: list[str] = field(default_factory=list)

    def as_dict(self) -> dict:
        return {
            "speech_rms_db": round(self.speech_rms_db, 1),
            "noise_floor_db": round(self.noise_floor_db, 1),
            "speech_to_noise_db": round(self.speech_rms_db - self.noise_floor_db, 1),
            "speech_centroid_hz": round(self.speech_centroid_hz),
            "analysis_rate": self.rate,
            "warnings": self.warnings,
        }


def load_audio(media: Path, rate: int = ANALYSIS_RATE) -> np.ndarray:
    """Decode to mono PCM once. int16 keeps an hour of audio near 115 MB."""
    completed = subprocess.run(
        ["ffmpeg", "-v", "error", "-nostdin", "-i", str(media),
         "-map", "0:a:0", "-ac", "1", "-ar", str(rate),
         "-f", "s16le", "-"],
        capture_output=True,
    )
    if completed.returncode or not completed.stdout:
        detail = (completed.stderr or b"").decode("utf-8", errors="replace").strip()[-300:]
        raise SystemExit(f"Could not decode audio from {media}: {detail}")
    return np.frombuffer(completed.stdout, dtype="<i2").astype(np.float32) / 32768.0


def db(value: float) -> float:
    return 20.0 * math.log10(max(value, 1e-9))


def window_stats(chunk: np.ndarray, rate: int) -> tuple[float, float, float, float]:
    """(rms_db, centroid_hz, high_ratio, crest_db) for one slice of audio."""
    if chunk.size < 32:
        return -120.0, 0.0, 0.0, 0.0
    rms = float(np.sqrt(np.mean(chunk ** 2)))
    peak = float(np.max(np.abs(chunk)))
    spectrum = np.abs(np.fft.rfft(chunk * np.hanning(chunk.size)))
    frequencies = np.fft.rfftfreq(chunk.size, 1.0 / rate)
    total = float(spectrum.sum()) or 1e-9
    centroid = float((spectrum * frequencies).sum() / total)
    high = float(spectrum[frequencies >= 1500].sum() / total)
    return db(rms), centroid, high, db(peak) - db(rms)


def profile_speech(audio: np.ndarray, words: list[dict], rate: int = ANALYSIS_RATE) -> SpeechProfile:
    """Measure the speaker's own level from the words, and the floor from silence."""
    speech_levels: list[float] = []
    centroids: list[float] = []
    for word in words:
        begin = int(float(word["start"]) * rate)
        finish = int(float(word["end"]) * rate)
        chunk = audio[begin:finish]
        if chunk.size < rate // 100:
            continue
        level, centroid, _high, _crest = window_stats(chunk, rate)
        speech_levels.append(level)
        centroids.append(centroid)
    if not speech_levels:
        raise SystemExit("No word intervals overlap the audio; check the transcript timings")
    speech_rms = float(np.median(speech_levels))

    # Noise floor from the quietest tenth of 100 ms frames, which lands in room
    # tone rather than in speech even when the recording is dense.
    frame = rate // 10
    usable = audio[: (audio.size // frame) * frame].reshape(-1, frame)
    if usable.size:
        frame_levels = 20 * np.log10(np.maximum(np.sqrt((usable ** 2).mean(axis=1)), 1e-9))
        floor = float(np.percentile(frame_levels, 10))
    else:
        floor = -90.0

    warnings: list[str] = []
    if speech_rms - floor < 18:
        warnings.append(
            f"Отношение речь/шум всего {speech_rms - floor:.0f} дБ. Классификация вздохов "
            f"будет неуверенной — проверь спорные паузы вручную."
        )
    if speech_rms > -8:
        warnings.append("Речь записана почти в клиппинг: проверь исходник на перегруз.")
    return SpeechProfile(
        speech_rms_db=speech_rms,
        noise_floor_db=floor,
        speech_centroid_hz=float(np.median(centroids)) if centroids else 0.0,
        samples=int(audio.size),
        rate=rate,
        warnings=warnings,
    )


SENTENCE_ENDINGS = ".!?…"


def build_gaps(words: list[dict], audio_duration: float) -> list[Gap]:
    gaps: list[Gap] = []
    for index in range(len(words) - 1):
        current = words[index]
        following = words[index + 1]
        start = float(current["end"])
        end = float(following["start"])
        if end - start <= 0.001:
            continue
        gaps.append(Gap(
            index=index,
            start=start,
            end=end,
            after_word=str(current.get("word", "")),
            before_word=str(following.get("word", "")),
            after_sentence_end=str(current.get("word", "")).rstrip()[-1:] in SENTENCE_ENDINGS,
        ))
    _ = audio_duration
    return gaps


def count_onsets(chunk: np.ndarray, rate: int) -> int:
    """Number of sharp energy rises — separates a click from a sustained breath."""
    if chunk.size < rate // 50:
        return 0
    window = max(8, rate // 500)
    energy = np.convolve(chunk ** 2, np.ones(window) / window, mode="same")
    smoothed = np.sqrt(np.maximum(energy, 1e-12))
    threshold = float(np.percentile(smoothed, 60)) * 2.4
    above = smoothed > max(threshold, 1e-6)
    return int(np.count_nonzero(np.diff(above.astype(np.int8)) == 1))


# How much of each kind of gap survives, per editing style. Measured in seconds.
STYLES: dict[str, dict[str, float]] = {
    "gentle": {"breath": 0.18, "breath_long": 0.26, "click": 0.08,
               "sentence_pause": 0.38, "thinking_pause": 0.30, "dead_air": 0.30,
               "uncertain": 0.45},
    "balanced": {"breath": 0.11, "breath_long": 0.18, "click": 0.05,
                 "sentence_pause": 0.26, "thinking_pause": 0.22, "dead_air": 0.18,
                 "uncertain": 0.34},
    "tight": {"breath": 0.07, "breath_long": 0.12, "click": 0.03,
              "sentence_pause": 0.18, "thinking_pause": 0.14, "dead_air": 0.12,
              "uncertain": 0.26},
}
DEFAULT_STYLE = "balanced"

# Level bands relative to the speaker's own median speech level.
#
# Level turned out to be the only reliable discriminator on real footage: a
# low-frequency energy ratio overlaps heavily between voiced speech (0.31 median)
# and breaths (0.16-0.30), while level separates them cleanly — genuine breaths sat
# 9-19 dB below speech, whereas every "breath" within 3 dB of speech level was
# actually speech that the transcript's word boundaries had missed. Cutting those
# would clip words, so anything near speech level is kept whole.
SPEECH_BLEED_DB = -6.0
QUIET_DB = -18.0


def classify(gaps: list[Gap], audio: np.ndarray, profile: SpeechProfile,
             style: str = DEFAULT_STYLE) -> list[Gap]:
    """Label each gap and decide how much of it to keep.

    Ordered so that the safe answer wins ties: anything that might be speech is
    kept in full, and only clearly non-verbal material is shortened.
    """
    budget = STYLES.get(style) or STYLES[DEFAULT_STYLE]
    rate = profile.rate
    for gap in gaps:
        begin = int(gap.start * rate)
        finish = int(gap.end * rate)
        chunk = audio[begin:finish]
        gap.rms_db, gap.centroid_hz, gap.high_ratio, gap.crest_db = window_stats(chunk, rate)
        gap.relative_db = gap.rms_db - profile.speech_rms_db
        gap.onsets = count_onsets(chunk, rate)
        duration = gap.duration

        if duration < 0.12:
            gap.kind = "coarticulation"
            gap.keep = duration
            gap.reason = "Короче 120 мс: это часть произношения, резать нельзя."
            continue

        if gap.relative_db > SPEECH_BLEED_DB:
            gap.kind = "speech_bleed"
            gap.keep = duration
            gap.reason = (
                f"Уровень всего {gap.relative_db:+.1f} дБ от речи — это почти наверняка "
                f"речь, не попавшая в границы слов. Не режем."
            )
            continue

        breathy = gap.high_ratio >= 0.35 or gap.centroid_hz >= 1800
        clicky = duration < 0.3 and gap.crest_db > 18 and gap.onsets <= 2 and not breathy
        quiet = gap.relative_db <= QUIET_DB or gap.rms_db < profile.noise_floor_db + 6

        if clicky:
            gap.kind = "click"
            gap.keep = budget["click"]
            gap.reason = "Щелчок или причмокивание: убирается почти полностью."
        elif gap.after_sentence_end and duration <= 0.9:
            # A breath after a full stop doubles as the sentence's beat; treating
            # it as a breath would strip the punctuation's rhythm.
            gap.kind = "sentence_pause"
            gap.keep = budget["sentence_pause"]
            gap.reason = "Пауза после точки: держим ритм, но не тянем."
        elif breathy:
            gap.kind = "breath"
            # Never removed outright: with no air at all the next phrase sounds
            # spliced. A short remnant keeps the delivery human.
            gap.keep = budget["breath_long"] if duration > 0.6 else budget["breath"]
            gap.reason = "Вздох: оставляем короткий остаток, чтобы стык не звучал склеенным."
        elif quiet and duration > 0.7:
            gap.kind = "dead_air"
            gap.keep = budget["dead_air"]
            gap.reason = "Долгая тихая пауза: сжимаем сильно."
        elif quiet:
            gap.kind = "thinking_pause"
            gap.keep = budget["thinking_pause"]
            gap.reason = "Тихая пауза раздумья: слегка подрезаем."
        else:
            gap.kind = "uncertain"
            gap.keep = budget["uncertain"]
            gap.reason = (
                f"Не тихо ({gap.relative_db:+.1f} дБ) и не похоже на вздох "
                f"(ВЧ {gap.high_ratio:.2f}). Подрезаем осторожно, проверь на слух."
            )
        gap.keep = max(0.0, min(duration, gap.keep))
    return gaps


def normalise_token(word: str) -> str:
    return "".join(
        character for character in word.casefold()
        if character.isalnum() or character in "-"
    )


def find_hesitations(words: list[dict], gaps: list[Gap],
                     probability_floor: float = 0.55) -> list[dict]:
    """Standalone hesitation sounds that are safe to remove."""
    by_index = {gap.index: gap for gap in gaps}
    found = []
    for index, word in enumerate(words):
        token = normalise_token(str(word.get("word", "")))
        if token not in HESITATION_SOUNDS:
            continue
        before = by_index.get(index - 1)
        after = by_index.get(index)
        # Only when it stands apart from the words around it; a hesitation glued
        # to the next word cannot be lifted out without clipping that word.
        isolated = (before is None or before.duration >= 0.06) and \
                   (after is None or after.duration >= 0.06)
        found.append({
            "word_id": word.get("id", index),
            "word": word.get("word"),
            "start": round(float(word["start"]), 3),
            "end": round(float(word["end"]), 3),
            "probability": word.get("probability"),
            "isolated": isolated,
            "safe_to_remove": isolated,
            "reason": (
                "Изолированный звук хезитации." if isolated
                else "Примыкает к соседнему слову — удаление обрежет речь."
            ),
        })
    _ = probability_floor
    return found


def find_filler_candidates(words: list[dict]) -> list[dict]:
    """Discourse markers. Reported for review, never removed automatically."""
    candidates = []
    for index, word in enumerate(words):
        token = normalise_token(str(word.get("word", "")))
        if token in FILLER_WORD_CANDIDATES:
            candidates.append({
                "word_id": word.get("id", index),
                "word": word.get("word"),
                "start": round(float(word["start"]), 3),
                "end": round(float(word["end"]), 3),
                "note": (
                    "Слово-филлер. Часто лишнее, но может держать ритм или менять "
                    "смысл — решает человек, автоматически не удаляется."
                ),
            })
    return candidates


def summarise(gaps: list[Gap]) -> dict:
    kinds: dict[str, dict] = {}
    for gap in gaps:
        entry = kinds.setdefault(gap.kind, {"count": 0, "seconds": 0.0, "removed": 0.0})
        entry["count"] += 1
        entry["seconds"] += gap.duration
        entry["removed"] += max(0.0, gap.duration - gap.keep)
    for entry in kinds.values():
        entry["seconds"] = round(entry["seconds"], 2)
        entry["removed"] = round(entry["removed"], 2)
    return kinds


def analyse(media: Path, words: list[dict],
            style: str = DEFAULT_STYLE) -> tuple[SpeechProfile, list[Gap], dict]:
    audio = load_audio(media)
    profile = profile_speech(audio, words)
    gaps = classify(build_gaps(words, audio.size / profile.rate), audio, profile, style)
    extras = {
        "style": style,
        "thresholds": {
            "speech_bleed_db": SPEECH_BLEED_DB,
            "quiet_db": QUIET_DB,
            "keep_budget": STYLES.get(style) or STYLES[DEFAULT_STYLE],
        },
        "hesitations": find_hesitations(words, gaps),
        "filler_candidates": find_filler_candidates(words),
        "gap_summary": summarise(gaps),
    }
    return profile, gaps, extras