segment_roles.py

#!/usr/bin/env python3

from __future__ import annotations

import json
import re
from pathlib import Path

ROLE_NARRATION = "narration"
ROLE_DIALOGUE = "dialogue"

SHORT_DIALOGUE_THRESHOLD = 40

def _split_paragraphs(text: str) -> list[str]:
    
    text = text.strip()
    
    parts = re.split(r'\n(?:\s*\n)+|\n---\n', text)
    return [p.strip() for p in parts if p.strip()]

def _classify_paragraph(paragraph: str) -> str:
    
    stripped = paragraph.strip()
    
    if stripped.startswith('"') and stripped.endswith('"'):
        return ROLE_DIALOGUE
    return ROLE_NARRATION

def _split_paragraph_roles(paragraph: str) -> list[dict]:
    
    stripped = paragraph.strip()
    if not stripped:
        return []

    if stripped.startswith('"') and stripped.endswith('"') and stripped.count('"') == 2:
        return [{"role": ROLE_DIALOGUE, "text": stripped}]
    if '"' not in stripped:
        return [{"role": ROLE_NARRATION, "text": stripped}]

    segments: list[dict] = []
    pos = 0
    while pos < len(stripped):
        quote_start = stripped.find('"', pos)
        if quote_start < 0:
            rest = stripped[pos:].strip()
            if rest:
                segments.append({"role": ROLE_NARRATION, "text": rest})
            break

        before = stripped[pos:quote_start].strip()
        if before:
            segments.append({"role": ROLE_NARRATION, "text": before})

        quote_end = stripped.find('"', quote_start + 1)
        if quote_end < 0:
            
            rest = stripped[quote_start:].strip()
            if rest:
                segments.append({"role": ROLE_NARRATION, "text": rest})
            break

        quoted = stripped[quote_start:quote_end + 1]
        segments.append({"role": ROLE_DIALOGUE, "text": quoted})
        pos = quote_end + 1

    return segments

def segment_roles(text: str) -> list[dict]:
    
    paragraphs = _split_paragraphs(text)
    if not paragraphs:
        return []

    spans: list[dict] = []
    for para in paragraphs:
        for seg in _split_paragraph_roles(para):
            if spans and spans[-1]["role"] == seg["role"]:
                
                if (seg["role"] == ROLE_DIALOGUE
                        and len(seg["text"]) <= SHORT_DIALOGUE_THRESHOLD):
                    spans.append(dict(seg))
                else:
                    spans[-1]["text"] += "\n\n" + seg["text"]
            else:
                spans.append(dict(seg))

    return spans

def chunk_text_multi(
    text: str,
    max_chars: int,
    chunk_fn,
) -> list[dict]:
    
    spans = segment_roles(text)
    result: list[dict] = []
    for span in spans:
        chunks = chunk_fn(span["text"], max_chars)
        for chunk_text in chunks:
            result.append({"role": span["role"], "text": chunk_text})
    return result

def has_dialogue(text: str, min_paragraphs: int = 2) -> bool:
    
    paragraphs = _split_paragraphs(text)
    count = sum(1 for p in paragraphs if _classify_paragraph(p) == ROLE_DIALOGUE)
    return count >= min_paragraphs

DEFAULT_ENGLISH_VOICE_ROLES = {
    "roles": {
        "dialogue": {
            "seed_audio": "refs/en_whisper.mp3",
            "seed_text_file": "supporting_scripts/voice_test_seed_ref_en.txt",
        },
        "narration": {
            "seed_audio": "refs/en_narrator_d.wav",
            "seed_text_file": "refs/en_narrator_d.txt",
        },
    },
    "segmentation": {"strategy": "quotes"},
}

def load_voice_roles(script_path: str | Path, *, text: str | None = None, lang: str | None = None) -> dict | None:
    
    script_path = Path(script_path).resolve()
    
    collection_dir = script_path.parent
    config_path = collection_dir / "voice_roles.json"
    if config_path.exists():
        try:
            return json.loads(config_path.read_text(encoding="utf-8"))
        except (json.JSONDecodeError, OSError):
            pass

    if lang and lang == "english" and text and has_dialogue(text):
        return DEFAULT_ENGLISH_VOICE_ROLES

    return None