"""
Audio parser — whisperx pipeline with speaker diarization.

Strategy:
1. Load audio with whisperx.load_audio() (handles mp3/wav/m4a via ffmpeg)
2. Transcribe with faster-whisper via whisperx model
3. Align word-level timestamps with whisperx.align()
4. Diarize speakers with pyannote.audio via whisperx.DiarizationPipeline
5. Assign speaker labels with whisperx.assign_word_speakers()
6. Yield one text:transcript_segment per speaker segment

Falls back gracefully at each step:
- No alignment model for language → skip alignment, use segment timestamps
- No HF_TOKEN / pyannote unavailable → skip diarization, emit segments without speaker
- ffmpeg missing → raise ImportError with install hint

Environment variables:
  HF_TOKEN        Hugging Face token for pyannote.audio (required for diarization)
  WHISPER_MODEL   Model size: tiny/base/small/medium/large-v2 (default: base)
  WHISPER_DEVICE  cpu/cuda/mps (default: auto-detect)
  WHISPER_COMPUTE float16/int8/float32 (default: int8 on cpu, float16 on gpu)

Output elements:
  text:transcript_segment with TranscriptData (speaker, start_time_s, end_time_s,
  word_timestamps). section_path = [filename_stem]. One element per segment.

Stats note: token count reflects transcript word count, not model tokens.
"""

from __future__ import annotations

import os
from collections.abc import Iterator
from pathlib import Path

from .base import BaseParser, FileRecord, RawElement

_SUPPORTED_MIMES = {"audio/mpeg", "audio/wav", "audio/mp4"}

# Whisper model to use — balance speed vs accuracy
_DEFAULT_MODEL = "base"
_DEFAULT_BATCH_SIZE = 16


def _device_and_compute() -> tuple[str, str]:
    """Auto-detect best device and compute type."""
    device = os.environ.get("WHISPER_DEVICE", "").lower()
    compute = os.environ.get("WHISPER_COMPUTE", "").lower()

    if not device:
        try:
            import torch  # type: ignore[import-untyped]

            if torch.cuda.is_available():
                device = "cuda"
            elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
                device = "mps"
            else:
                device = "cpu"
        except ImportError:
            device = "cpu"

    if not compute:
        compute = "int8" if device == "cpu" else "float16"

    return device, compute


def _load_whisperx():  # type: ignore[no-untyped-def]
    """Import whisperx; raise ImportError with install hint if missing."""
    try:
        import whisperx  # type: ignore[import-untyped]

        return whisperx
    except ImportError as e:
        raise ImportError(
            "whisperx required for audio parsing: "
            "pip install 'extractor[audio]'  # or: pip install whisperx faster-whisper"
        ) from e


def _transcribe(path: Path) -> tuple[dict, object, str]:
    """
    Run whisperx transcription on path.
    Returns (result_dict, whisperx_module, detected_language).
    result_dict has 'segments' and 'language'.
    """
    wx = _load_whisperx()
    device, compute = _device_and_compute()
    model_name = os.environ.get("WHISPER_MODEL", _DEFAULT_MODEL)

    audio = wx.load_audio(str(path))
    model = wx.load_model(model_name, device, compute_type=compute)
    result = model.transcribe(audio, batch_size=_DEFAULT_BATCH_SIZE)
    language = result.get("language", "en")

    # Word-level alignment
    try:
        align_model, metadata = wx.load_align_model(language_code=language, device=device)
        result = wx.align(
            result["segments"],
            align_model,
            metadata,
            audio,
            device,
            return_char_alignments=False,
        )
    except Exception:
        pass  # Alignment unavailable for this language — use segment-level timestamps

    return result, wx, language


def _diarize(path: Path, wx, result: dict, language: str) -> dict:
    """
    Run speaker diarization and assign speakers to words.
    Returns updated result dict (may be unchanged if diarization unavailable).
    """
    hf_token = os.environ.get("HF_TOKEN", "")
    if not hf_token:
        return result  # No token — skip diarization

    try:
        device, _ = _device_and_compute()
        audio = wx.load_audio(str(path))
        diarize_model = wx.DiarizationPipeline(use_auth_token=hf_token, device=device)
        diarize_segments = diarize_model(audio)
        result = wx.assign_word_speakers(diarize_segments, result)
    except Exception:
        pass  # Diarization failed — continue without speaker labels

    return result


