annotator.py

"""
Annotation writer — merges pipeline outputs into structured files.

Produces three output formats:
    1. **Per-file JSON** — full segment-level annotations.
    2. **Per-file RTTM** — standard diarization format (written by
       :mod:`diarizer`, not this module).
    3. **Master CSV manifest** — one row per processed file with summary
       statistics.

All file writes are atomic (write-to-tmp → rename) to prevent corrupt
output on crashes.
"""

from __future__ import annotations

import csv
import json
import threading
from datetime import datetime, timezone
from pathlib import Path

import soundfile as sf
from loguru import logger

# Lock for thread-safe manifest CSV appends
_manifest_lock = threading.Lock()


# ── per-file JSON annotations ──────────────────────────────────


def write_json_annotations(
    segments: list[dict],
    file_id: str,
    out_dir: Path,
) -> Path:
    """Write enriched segment annotations to a per-file JSON file.

    Each element in the JSON array contains:
        ``speaker``, ``start``, ``end``, ``emotion``,
        ``emotion_confidence``, ``language``, ``language_confidence``.

    Args:
        segments: Enriched segment list (output of emotion + language stages).
        file_id:  Identifier used for the filename.
        out_dir:  Annotations directory.

    Returns:
        Path to the written JSON file.
    """
    json_path = out_dir / f"{file_id}.json"
    tmp_path = json_path.with_suffix(".json.tmp")

    # Ensure consistent key ordering for readability
    ordered_segments: list[dict] = []
    for seg in segments:
        ordered_segments.append(
            {
                "speaker": seg.get("speaker", "UNKNOWN"),
                "start": seg.get("start", 0.0),
                "end": seg.get("end", 0.0),
                "emotion": seg.get("emotion", "unknown"),
                "emotion_confidence": seg.get("emotion_confidence", 0.0),
                "language": seg.get("language", "unknown"),
                "language_confidence": seg.get("language_confidence", 0.0),
            }
        )

    tmp_path.write_text(
        json.dumps(ordered_segments, indent=2, ensure_ascii=False),
        encoding="utf-8",
    )
    tmp_path.replace(json_path)

    logger.info("[{}] Wrote JSON annotations ({} segments)", file_id, len(segments))
    return json_path


# ── master CSV manifest ─────────────────────────────────────────

MANIFEST_COLUMNS: list[str] = [
    "file_id",
    "url",
    "status",
    "num_speakers",
    "num_segments",
    "duration_sec",
    "languages_detected",
    "processed_at",
]


def update_manifest(
    file_id: str,
    url: str,
    status: str,
    segments: list[dict],
    wav_path: Path | None,
    out_dir: Path,
) -> Path:
    """Append or update a row in the master CSV manifest.

    Thread-safe via ``_manifest_lock``.

    Args:
        file_id:   Unique file identifier.
        url:       Original source URL or path.
        status:    ``"completed"`` or ``"failed"``.
        segments:  Enriched segment list (may be empty on failure).
        wav_path:  Path to WAV for duration calculation (``None`` on failure).
        out_dir:   Directory containing ``manifest.csv``.

    Returns:
        Path to the manifest CSV.
    """
    manifest_path = out_dir / "manifest.csv"
    summary = get_file_summary(segments, wav_path)

    row = {
        "file_id": file_id,
        "url": url,
        "status": status,
        "num_speakers": summary["num_speakers"],
        "num_segments": summary["num_segments"],
        "duration_sec": summary["duration_sec"],
        "languages_detected": summary["languages_detected"],
        "processed_at": datetime.now(timezone.utc).isoformat(),
    }

    with _manifest_lock:
        write_header = not manifest_path.exists()

        with open(manifest_path, "a", newline="", encoding="utf-8") as fh:
            writer = csv.DictWriter(fh, fieldnames=MANIFEST_COLUMNS)
            if write_header:
                writer.writeheader()
            writer.writerow(row)

    logger.debug("[{}] Manifest updated (status={})", file_id, status)
    return manifest_path


# ── summary helpers ─────────────────────────────────────────────


def get_file_summary(
    segments: list[dict],
    wav_path: Path | None,
) -> dict:
    """Compute summary statistics for one processed file.

    Args:
        segments: Enriched segment list.
        wav_path: WAV file path for duration (can be ``None``).

    Returns:
        Dict with ``num_speakers``, ``num_segments``, ``duration_sec``,
        ``languages_detected``.
    """
    num_speakers = len({s.get("speaker") for s in segments}) if segments else 0
    num_segments = len(segments)

    # Duration from the WAV file itself (more reliable than segment spans)
    duration_sec = 0.0
    if wav_path and wav_path.exists():
        try:
            info = sf.info(str(wav_path))
            duration_sec = round(info.frames / info.samplerate, 2)
        except Exception:
            # Fallback: use segment timestamps
            if segments:
                duration_sec = round(max(s.get("end", 0) for s in segments), 2)

    # Unique languages detected (excluding error / too_short)
    languages = sorted(
        {
            s.get("language", "unknown")
            for s in segments
            if s.get("language") not in ("too_short", "error", "unknown", None)
        }
    )
    languages_str = ",".join(languages) if languages else "unknown"

    return {
        "num_speakers": num_speakers,
        "num_segments": num_segments,
        "duration_sec": duration_sec,
        "languages_detected": languages_str,
    }


def get_output_as_dict(json_path: Path) -> list[dict]:
    """Read a previously written JSON annotation file back into memory.

    Args:
        json_path: Path to the ``.json`` annotation file.

    Returns:
        List of segment dicts, or empty list if file is missing/corrupt.
    """
    if not json_path.exists():
        logger.warning("Annotation file not found: {}", json_path)
        return []
    try:
        data = json.loads(json_path.read_text(encoding="utf-8"))
        if isinstance(data, list):
            return data
        logger.warning("Unexpected JSON structure in {}", json_path)
        return []
    except (json.JSONDecodeError, OSError) as exc:
        logger.error("Failed to read annotation {}: {}", json_path, exc)
        return []