services/whisper/utils/caption_director.py

"""
caption_director.py
---------------------------------------
AI Caption Intelligence System

Responsibilities:
- Convert transcript → styled caption segments
- Decide emphasis words
- Break text into readable chunks
- Optimize for TikTok/Reels retention
- Support hook-style captions

INPUT:
    words = [
        {"word": "hello", "start": 0.2, "end": 0.5},
        ...
    ]

OUTPUT:
    caption blocks:
    [
        {
            "text": "THIS IS CRAZY",
            "start": 0.2,
            "end": 2.1,
            "style": "hook"
        }
    ]
"""

import re


# -----------------------------
# CONFIG
# -----------------------------

MAX_WORDS_PER_CAPTION = 6
HOOK_KEYWORDS = [
    "listen", "wait", "you", "this", "crazy",
    "insane", "important", "stop", "secret"
]


# -----------------------------
# UTIL: CLEAN TEXT
# -----------------------------

def clean_word(word):
    return re.sub(r"[^a-zA-Z0-9']", "", word).lower()


# -----------------------------
# DETECT EMPHASIS
# -----------------------------

def is_emphasis(word):
    w = clean_word(word)
    return w in HOOK_KEYWORDS or len(word) > 8


# -----------------------------
# GROUP WORDS INTO CAPTIONS
# -----------------------------

def group_words(words):
    captions = []
    buffer = []

    for w in words:
        buffer.append(w)

        if len(buffer) >= MAX_WORDS_PER_CAPTION:
            captions.append(buffer)
            buffer = []

    if buffer:
        captions.append(buffer)

    return captions


# -----------------------------
# BUILD CAPTION BLOCK
# -----------------------------

def build_caption_block(group):
    text = []
    start = group[0]["start"]
    end = group[-1]["end"]

    emphasis_count = 0

    for w in group:
        word = w["word"]

        if is_emphasis(word):
            text.append(word.upper())
            emphasis_count += 1
        else:
            text.append(word)

    caption_text = " ".join(text)

    style = "hook" if emphasis_count > 0 else "normal"

    return {
        "text": caption_text,
        "start": start,
        "end": end,
        "style": style
    }


# -----------------------------
# MAIN DIRECTOR
# -----------------------------

def caption_director(words):
    """
    Main caption intelligence engine
    """

    if not words:
        return []

    grouped = group_words(words)

    captions = []

    for group in grouped:
        captions.append(build_caption_block(group))

    return captions


# -----------------------------
# HOOK CAPTION GENERATOR
# -----------------------------

def generate_hook_caption(words):
    """
    Extracts first high-impact caption
    """

    for w in words[:20]:
        if is_emphasis(w["word"]):
            return {
                "text": w["word"].upper(),
                "start": w["start"],
                "end": w["end"],
                "style": "hook"
            }

    return None


# -----------------------------
# AUTO CAPTION PIPELINE
# -----------------------------

def auto_captions(words):
    """
    Full pipeline:
    - detect hook
    - generate captions
    """

    captions = caption_director(words)

    hook = generate_hook_caption(words)

    if hook:
        captions.insert(0, hook)

    return captions


# -----------------------------
# STYLE DECISION ENGINE
# -----------------------------

def decide_style(caption):
    text = caption["text"]

    if caption["style"] == "hook":
        return "large_bold_center"

    if len(text) > 40:
        return "small_multi_line"

    if text.isupper():
        return "emphasis"

    return "standard"


# -----------------------------
# EXPORT HELPERS
# -----------------------------

def format_for_render(captions):
    """
    Converts captions into render-friendly format
    """

    formatted = []

    for c in captions:
        formatted.append({
            "text": c["text"],
            "start": c["start"],
            "end": c["end"],
            "style": decide_style(c)
        })

    return formatted


# -----------------------------
# PUBLIC API
# -----------------------------

def process_captions(words):
    """
    Full external API
    """

    captions = auto_captions(words)

    return format_for_render(captions)