from __future__ import annotations

import logging
import threading
import wave
from pathlib import Path

import numpy as np

from app.config import get_settings
from app.errors import JobCancelled, VaaniError
from app.languages import PROMPTS
from app.transcript import chunks_from_segments, drop_prompt, words_from_segments

log = logging.getLogger("vaani.transcribe")

_model = None
_loaded_name: str | None = None
_lock = threading.Lock()


def whisper_installed() -> bool:
    try:
        import faster_whisper  # noqa: F401
    except ImportError:
        return False
    return True


def get_model():
    global _model, _loaded_name
    settings = get_settings()
    with _lock:
        if _model is not None and _loaded_name == settings.whisper_model:
            return _model
        _model = _load(settings.whisper_model, settings.whisper_device, settings.whisper_compute)
        _loaded_name = settings.whisper_model
        return _model


def _load(name: str, device: str, compute: str):
    try:
        from faster_whisper import WhisperModel
    except ImportError as exc:
        raise VaaniError(
            "The speech engine is not installed. Run the setup in the README again.",
            "tool_missing",
            500,
        ) from exc
    settings = get_settings()
    settings.model_dir.mkdir(parents=True, exist_ok=True)
    try:
        return WhisperModel(
            name,
            device=device,
            compute_type=compute,
            download_root=str(settings.model_dir),
        )
    except Exception:
        if device != "cpu":
            log.exception("speech model failed on %s, trying cpu", device)
            return WhisperModel(
                name,
                device="cpu",
                compute_type="int8",
                download_root=str(settings.model_dir),
            )
        raise


def transcribe_wav(wav: Path, language: str | None, on_segment, cancel) -> tuple[list[dict], str]:
    return _run(wav, language, "transcribe", on_segment, cancel, as_words=False)


def transcribe_words(wav: Path, language: str | None, on_segment, cancel) -> tuple[list[dict], str]:
    return _run(wav, language, "transcribe", on_segment, cancel, as_words=True)


def translate_wav(wav: Path, language: str | None, cancel) -> list[dict]:
    chunks, _detected = _run(wav, language, "translate", None, cancel, as_words=False)
    return chunks


def _run(wav: Path, language: str | None, task: str, on_segment, cancel, as_words: bool) -> tuple[list[dict], str]:
    if cancel():
        raise JobCancelled()
    model = get_model()
    pieces = []
    offset = 0.0
    segment_count = 0

    def report(index: int, detected: str) -> None:
        nonlocal segment_count
        if index:
            segment_count += 1
        if on_segment and (index or not segment_count):
            on_segment(segment_count, detected)

    # FFmpeg prepares mono, 16-bit PCM at 16 kHz for this pipeline.
    with wave.open(str(wav), "rb") as handle:
        rate = handle.getframerate()
        if language is None:
            language = _detect_language(model, handle, cancel)
        while raw := handle.readframes(rate * 300):
            if cancel():
                raise JobCancelled()
            audio = np.frombuffer(raw, dtype="<i2").astype(np.float32) / 32768.0
            chunk, _ = _run_chunk(model, audio, language, task, report, cancel, as_words)
            for piece in chunk:
                piece["start"] += offset
                piece["end"] += offset
            pieces.extend(chunk)
            offset += audio.size / rate
    return pieces, language or "en"


def _detect_language(model, handle, cancel) -> str:
    rate = handle.getframerate()
    frames = handle.getnframes()
    window = rate * 30
    last_start = frames - window
    starts = (0, last_start // 2, last_start) if frames >= window * 3 else (0,)
    samples = []
    for start in starts:
        if cancel():
            raise JobCancelled()
        handle.setpos(start)
        samples.append(handle.readframes(window if len(starts) > 1 else frames))
    handle.rewind()
    audio = np.frombuffer(b"".join(samples), dtype="<i2").astype(np.float32) / 32768.0
    detected, _, _ = model.detect_language(
        audio=audio, vad_filter=True, language_detection_segments=3,
    )
    return detected


def _run_chunk(model, audio, language, task, on_segment, cancel, as_words):
    prompt = PROMPTS.get(language or "") if task == "transcribe" else None
    attempts = [(True, True), (True, False), (False, False)]
    last: Exception | None = None
    for use_vad, word_times in attempts:
        try:
            return _decode(
                model,
                audio,
                language,
                task,
                prompt,
                use_vad,
                word_times,
                on_segment,
                cancel,
                as_words,
            )
        except JobCancelled:
            raise
        except Exception as exc:
            last = exc
            text = str(exc).lower()
            if any(token in text for token in ("vad", "onnx", "word timestamp", "word_timestamps")):
                log.warning("retrying transcription without %s", "vad" if "vad" in text or "onnx" in text else "word times")
                continue
            break
    raise VaaniError(
        "The transcript could not be written.",
        "transcript_failed",
        500,
        "retry",
    ) from last


def _decode(model, audio, language, task, prompt, use_vad, word_times, on_segment, cancel, as_words: bool):
    # Hindi listening must not carry a bad line into the next window. That loop
    # was writing the same mark again and again on mixed speech.
    hold_context = language != "hi"
    segments, info = model.transcribe(
        audio,
        language=language,
        task=task,
        vad_filter=use_vad,
        word_timestamps=word_times,
        beam_size=5,
        initial_prompt=prompt,
        condition_on_previous_text=hold_context,
        hallucination_silence_threshold=None if hold_context else 2.0,
    )
    detected = getattr(info, "language", None) or language or "en"
    if on_segment:
        on_segment(0, detected)
    pieces: list[dict] = []
    for index, segment in enumerate(segments, start=1):
        if cancel():
            raise JobCancelled()
        if on_segment:
            on_segment(index, detected)
        if as_words:
            pieces.extend(words_from_segments([segment]))
        else:
            pieces.extend(chunks_from_segments([segment]))
    return drop_prompt(pieces, prompt), detected