"""Распознавание речи с разделением по говорящим.

Модели загружаются один раз при старте сервиса и живут в памяти: холодная
загрузка GigaAM занимает секунды, и делать её на каждый запрос недопустимо.
"""
import logging
import shutil
import subprocess
import time
import wave
from pathlib import Path

import numpy as np

from app.acoustics import segment_acoustics, separation_quality
from app.text import (apply_replacements, load_replacements,
                      normalize_typography, split_sentences)
from app.turns import Segment, chunk_ranges, merge_turns, speech_stats

log = logging.getLogger(__name__)

SAMPLE_RATE = 16000
# GigaAM падает на кусках длиннее ~200 с, держим запас.
# Предел зависит от сборки onnxruntime: на macOS реплика в 190 с проходит,
# на Windows та же падает в свёртке энкодера. Берём с большим запасом, а на
# случай сюрпризов есть деление куска пополам при ошибке.
MAX_CHUNK_SEC = 60.0
# Ниже этого делить бессмысленно: явно не длина виновата.
MIN_SPLIT_SEC = 2.0
# Модель отпечатков рассчитана на короткий фрагмент и на длинном падает.
# Для оценки голоса секунд более чем достаточно.
EMBED_SEC = 8.0
# Ниже этого значения голоса практически неразличимы и разметка по говорящим
# случайна. На записях с одним микрофоном в комнате так бывает часто.
RELIABLE_SEPARATION = 0.35
# Промежутки короче этого - обычные паузы в разговоре, туда лезть незачем.
GAP_MIN_SEC = 3.0
# Короткий обрывок из промежутка почти всегда шум, принятый за слово:
# шелест бумаг и стук клавиш на записи со стола дают почти ту же громкость,
# что и речь. Порог отсекает их, не трогая настоящие фразы.
MIN_RECOVERED_CHARS = 15
ASR_MODEL_NAME = "gigaam-v3-e2e-rnnt"
ASR_MODEL_REL = Path("gigaam")
SEG_MODEL_REL = Path("sherpa-onnx-pyannote-segmentation-3-0") / "model.onnx"
EMB_MODEL_REL = Path("nemo_en_titanet_small.onnx")


class ModelsMissing(RuntimeError):
    """Модели не скачаны - сервис не может работать."""


def find_ffmpeg(base_dir: Path) -> str:
    """Ищет ffmpeg сначала в папке пакета, потом в системе."""
    for candidate in (base_dir / "bin" / "ffmpeg.exe", base_dir / "bin" / "ffmpeg"):
        if candidate.is_file():
            return str(candidate)
    found = shutil.which("ffmpeg")
    if not found:
        raise ModelsMissing(
            "не найден ffmpeg - положите ffmpeg.exe в папку bin рядом с сервисом")
    return found


# Выравнивание громкости - решающий шаг для записей с одним микрофоном на столе.
# На проверочных разговорах доля второго участника выросла с 1.8 до 24.3 процента,
# а переключений между репликами - с 21 до 71 процента. Распознанный текст при
# этом не меняется: сравнение фрагментов дало полностью совпадающие расшифровки.
# Параметры подобраны замером: f=400:g=3 даёт лишь 8.6 процента вместо 43.2,
# так что менять их наугад не стоит.
NORMALIZE_FILTER = "dynaudnorm=f=200:g=5"


def to_wav16k(src: Path, dst: Path, ffmpeg: str, normalize: bool = True) -> float:
    """Приводит любой аудиофайл к 16 кГц моно WAV. Возвращает длительность."""
    command = [ffmpeg, "-nostdin", "-v", "error", "-y", "-i", str(src), "-vn"]
    if normalize:
        command += ["-af", NORMALIZE_FILTER]
    command += ["-ar", str(SAMPLE_RATE), "-ac", "1", "-c:a", "pcm_s16le", str(dst)]
    subprocess.run(command, check=True, capture_output=True)
    with wave.open(str(dst)) as w:
        return w.getnframes() / w.getframerate()


