from typing import List, Dict
from faster_whisper import WhisperModel
import os
import torch
from pyannote.audio import Pipeline

# Load environment variables
WHISPER_MODEL_SIZE = os.getenv("WHISPER_MODEL_SIZE", "base")
HF_TOKEN = os.getenv("HF_TOKEN")  # For pyannote

# Initialize Whisper model
try:
    _whisper = WhisperModel(WHISPER_MODEL_SIZE, device="cpu", compute_type="int8")
except Exception as e:
    print(f"Error loading Whisper model: {e}")
    _whisper = WhisperModel("base", device="cpu", compute_type="int8")

# Initialize pyannote diarization pipeline
if HF_TOKEN:
    try:
        diarization_pipeline = Pipeline.from_pretrained(
            "pyannote/speaker-diarization-3.1", token=HF_TOKEN
        )
        # Move pipeline to CPU if no GPU is available
        if not torch.cuda.is_available():
            diarization_pipeline = diarization_pipeline.to(torch.device("cpu"))
    except Exception as e:
        print(f"Error loading pyannote pipeline: {e}")
        diarization_pipeline = None
else:
    print("HF_TOKEN not set, skipping diarization.")
    diarization_pipeline = None


def transcribe(audio_path: str) -> List[Dict]:
    """
    Transcribe an audio file using faster-whisper and combine with
    pyannote.audio for speaker diarization.
    """
    # 1. Transcribe with Whisper
    segments, _ = _whisper.transcribe(
        audio_path, language="en", vad_filter=True, beam_size=1
    )

    whisper_segments = []
    for seg in segments:
        whisper_segments.append(
            {"start": float(seg.start), "end": float(seg.end), "text": seg.text.strip()}
        )

    if not diarization_pipeline:
        # If diarization is not available, return with a single speaker
        for seg in whisper_segments:
            seg["speaker"] = "A"
        return whisper_segments

    # 2. Perform Diarization
    try:
        diarization = diarization_pipeline(audio_path)

        # --- DEBUGGING LINES (optional) ---
        print("\n--- Diarization Output ---")
        print(f"Type of diarization object: {type(diarization)}")
        print("Diarization object content:")
        print(diarization)
        print("--- End Diarization Output ---\n")
        # ----------------------------------

    except Exception as e:
        print(f"Error during diarization: {e}")
        for seg in whisper_segments:
            seg["speaker"] = "A"
        return whisper_segments

    # 3. Assign Speaker to Segments
    out_segments = []

    # For pyannote 4.x, use diarization.speaker_diarization
    annotation = diarization.speaker_diarization  # This is an Annotation object

    for seg in whisper_segments:
        midpoint = seg["start"] + (seg["end"] - seg["start"]) / 2
        speaker = "UNKNOWN"
        for turn, _, speaker_label in annotation.itertracks(yield_label=True):
            if turn.start <= midpoint <= turn.end:
                speaker = speaker_label
                break

        out_segments.append(
            {
                "start": seg["start"],
                "end": seg["end"],
                "speaker": speaker,
                "text": seg["text"],
            }
        )

    if not out_segments:
        return [{"start": 0.0, "end": 0.0, "speaker": "A", "text": ""}]
    return out_segments