app/services/pause_detection_service.py

"""Pause detection and pause-aware transcript assembly for STT dictation.

This module is intentionally separate from VAD preprocessing. Pause detectors
observe full audio and produce punctuation metadata; they never trim, reject, or
replace audio before Whisper sees it.
"""
from __future__ import annotations

import os
import time
import wave
from array import array
from dataclasses import asdict, dataclass
from typing import Any, Literal

from app.services.transcription_service import STTSegment
from app.services.voice_activity_service import VoiceActivityService, cleanup_vad_result

PauseDetector = Literal["rms_energy", "silero"]

FRAME_MS = 20
DEFAULT_RMS_THRESHOLD = 0.018
MIN_SPEECH_REGION_MS = 120.0
MERGE_SPEECH_GAP_MS = 300.0
MIN_PAUSE_MS = 1200.0
SHORT_PAUSE_MAX_MS = 1800.0
MEDIUM_PAUSE_MAX_MS = 3000.0


@dataclass
class Pause:
    start_ms: float
    end_ms: float
    duration_ms: float
    symbol: str
    type: str = "internal_pause"

    def model_dump(self) -> dict[str, Any]:
        return asdict(self)


@dataclass
class SpeechRegion:
    start_ms: float
    end_ms: float

    def model_dump(self) -> dict[str, Any]:
        return asdict(self)


@dataclass
class PauseDetectionResult:
    enabled: bool
    detector: str
    detection_time_ms: float
    pauses: list[Pause]
    speech_regions: list[SpeechRegion]
    pause_count: int
    speech_regions_count: int
    threshold: float | None = None
    frame_ms: int = FRAME_MS
    fallback_used: bool = False
    error: str | None = None

    def model_dump(self) -> dict[str, Any]:
        data = asdict(self)
        data["pauses"] = [pause.model_dump() for pause in self.pauses]
        data["speech_regions"] = [region.model_dump() for region in self.speech_regions]
        return data


@dataclass
class PauseAwareTranscript:
    pause_text: str
    inserted_pause_count: int
    insertion_strategy: str

    def model_dump(self) -> dict[str, Any]:
        return asdict(self)


