"""WhisperX transcription with speaker diarization."""

import os
import tempfile
from pathlib import Path
from typing import Any, Optional

import numpy as np
import soundfile as sf
import torch
import torch.serialization
import whisperx
from simple_diarizer.diarizer import Diarizer

# Add safe globals for PyTorch 2.8+ weights loading
torch.serialization.add_safe_globals([
    "omegaconf.listconfig.ListConfig",
    "omegaconf.dictconfig.DictConfig"
])


class WhisperXTranscriber:
    """Handles audio transcription and speaker diarization using WhisperX."""

    def __init__(self, model_dir: Path, model_size: str = "base"):
        """Initialize WhisperX with local models."""
        self.model_dir = model_dir
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.compute_type = "float16" if self.device == "cuda" else "int8"
        self.is_loaded = False
        self.diarizer = None

        # Delay loading until first use to avoid startup issues
        self.model = None
        self.align_model = None
        self.align_metadata = None
        self.model_size = model_size

        try:
            # Load simple_diarizer (no authentication required!)
            self.diarizer = Diarizer(
                embed_model="xvec",
                cluster_method="sc",
            )
            print("Speaker diarization enabled (using simple_diarizer)")
        except Exception as e:
            print(f"Warning: Failed to load diarizer: {e}")
            self.diarizer = None

    def _ensure_models_loaded(self):
        """Load models on first use."""
        if self.model is None:
            try:
                print("Loading WhisperX models...")
                # Load ASR model
                self.model = whisperx.load_model(
                    self.model_size,
                    self.device,
                    compute_type=self.compute_type,
                    download_root=str(self.model_dir / "whisper"),
                )

                # Load alignment model
                self.align_model, self.align_metadata = whisperx.load_align_model(
                    language_code="en",
                    device=self.device,
                    model_dir=str(self.model_dir / "alignment"),
                )
                self.is_loaded = True
                print("✓ WhisperX models loaded successfully")
            except Exception as e:
                print(f"Failed to load WhisperX models: {e}")
                self.is_loaded = False
                raise

    async def transcribe_chunk(self, audio_chunk: np.ndarray) -> Optional[dict[str, Any]]:
        """Transcribe an audio chunk with speaker diarization."""
        # Load models on first use
        if not self.is_loaded:
            try:
                self._ensure_models_loaded()
            except Exception:
                return None

        try:
            # Transcribe
            result = self.model.transcribe(
                audio_chunk,
                batch_size=16,
            )

            # Align whisper output
            result = whisperx.align(
                result["segments"],
                self.align_model,
                self.align_metadata,
                audio_chunk,
                self.device,
            )

            # Diarize using simple_diarizer
            if self.diarizer:
                result = self._apply_diarization(audio_chunk, result)

            # Format output
            formatted_result = []
            for segment in result["segments"]:
                formatted_result.append({
                    "start": segment["start"],
                    "end": segment["end"],
                    "text": segment["text"],
                    "speaker": segment.get("speaker", "Unknown"),
                })

            return {"segments": formatted_result}

        except Exception as e:
            print(f"Transcription error: {e}")
            return None

    def _apply_diarization(
        self, audio_chunk: np.ndarray, transcript_result: dict
    ) -> dict:
        """Apply speaker diarization to transcript segments."""
        try:
            # simple_diarizer needs a file path, so write to temp file
            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
                temp_path = f.name
                sf.write(temp_path, audio_chunk, 16000)

            # Run diarization
            diarization = self.diarizer.diarize(temp_path, num_speakers=None)

            # Clean up temp file
            os.unlink(temp_path)

            # Assign speakers to segments based on time overlap
            for segment in transcript_result["segments"]:
                segment_mid = (segment["start"] + segment["end"]) / 2
                speaker = "Unknown"

                for diar_seg in diarization:
                    if diar_seg["start"] <= segment_mid <= diar_seg["end"]:
                        speaker = f"Speaker {diar_seg['label']}"
                        break

                segment["speaker"] = speaker

        except Exception as e:
            print(f"Diarization error: {e}")

        return transcript_result
