"""Sarvam AI batch speech-to-text with diarization."""

from __future__ import annotations

import glob
import json
import os
import tempfile
from dataclasses import dataclass

from sarvamai import SarvamAI


@dataclass
class Segment:
    start: float
    end: float
    speaker_id: str
    text: str


@dataclass
class Transcription:
    language_code: str | None
    segments: list[Segment]
    full_transcript: str


def transcribe_with_diarization(
    audio_path: str,
    api_key: str,
    num_speakers: int = 2,
) -> Transcription:
    """Run a Sarvam batch STT job with diarization and return parsed segments."""
    client = SarvamAI(api_subscription_key=api_key)

    job = client.speech_to_text_job.create_job(
        model="saaras:v3",
        mode="verbatim",
        language_code="unknown",
        with_diarization=True,
        num_speakers=num_speakers,
    )

    job.upload_files(file_paths=[audio_path])
    job.start()
    # Default timeout is 600s; raise it so longer recordings don't get cut off.
    job.wait_until_complete(poll_interval=5, timeout=3600)

    file_results = job.get_file_results()
    if not file_results.get("successful"):
        failed = file_results.get("failed", [])
        msg = failed[0].get("error_message") if failed else "unknown error"
        raise RuntimeError(f"Sarvam transcription failed: {msg}")

    with tempfile.TemporaryDirectory() as out_dir:
        job.download_outputs(output_dir=out_dir)
        json_files = sorted(glob.glob(os.path.join(out_dir, "*.json")))
        if not json_files:
            raise RuntimeError("No output JSON returned by Sarvam.")
        with open(json_files[0], "r", encoding="utf-8") as fh:
            data = json.load(fh)

    return _parse(data)


def _parse(data: dict) -> Transcription:
    language_code = data.get("language_code")
    full_transcript = data.get("transcript", "") or ""

    segments: list[Segment] = []
    diarized = data.get("diarized_transcript") or {}
    for entry in diarized.get("entries", []) or []:
        text = (entry.get("transcript") or "").strip()
        if not text:
            continue
        segments.append(
            Segment(
                start=float(entry.get("start_time_seconds", 0.0)),
                end=float(entry.get("end_time_seconds", 0.0)),
                speaker_id=str(entry.get("speaker_id", "0")),
                text=text,
            )
        )

    # Fallback: no diarization entries -> single segment from the full transcript.
    if not segments and full_transcript:
        segments.append(Segment(0.0, 0.0, "0", full_transcript))

    segments.sort(key=lambda s: s.start)
    return Transcription(language_code, segments, full_transcript)