class PauseDetectionService:
    """Detect pauses on full 16 kHz mono PCM WAV audio."""

    def __init__(self) -> None:
        self._vad = VoiceActivityService()

    def detect(self, wav_path: str, detector: str = "rms_energy") -> PauseDetectionResult:
        normalized = (detector or "rms_energy").strip().lower()
        if normalized == "rms":
            normalized = "rms_energy"
        if normalized == "rms_energy":
            return self._rms_energy(wav_path)
        if normalized == "silero":
            return self._vad_regions(wav_path, normalized)
        start = time.perf_counter()
        return PauseDetectionResult(
            enabled=False,
            detector=normalized,
            detection_time_ms=round((time.perf_counter() - start) * 1000, 2),
            pauses=[],
            speech_regions=[],
            pause_count=0,
            speech_regions_count=0,
            fallback_used=True,
            error=f"unsupported_pause_detector:{normalized}",
        )

    def build_pause_text(
        self,
        transcript: str,
        stt_segments: list[STTSegment] | None,
        pauses: list[Pause],
    ) -> PauseAwareTranscript:
        transcript = (transcript or "").strip()
        usable_pauses = [
            pause
            for pause in pauses
            if pause.type == "internal_pause" and pause.duration_ms >= MIN_PAUSE_MS
        ]
        if not transcript or not usable_pauses:
            return PauseAwareTranscript(transcript, 0, "none")

        timed_segments = [
            segment for segment in (stt_segments or [])
            if segment.text.strip() and segment.start_ms is not None and segment.end_ms is not None
        ]
        if not timed_segments:
            return PauseAwareTranscript(transcript, 0, "none")

        if len(timed_segments) == 1:
            segment = timed_segments[0]
            pieces, used = _split_segment_with_pauses(
                segment.text,
                float(segment.start_ms or 0.0),
                float(segment.end_ms or segment.start_ms or 0.0),
                list(enumerate(usable_pauses)),
            )
            if not used:
                return PauseAwareTranscript(transcript, 0, "single_segment")
            return PauseAwareTranscript(_render_tokens(pieces), len(used), "single_segment")

        tokens: list[str] = []
        inserted: set[int] = set()
        for index, segment in enumerate(timed_segments):
            segment_start = float(segment.start_ms or 0.0)
            segment_end = float(segment.end_ms or segment_start)
            tokens.append(segment.text)

            next_segment_start = float(timed_segments[index + 1].start_ms or segment_end) if index + 1 < len(timed_segments) else None
            if next_segment_start is None:
                continue

            boundary_pauses = [
                (pause_index, pause)
                for pause_index, pause in enumerate(usable_pauses)
                if pause_index not in inserted
                and pause.start_ms >= segment_end
                and pause.end_ms <= next_segment_start
            ]
            if not boundary_pauses:
                continue

            boundary_pauses.sort(key=lambda item: item[1].start_ms)
            for pause_index, pause in boundary_pauses:
                tokens.append(pause.symbol)
                inserted.add(pause_index)

        return PauseAwareTranscript(_render_tokens(tokens), len(inserted), "segment_boundary")

    def _rms_energy(self, wav_path: str) -> PauseDetectionResult:
        start = time.perf_counter()
        threshold = float(os.getenv("PAUSE_RMS_THRESHOLD", str(DEFAULT_RMS_THRESHOLD)))
        try:
            samples, sample_rate, channels, sample_width = _read_pcm_wav(wav_path)
            if sample_rate != 16000 or channels != 1 or sample_width != 2:
                raise ValueError("pause_detection_requires_16khz_mono_s16_wav")
            regions = _speech_regions_from_rms(samples, sample_rate, threshold)
            pauses = _pauses_from_regions(regions)
            return PauseDetectionResult(
                enabled=True,
                detector="rms_energy",
                detection_time_ms=round((time.perf_counter() - start) * 1000, 2),
                pauses=pauses,
                speech_regions=regions,
                pause_count=len(pauses),
                speech_regions_count=len(regions),
                threshold=threshold,
                fallback_used=False,
            )
        except Exception as exc:
            return PauseDetectionResult(
                enabled=True,
                detector="rms_energy",
                detection_time_ms=round((time.perf_counter() - start) * 1000, 2),
                pauses=[],
                speech_regions=[],
                pause_count=0,
                speech_regions_count=0,
                threshold=threshold,
                fallback_used=True,
                error=repr(exc),
            )

    def _vad_regions(self, wav_path: str, detector: str) -> PauseDetectionResult:
        start = time.perf_counter()
        vad_result = None
        try:
            vad_result = self._vad.process(wav_path, detector)
            regions = [
                SpeechRegion(
                    start_ms=round(float(region.get("start_ms", 0.0)), 2),
                    end_ms=round(float(region.get("end_ms", 0.0)), 2),
                )
                for region in (vad_result.speech_regions or [])
                if region.get("end_ms") is not None and region.get("start_ms") is not None
            ]
            pauses = _pauses_from_regions(regions)
            return PauseDetectionResult(
                enabled=True,
                detector=detector,
                detection_time_ms=round((time.perf_counter() - start) * 1000, 2),
                pauses=pauses,
                speech_regions=regions,
                pause_count=len(pauses),
                speech_regions_count=len(regions),
                threshold=None,
                fallback_used=bool(vad_result.fallback_used),
                error=vad_result.error,
            )
        except Exception as exc:
            return PauseDetectionResult(
                enabled=True,
                detector=detector,
                detection_time_ms=round((time.perf_counter() - start) * 1000, 2),
                pauses=[],
                speech_regions=[],
                pause_count=0,
                speech_regions_count=0,
                fallback_used=True,
                error=repr(exc),
            )
        finally:
            if vad_result:
                cleanup_vad_result(vad_result)

    def _build_proportional_transcript(self, transcript: str, pauses: list[Pause]) -> PauseAwareTranscript:
        words = transcript.split()
        if len(words) < 2:
            return PauseAwareTranscript(transcript, 0, "none")
        ordered = sorted(pauses, key=lambda pause: pause.start_ms)
        tokens: list[str] = []
        pause_index = 0
        for index, word in enumerate(words):
            tokens.append(word)
            proportion = (index + 1) / max(1, len(words))
            while pause_index < len(ordered) and pause_index / max(1, len(ordered)) < proportion:
                tokens.append(ordered[pause_index].symbol)
                pause_index += 1
        return PauseAwareTranscript(_render_tokens(tokens), pause_index, "proportional")