def _segments_to_elements(
    result: dict,
    source_stem: str,
) -> Iterator[RawElement]:
    """
    Convert whisperx result segments to text:transcript_segment RawElements.

    Merges consecutive same-speaker segments into a single element up to
    ~30 seconds to keep chunk sizes reasonable. If no speakers are assigned,
    emits each segment individually.
    """
    segments = result.get("segments", [])
    if not segments:
        return

    has_speakers = any(seg.get("speaker") for seg in segments)

    if has_speakers:
        # Merge consecutive same-speaker segments (max 30s window)
        yield from _merge_speaker_segments(segments)
    else:
        # No diarization — emit each segment as-is
        for seg in segments:
            text = seg.get("text", "").strip()
            if not text:
                continue
            start = float(seg.get("start", 0.0))
            end = float(seg.get("end", start))
            words = seg.get("words", [])
            yield RawElement(
                element_type="text:transcript_segment",
                text=text,
                transcript_data={
                    "speaker": None,
                    "start_time_s": round(start, 3),
                    "end_time_s": round(end, 3),
                    "word_timestamps": _clean_words(words) if words else None,
                },
                extra={
                    "extraction_method": "whisperx:transcribe",
                    "confidence": _avg_word_confidence(words),
                    "section_path": [source_stem],
                },
            )


def _merge_speaker_segments(segments: list[dict]) -> Iterator[RawElement]:
    """Merge consecutive same-speaker segments into transcript elements."""
    _MAX_WINDOW_S = 30.0

    buffer_text: list[str] = []
    buffer_words: list[dict] = []
    current_speaker: str | None = None
    seg_start: float = 0.0
    seg_end: float = 0.0

    def flush() -> RawElement | None:
        if not buffer_text:
            return None
        return RawElement(
            element_type="text:transcript_segment",
            text=" ".join(buffer_text).strip(),
            transcript_data={
                "speaker": current_speaker,
                "start_time_s": round(seg_start, 3),
                "end_time_s": round(seg_end, 3),
                "word_timestamps": buffer_words if buffer_words else None,
            },
            extra={
                "extraction_method": "whisperx:diarize",
                "confidence": _avg_word_confidence(buffer_words),
            },
        )

    for seg in segments:
        text = seg.get("text", "").strip()
        if not text:
            continue
        speaker = seg.get("speaker") or None
        start = float(seg.get("start", 0.0))
        end = float(seg.get("end", start))
        words = seg.get("words", [])

        # Flush on speaker change or window overflow
        if buffer_text and (speaker != current_speaker or (end - seg_start) > _MAX_WINDOW_S):
            el = flush()
            if el:
                yield el
            buffer_text = []
            buffer_words = []
            seg_start = start

        if not buffer_text:
            seg_start = start
            current_speaker = speaker

        buffer_text.append(text)
        buffer_words.extend(_clean_words(words))
        seg_end = end

    el = flush()
    if el:
        yield el


def _clean_words(words: list[dict]) -> list[dict]:
    """Normalize word timestamp dicts to minimal schema."""
    cleaned = []
    for w in words:
        entry: dict = {"word": w.get("word", "").strip()}
        if "start" in w:
            entry["start"] = round(float(w["start"]), 3)
        if "end" in w:
            entry["end"] = round(float(w["end"]), 3)
        if "score" in w:
            entry["confidence"] = round(float(w["score"]), 4)
        if entry["word"]:
            cleaned.append(entry)
    return cleaned


def _avg_word_confidence(words: list[dict]) -> float:
    """Compute average word confidence from word list."""
    scores = [float(w.get("score", 1.0)) for w in words if "score" in w]
    return round(sum(scores) / len(scores), 4) if scores else 1.0


class AudioParser(BaseParser):
    handles = [
        ("audio/mpeg", None),
        ("audio/wav", None),
        ("audio/mp4", None),
    ]

    def parse(self, file_record: FileRecord) -> Iterator[RawElement]:
        # Emit document title element from filename
        stem = file_record.path.stem
        yield RawElement(
            element_type="meta:document_title",
            text=stem,
            extra={"extraction_method": "whisperx:filename", "confidence": 1.0},
        )

        result, wx, language = _transcribe(file_record.path)
        result = _diarize(file_record.path, wx, result, language)

        yield from _segments_to_elements(result, stem)
