services/whisper/utils/srt.py

from typing import List, Dict, Any, Union


# =====================================================
# PUBLIC API (USED BY MAIN.PY)
# =====================================================

def generate_srt(data: List[Dict[str, Any]]) -> str:
    """
    Universal SRT generator for:
    - Whisper word output (V1–V7)
    - Highlight segments (start/end grouped words)
    - Mixed/partial structures

    Expected input formats:
    1. Word-level:
        {"text": "...", "start": float, "end": float}

    2. Segment-level:
        [{"start": float, "end": float, "text": "..."}]

    Returns:
        SRT formatted string
    """

    if not data:
        return ""

    normalized = _normalize_input(data)
    return _build_srt(normalized)


# =====================================================
# NORMALIZATION LAYER (CRITICAL FOR V1–V7 COMPATIBILITY)
# =====================================================

def _normalize_input(data: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    """
    Converts any supported structure into unified subtitle blocks
    """

    normalized = []

    # CASE 1: Already segment-based
    if isinstance(data[0], dict) and "start" in data[0] and "end" in data[0] and "text" in data[0]:
        for item in data:
            normalized.append({
                "start": float(item.get("start", 0)),
                "end": float(item.get("end", 0)),
                "text": str(item.get("text", "")).strip()
            })
        return normalized

    # CASE 2: Whisper word-level output
    buffer = []
    current_start = None

    for w in data:

        if not isinstance(w, dict):
            continue

        text = str(w.get("text", "")).strip()
        start = w.get("start", None)
        end = w.get("end", None)

        if start is None or end is None:
            continue

        if current_start is None:
            current_start = start

        buffer.append(text)

        # Chunking strategy: group every ~8–12 words
        if len(buffer) >= 10:

            normalized.append({
                "start": current_start,
                "end": end,
                "text": " ".join(buffer)
            })

            buffer = []
            current_start = None

    # flush remaining buffer
    if buffer:
        normalized.append({
            "start": current_start or 0,
            "end": data[-1].get("end", 0),
            "text": " ".join(buffer)
        })

    return normalized


# =====================================================
# SRT BUILDER
# =====================================================

def _build_srt(items: List[Dict[str, Any]]) -> str:
    """
    Converts normalized subtitle blocks → SRT format
    """

    output = []
    index = 1

    for item in items:

        start = _format_time(item["start"])
        end = _format_time(item["end"])
        text = _clean_text(item["text"])

        if not text:
            continue

        output.append(f"{index}")
        output.append(f"{start} --> {end}")
        output.append(f"{text}")
        output.append("")  # blank line separator

        index += 1

    return "\n".join(output).strip()


# =====================================================
# TIME FORMATTER
# =====================================================

def _format_time(seconds: Union[int, float]) -> str:
    """
    Converts seconds → SRT timestamp format
    HH:MM:SS,mmm
    """

    try:
        seconds = float(seconds)
    except:
        seconds = 0.0

    hrs = int(seconds // 3600)
    mins = int((seconds % 3600) // 60)
    secs = int(seconds % 60)
    ms = int((seconds - int(seconds)) * 1000)

    return f"{hrs:02}:{mins:02}:{secs:02},{ms:03}"


# =====================================================
# TEXT CLEANER (IMPORTANT FOR VIDEO RENDERING STABILITY)
# =====================================================

def _clean_text(text: str) -> str:
    """
    Sanitizes subtitle text for rendering engines
    """

    if not text:
        return ""

    text = text.replace("\n", " ")
    text = text.replace("\r", " ")

    # remove excessive spacing
    text = " ".join(text.split())

    return text.strip()


# =====================================================
# OPTIONAL DEBUG HELPER (SAFE IN PRODUCTION)
# =====================================================

def debug_srt(data: List[Dict[str, Any]]) -> dict:
    """
    Returns structured preview for debugging pipelines
    """

    normalized = _normalize_input(data)

    return {
        "blocks": len(normalized),
        "sample": normalized[:3],
        "duration": normalized[-1]["end"] if normalized else 0
    }