def read_wav(path: Path) -> np.ndarray:
    with wave.open(str(path)) as w:
        if w.getframerate() != SAMPLE_RATE or w.getnchannels() != 1:
            raise ValueError(f"ожидался {SAMPLE_RATE} Гц моно, получено "
                             f"{w.getframerate()} Гц / {w.getnchannels()} кан.")
        raw = w.readframes(w.getnframes())
    return np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0


def gap_ranges(turns: list[dict], duration: float):
    """Промежутки между репликами, куда распознавание не заглядывало."""
    previous = 0.0
    for turn in turns:
        if turn["start"] - previous >= GAP_MIN_SEC:
            yield previous, turn["start"]
        previous = max(previous, turn["end"])
    if duration - previous >= GAP_MIN_SEC:
        yield previous, duration


class Pipeline:
    """Диаризация плюс распознавание. Потокобезопасен на уровне одного воркера."""

    def __init__(self, models_dir: Path, threads: int, replacements_path: Path,
                 base_dir: Path):
        self.models_dir = Path(models_dir)
        self.threads = threads
        self.replacements_path = Path(replacements_path)
        self.base_dir = Path(base_dir)
        self._ffmpeg: str | None = None
        self._asr = None
        self._diar_cache: dict[int, object] = {}
        self._embedder = None
        self._replacements: dict[str, str] = {}

    @property
    def ffmpeg(self) -> str:
        """Ищется при первом обращении, а не в конструкторе.

        Иначе отсутствие ffmpeg роняло бы сам импорт модуля - и вместо внятного
        сообщения пользователь получал бы стек при старте, а обновление кода
        не могло бы проверить работоспособность новой версии.
        """
        if self._ffmpeg is None:
            self._ffmpeg = find_ffmpeg(self.base_dir)
        return self._ffmpeg

    def warmup(self) -> None:
        """Грузит модели заранее, чтобы первый запрос не ждал."""
        import onnx_asr

        missing = [str(p) for p in (self.models_dir / SEG_MODEL_REL,
                                    self.models_dir / EMB_MODEL_REL,
                                    self.models_dir / ASR_MODEL_REL / "config.json")
                   if not p.is_file()]
        if missing:
            raise ModelsMissing(
                f"не найдены модели: {', '.join(missing)}. Запустите download_models.bat")

        t0 = time.time()
        # CoreML и DirectML на GigaAM падают на отдельных узлах графа, поэтому явно CPU:
        # на 16-ядерном процессоре этого более чем достаточно.
        self._asr = onnx_asr.load_model(
            ASR_MODEL_NAME, str(self.models_dir / ASR_MODEL_REL),
            providers=["CPUExecutionProvider"],
            sess_options=self._session_options(),
        )
        self._reload_replacements()
        log.info("модели загружены за %.1f с, потоков: %d", time.time() - t0, self.threads)

    def _session_options(self):
        import onnxruntime as ort

        opts = ort.SessionOptions()
        opts.intra_op_num_threads = self.threads
        opts.inter_op_num_threads = 1
        return opts

    def _reload_replacements(self) -> None:
        self._replacements = load_replacements(self.replacements_path)

    def _diarizer(self, num_speakers: int):
        """Диаризатор кэшируется по числу говорящих - его создание не бесплатно."""
        import sherpa_onnx

        if num_speakers in self._diar_cache:
            return self._diar_cache[num_speakers]

        clustering = (sherpa_onnx.FastClusteringConfig(num_clusters=num_speakers)
                      if num_speakers > 0 else sherpa_onnx.FastClusteringConfig(threshold=0.5))
        cfg = sherpa_onnx.OfflineSpeakerDiarizationConfig(
            segmentation=sherpa_onnx.OfflineSpeakerSegmentationModelConfig(
                pyannote=sherpa_onnx.OfflineSpeakerSegmentationPyannoteModelConfig(
                    model=str(self.models_dir / SEG_MODEL_REL)),
                num_threads=self.threads),
            embedding=sherpa_onnx.SpeakerEmbeddingExtractorConfig(
                model=str(self.models_dir / EMB_MODEL_REL), num_threads=self.threads),
            clustering=clustering,
            # Подобрано на реальных звонках: отсекает обрывки на перекрывающейся
            # речи, из-за которых в расшифровку лезли огрызки чужих слов.
            min_duration_on=1.0,
            min_duration_off=0.7,
        )
        diar = sherpa_onnx.OfflineSpeakerDiarization(cfg)
        self._diar_cache[num_speakers] = diar
        return diar

    def _separation_quality(self, samples: np.ndarray, segments: list[Segment]) -> float:
        """Оценивает, действительно ли голоса разошлись, или разметка случайна."""
        import sherpa_onnx

        usable = [s for s in segments if s.end - s.start >= 1.0][:120]
        if len(usable) < 4:
            return 0.0
        if self._embedder is None:
            self._embedder = sherpa_onnx.SpeakerEmbeddingExtractor(
                sherpa_onnx.SpeakerEmbeddingExtractorConfig(
                    model=str(self.models_dir / EMB_MODEL_REL), num_threads=self.threads))
        vectors, labels = [], []
        for seg in usable:
            # Берём кусок из середины реплики: там речь устойчивее, чем на краях,
            # а длинный фрагмент модель отпечатков просто не переваривает.
            middle = (seg.start + seg.end) / 2
            half = min(EMBED_SEC, seg.end - seg.start) / 2
            piece = samples[int((middle - half) * SAMPLE_RATE):int((middle + half) * SAMPLE_RATE)]
            if len(piece) < SAMPLE_RATE // 2:
                continue
            stream = self._embedder.create_stream()
            stream.accept_waveform(SAMPLE_RATE, piece)
            stream.input_finished()
            vectors.append(np.array(self._embedder.compute(stream)))
            labels.append(seg.speaker)
        if len(vectors) < 4:
            return 0.0
        return separation_quality(np.array(vectors), np.array(labels))

    def _recognize_safely(self, audio: np.ndarray, depth: int = 0) -> str:
        """Распознаёт кусок, при ошибке деля его пополам.

        Пределы длины у разных сборок onnxruntime разные и заранее неизвестны,
        поэтому вместо подбора константы просто уменьшаем кусок и пробуем снова.
        """
        try:
            return self._asr.recognize(audio, sample_rate=SAMPLE_RATE)
        except Exception:
            seconds = len(audio) / SAMPLE_RATE
            if seconds <= MIN_SPLIT_SEC or depth >= 5:
                log.warning("не удалось распознать кусок в %.1f с, пропускаю", seconds)
                return ""
            log.info("кусок в %.1f с не прошёл, делю пополам", seconds)
            middle = len(audio) // 2
            left = self._recognize_safely(audio[:middle], depth + 1)
            right = self._recognize_safely(audio[middle:], depth + 1)
            return " ".join(part for part in (left, right) if part)

    def _recover_gaps(self, samples: np.ndarray, turns: list[dict],
                      duration: float) -> list[dict]:
        """Распознаёт то, что детектор речи не отметил как речь.

        Диаризация задаёт не только «кто говорит», но и границы того, что
        вообще попадает в распознавание. Модель pyannote обучена разделять
        говорящих, а не искать речь, и на тихой записи со стола часть фраз
        в эти границы не попадает - они пропадают молча. Замер на реальной
        записи: 4 промежутка из 35 содержали речь, около 2 % слов.

        Кто произнёс найденное, неизвестно: диаризация этих участков не
        видела. Поэтому speaker = 0, а роль потом расставит LLM по смыслу.
        """
        out: list[dict] = []
        for start, end in gap_ranges(turns, duration):
            pieces = []
            for chunk_start, chunk_stop in chunk_ranges(start, end, MAX_CHUNK_SEC):
                audio = samples[int(chunk_start * SAMPLE_RATE):int(chunk_stop * SAMPLE_RATE)]
                if len(audio) < SAMPLE_RATE * 0.2:
                    continue
                pieces.append(self._recognize_safely(audio))
            text = " ".join(p.strip() for p in pieces if p and p.strip())
            if len(text) < MIN_RECOVERED_CHARS:
                continue
            text = normalize_typography(apply_replacements(text, self._replacements))
            audio = samples[int(start * SAMPLE_RATE):int(end * SAMPLE_RATE)]
            out.append({
                "speaker": 0,
                "start": round(start, 2),
                "end": round(end, 2),
                "text": text,
                "acoustics": segment_acoustics(audio),
                # Признак второго прохода: говорящий неизвестен, и потребителю
                # стоит знать, что это не результат диаризации.
                "recovered": True,
            })
        return out

    def transcribe(self, wav_path: Path, num_speakers: int = 2,
                   recover_gaps: bool = True) -> dict:
        """Полный проход: диаризация, распознавание реплик, постобработка.

        Модели грузятся при первой задаче, а не при создании: воркеров несколько,
        и держать копии моделей у тех, кому работа ещё не досталась, незачем.
        """
        if self._asr is None:
            self.warmup()
        self._reload_replacements()

        samples = read_wav(wav_path)
        duration = len(samples) / SAMPLE_RATE

        t0 = time.time()
        raw = self._diarizer(num_speakers).process(samples).sort_by_start_time()
        segments = [Segment(start=s.start, end=s.end, speaker=s.speaker) for s in raw]
        try:
            quality = self._separation_quality(samples, segments)
        except Exception as exc:  # noqa: BLE001 - оценка вспомогательная
            # Метрика не должна ронять задачу: без неё расшифровка всё равно нужна.
            log.warning("не удалось оценить разделение говорящих: %s", exc)
            quality = 0.0
        t_diar = time.time() - t0

        t0 = time.time()
        turns_out = []
        for turn in merge_turns(segments, gap=0.8):
            pieces = []
            for start, stop in chunk_ranges(turn.start, turn.end, MAX_CHUNK_SEC):
                audio = samples[int(start * SAMPLE_RATE):int(stop * SAMPLE_RATE)]
                if len(audio) < SAMPLE_RATE * 0.2:
                    continue
                pieces.append(self._recognize_safely(audio))
            text = " ".join(p.strip() for p in pieces if p and p.strip())
            if not text:
                continue
            text = normalize_typography(apply_replacements(text, self._replacements))
            audio = samples[int(turn.start * SAMPLE_RATE):int(turn.end * SAMPLE_RATE)]
            turns_out.append({
                "speaker": turn.speaker + 1,
                "start": round(turn.start, 2),
                "end": round(turn.end, 2),
                "text": text,
                # Подсказка для LLM: у говорящего ближе к микрофону громкость
                # и доля высоких выше. Когда разделение по голосу ненадёжно,
                # это единственный объективный признак, кто есть кто.
                "acoustics": segment_acoustics(audio),
            })
        t_asr = time.time() - t0

        t0 = time.time()
        recovered = (self._recover_gaps(samples, turns_out, duration)
                     if recover_gaps else [])
        t_recovery = time.time() - t0
        if recovered:
            turns_out = sorted(turns_out + recovered, key=lambda t: t["start"])
            log.info("во втором проходе найдено реплик: %d", len(recovered))

        return {
            "duration_sec": round(duration, 1),
            "turns": turns_out,
            # Готовый материал для разметки ролей на стороне Talkscore: делить
            # реплики самостоятельно ей незачем, пунктуацию ставит ASR.
            "sentences": split_sentences(turns_out),
            "stats": {
                **speech_stats(segments),
                "separation_quality": quality,
                "speakers_reliable": quality >= RELIABLE_SEPARATION,
                "silence_sec": round(max(0.0, duration - sum(s.end - s.start for s in segments)), 1),
                "turns_count": len(turns_out),
                "recovered_turns": len(recovered),
                "recovered_sec": round(sum(t["end"] - t["start"] for t in recovered), 1),
            },
            "timing": {
                "diarization_sec": round(t_diar, 1),
                "asr_sec": round(t_asr, 1),
                "recovery_sec": round(t_recovery, 1),
                "realtime_factor": round(duration / max(t_diar + t_asr, 0.001), 1),
            },
        }