def _read_pcm_wav(wav_path: str) -> tuple[array, int, int, int]:
    with wave.open(wav_path, "rb") as wav:
        sample_rate = wav.getframerate()
        channels = wav.getnchannels()
        sample_width = wav.getsampwidth()
        frames = wav.readframes(wav.getnframes())
    if sample_width != 2:
        raise ValueError("expected_16bit_pcm_wav")
    samples = array("h")
    samples.frombytes(frames)
    return samples, sample_rate, channels, sample_width


def _speech_regions_from_rms(samples: array, sample_rate: int, threshold: float) -> list[SpeechRegion]:
    frame_size = max(1, int(sample_rate * FRAME_MS / 1000))
    raw_regions: list[tuple[int, int]] = []
    speech_start: int | None = None
    for start in range(0, len(samples), frame_size):
        end = min(len(samples), start + frame_size)
        frame = samples[start:end]
        rms = _rms(frame)
        if rms >= threshold and speech_start is None:
            speech_start = start
        elif rms < threshold and speech_start is not None:
            raw_regions.append((speech_start, start))
            speech_start = None
    if speech_start is not None:
        raw_regions.append((speech_start, len(samples)))

    min_speech_samples = int(sample_rate * MIN_SPEECH_REGION_MS / 1000)
    merge_gap_samples = int(sample_rate * MERGE_SPEECH_GAP_MS / 1000)
    filtered = [(start, end) for start, end in raw_regions if end - start >= min_speech_samples]
    if not filtered:
        return []

    merged: list[tuple[int, int]] = [filtered[0]]
    for start, end in filtered[1:]:
        previous_start, previous_end = merged[-1]
        if start - previous_end <= merge_gap_samples:
            merged[-1] = (previous_start, end)
        else:
            merged.append((start, end))

    return [
        SpeechRegion(
            start_ms=_samples_to_ms(start, sample_rate),
            end_ms=_samples_to_ms(end, sample_rate),
        )
        for start, end in merged
    ]


def _pauses_from_regions(regions: list[SpeechRegion]) -> list[Pause]:
    pauses: list[Pause] = []
    ordered = sorted(regions, key=lambda region: region.start_ms)
    for left, right in zip(ordered, ordered[1:]):
        duration_ms = round(max(0.0, right.start_ms - left.end_ms), 2)
        symbol = _pause_symbol(duration_ms)
        if not symbol:
            continue
        pauses.append(Pause(
            start_ms=round(left.end_ms, 2),
            end_ms=round(right.start_ms, 2),
            duration_ms=duration_ms,
            symbol=symbol,
        ))
    return pauses


def _pause_symbol(duration_ms: float) -> str | None:
    if duration_ms < MIN_PAUSE_MS:
        return None
    if duration_ms < SHORT_PAUSE_MAX_MS:
        return "..."
    if duration_ms < MEDIUM_PAUSE_MAX_MS:
        return "......"
    return "........."


def _split_segment_with_pauses(
    text: str,
    start_ms: float,
    end_ms: float,
    pauses: list[tuple[int, Pause]],
) -> tuple[list[str], set[int]]:
    words = text.split()
    if len(words) < 2 or end_ms <= start_ms:
        return [text], set()
    indexed = sorted(pauses, key=lambda item: item[1].start_ms)
    insertions: dict[int, list[str]] = {}
    used: set[int] = set()
    for index, pause in indexed:
        midpoint = pause.start_ms + (pause.duration_ms / 2)
        ratio = max(0.0, min(1.0, (midpoint - start_ms) / (end_ms - start_ms)))
        word_index = max(1, min(len(words) - 1, round(ratio * len(words))))
        insertions.setdefault(word_index, []).append(pause.symbol)
        used.add(index)

    pieces: list[str] = []
    for index, word in enumerate(words):
        if index in insertions:
            pieces.extend(insertions[index])
        pieces.append(word)
    return pieces, used


def _render_tokens(tokens: list[str]) -> str:
    text = ""
    for token in tokens:
        clean = str(token or "").strip()
        if not clean:
            continue
        if set(clean) == {"."} and len(clean) >= 3:
            text = text.rstrip() + clean
        else:
            if text and not text.endswith(" "):
                text += " "
            text += clean
    return text.strip()


def _rms(samples: array) -> float:
    if not samples:
        return 0.0
    total = 0.0
    for sample in samples:
        value = sample / 32768.0
        total += value * value
    return (total / len(samples)) ** 0.5


def _samples_to_ms(sample_index: int, sample_rate: int) -> float:
    return round((sample_index / sample_rate) * 1000, 2) if sample_rate else 0.0