"""
Transcription + diarization using WhisperX as a Python library.
"""
import gc
import json
from pathlib import Path

import torch
import whisperx

from ..errors import TranscriptionError, DiarizationError
from ..logging_config import get_logger

log = get_logger(__name__)


def transcribe_and_diarize(
    audio_path: Path,
    output_json: Path,
    model_size: str,
    compute_type: str,
    hf_token: str,
    progress_cb=None,
) -> dict:
    device = "cuda" if torch.cuda.is_available() else "cpu"
    if device == "cpu":
        log.warning("No CUDA GPU available - running on CPU, this will be much slower")

    def report(msg):
        log.info(msg)
        if progress_cb:
            progress_cb(msg)

    report("Loading transcription model...")
    try:
        model = whisperx.load_model(model_size, device, compute_type=compute_type)
    except RuntimeError as e:
        if "out of memory" in str(e).lower():
            raise TranscriptionError(
                f"Ran out of GPU memory loading the '{model_size}' model. "
                f"Try a smaller model size in Settings (e.g. 'small' or 'medium'), "
                f"or check nothing else is using the GPU right now.", cause=e
            )
        raise TranscriptionError(f"Failed to load Whisper model '{model_size}': {e}", cause=e)
    except Exception as e:
        raise TranscriptionError(f"Failed to load Whisper model '{model_size}': {e}", cause=e)

    try:
        report("Loading audio...")
        audio = whisperx.load_audio(str(audio_path))

        report("Transcribing...")
        result = model.transcribe(audio, batch_size=16)
    except RuntimeError as e:
        if "out of memory" in str(e).lower():
            raise TranscriptionError(
                f"Ran out of GPU memory during transcription. Try a smaller Whisper model size in Settings.", cause=e
            )
        raise TranscriptionError(f"Transcription failed: {e}", cause=e)
    except Exception as e:
        raise TranscriptionError(f"Transcription failed: {e}", cause=e)
    finally:
        del model
        gc.collect()
        if device == "cuda":
            torch.cuda.empty_cache()

    try:
        report("Aligning...")
        align_model, metadata = whisperx.load_align_model(language_code=result["language"], device=device)
        result = whisperx.align(result["segments"], align_model, metadata, audio, device, return_char_alignments=False)
        del align_model
        gc.collect()
        if device == "cuda":
            torch.cuda.empty_cache()
    except Exception as e:
        raise TranscriptionError(f"Alignment step failed: {e}", cause=e)

    if hf_token:
        try:
            report("Diarizing...")
            diarize_model = whisperx.diarize.DiarizationPipeline(use_auth_token=hf_token, device=device)
            diarize_segments = diarize_model(audio)
            result = whisperx.assign_word_speakers(diarize_segments, result)
        except Exception as e:
            msg = str(e)
            if "gated" in msg.lower() or "403" in msg or "access" in msg.lower():
                raise DiarizationError(
                    "HuggingFace rejected access to the diarization model. Make sure you've accepted the "
                    "terms on the gated model page (linked in the error below) using the SAME account "
                    f"your HF token belongs to.\n{msg}", cause=e
                )
            raise DiarizationError(f"Diarization failed: {msg}", cause=e)
    else:
        report("No HF token set - skipping speaker diarization (all speech will be unattributed)")

    try:
        output_json.parent.mkdir(parents=True, exist_ok=True)
        output_json.write_text(json.dumps(result, ensure_ascii=False), encoding='utf-8')
    except OSError as e:
        raise TranscriptionError(f"Transcribed successfully, but failed to write output to {output_json}: {e}", cause=e)

    report("Done")
    return result