scripts/generate_social_science_proof_audio.py

"""Generate real bilingual narration for the Social Science exam-point proof.

English uses local Kokoro. Malayalam uses the locally cached AI4Bharat Indic
Parler model. The lesson JSON is already separately authored in each language;
this script never machine-translates one version into the other.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import os
import re
import subprocess
import sys
from pathlib import Path

import numpy as np
import soundfile as sf


ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "backend"))


def probe_duration(path: Path) -> float:
    run = subprocess.run(
        [
            "ffprobe",
            "-v",
            "error",
            "-show_entries",
            "format=duration",
            "-of",
            "default=nw=1:nk=1",
            str(path),
        ],
        check=True,
        capture_output=True,
        text=True,
    )
    return round(float(run.stdout.strip()), 6)


def append_silence(path: Path, seconds: float) -> None:
    if seconds <= 0:
        return
    audio, sample_rate = sf.read(path, dtype="float32", always_2d=True)
    silence = np.zeros((round(seconds * sample_rate), audio.shape[1]), dtype=np.float32)
    sf.write(path, np.concatenate([audio, silence], axis=0), sample_rate, subtype="PCM_16")


def retime_audio(path: Path, speed: float) -> None:
    if abs(speed - 1.0) < 0.001:
        return
    if speed < 0.5 or speed > 2.0:
        raise ValueError("Audio speed must be between 0.5 and 2.0.")
    temporary = path.with_name(f"{path.stem}.retimed.wav")
    subprocess.run(
        [
            "ffmpeg",
            "-hide_banner",
            "-loglevel",
            "error",
            "-y",
            "-i",
            str(path),
            "-filter:a",
            f"atempo={speed:.4f}",
            "-c:a",
            "pcm_s16le",
            str(temporary),
        ],
        check=True,
    )
    temporary.replace(path)


def kokoro_audio(text: str, output: Path, voice: str, speed: float) -> None:
    from kokoro import KPipeline

    pipeline = KPipeline(lang_code="b" if voice.startswith("b") else "a")
    parts = [
        audio
        for _graphemes, _phonemes, audio in pipeline(
            re.sub(r"\s+", " ", text).strip(),
            voice=voice,
            speed=speed,
            split_pattern=r"\n+",
        )
    ]
    if not parts:
        raise RuntimeError("Kokoro returned no audio.")
    sf.write(output, np.concatenate(parts).astype(np.float32), 24_000, subtype="PCM_16")


def ai4bharat_audio(text: str, output: Path, voice: str) -> None:
    from app.schemas.video import VideoSceneAudioInput
    from app.services.tts_provider import AI4BharatIndicParlerTTSProvider

    provider = AI4BharatIndicParlerTTSProvider()
    sentences = [
        part.strip()
        for part in re.split(r"(?<=[.!?])\s+", re.sub(r"\s+", " ", text).strip())
        if part.strip()
    ]
    narration_chunks: list[str] = []
    for sentence in sentences:
        words = sentence.split()
        chunk_count = max(1, math.ceil(len(words) / 7))
        base_size, larger_chunks = divmod(len(words), chunk_count)
        cursor = 0
        for chunk_index in range(chunk_count):
            chunk_size = base_size + (1 if chunk_index < larger_chunks else 0)
            narration_chunks.append(" ".join(words[cursor : cursor + chunk_size]))
            cursor += chunk_size

    rendered_parts: list[tuple[np.ndarray, int]] = []
    temporary_paths: list[Path] = []
    try:
        for chunk_index, narration_chunk in enumerate(narration_chunks, start=1):
            chunk_hash = hashlib.sha256(narration_chunk.encode("utf-8")).hexdigest()[:8]
            chunk_output = output.with_name(
                f"{output.stem}.part-{chunk_index:02d}-{chunk_hash}.wav"
            )
            temporary_paths.append(chunk_output)
            print(
                json.dumps(
                    {
                        "event": "social_science_proof_tts_chunk_start",
                        "chunk": chunk_index,
                        "chunkCount": len(narration_chunks),
                        "wordCount": len(narration_chunk.split()),
                        "text": narration_chunk,
                    },
                    ensure_ascii=True,
                ),
                flush=True,
            )
            if chunk_output.exists():
                actual_path = chunk_output
            else:
                scene = VideoSceneAudioInput(
                    scene_id=chunk_index,
                    type="concept",
                    duration_seconds=max(3.0, len(narration_chunk.split()) / 1.75),
                    voice_text=narration_chunk,
                )
                result = provider.generate_scene_audio(
                    scene=scene,
                    output_file=chunk_output,
                    voice_mode="malayalam_soft",
                    voice=voice,
                    language="ml",
                )
                actual_path = result.file_path
            audio, sample_rate = sf.read(actual_path, dtype="float32", always_2d=True)
            rendered_parts.append((audio, sample_rate))
            if actual_path != chunk_output:
                temporary_paths.append(actual_path)
        if not rendered_parts:
            raise RuntimeError("AI4Bharat returned no audio.")
        sample_rate = rendered_parts[0][1]
        if any(part_sample_rate != sample_rate for _audio, part_sample_rate in rendered_parts):
            raise RuntimeError("AI4Bharat chunks used inconsistent sample rates.")
        join_pause = np.zeros((round(0.16 * sample_rate), rendered_parts[0][0].shape[1]), dtype=np.float32)
        combined: list[np.ndarray] = []
        for part_index, (audio, _sample_rate) in enumerate(rendered_parts):
            if part_index:
                combined.append(join_pause)
            combined.append(audio)
        sf.write(output, np.concatenate(combined, axis=0), sample_rate, subtype="PCM_16")
    finally:
        if output.exists():
            for temporary_path in temporary_paths:
                temporary_path.unlink(missing_ok=True)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--lesson", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--provider", choices=("kokoro", "ai4bharat"), required=True)
    parser.add_argument("--language", choices=("en", "ml"))
    parser.add_argument("--voice")
    parser.add_argument("--speed", type=float, default=0.96)
    parser.add_argument("--audio-speed", type=float, default=1.0)
    parser.add_argument(
        "--retime-existing-factor",
        type=float,
        default=1.0,
        help="Post-process existing WAV chunks by this tempo factor without re-synthesising them.",
    )
    parser.add_argument("--force", action="store_true")
    args = parser.parse_args()

    os.environ.setdefault("USE_TF", "0")
    os.environ.setdefault("TRANSFORMERS_NO_TF", "1")
    os.environ.setdefault("HF_HUB_DISABLE_XET", "1")

    lesson_path = args.lesson if args.lesson.is_absolute() else ROOT / args.lesson
    output_dir = args.output_dir if args.output_dir.is_absolute() else ROOT / args.output_dir
    output_dir.mkdir(parents=True, exist_ok=True)
    lesson = json.loads(lesson_path.read_text(encoding="utf-8"))
    language = args.language or lesson.get("language")
    if language not in {"en", "ml"}:
        raise ValueError("Lesson language must be supplied with --language en|ml.")
    if language == "en" and args.provider != "kokoro":
        raise ValueError("English proof narration must use Kokoro.")
    if language == "ml" and args.provider != "ai4bharat":
        raise ValueError("Malayalam proof narration must use AI4Bharat.")
    voice = args.voice or ("af_heart" if args.provider == "kokoro" else "Anjali")
    items = lesson.get("units") or lesson.get("scenes")
    if not items:
        raise ValueError("Lesson contains no narration units.")

    chunks: list[dict] = []
    for index, scene in enumerate(items, start=1):
        narration = scene.get(f"tts_text_{language}") or scene.get("narration")
        if not narration:
            raise ValueError(f"Narration unit {scene['id']} is missing tts_text_{language}.")
        if language == "ml":
            latin_words = re.findall(r"\b[A-Za-z]{3,}\b", narration)
            if len(latin_words) > 1:
                raise ValueError(
                    f"Malayalam narration unit {scene['id']} contains excessive Latin-script prose: {latin_words}"
                )
        output = output_dir / f"{index:02d}-{scene['id']}.wav"
        retrieval_pause = float(scene.get("retrievalPauseSeconds", 0))
        tail_pause = float(scene.get("pauseAfterSeconds", 0.35))
        if args.force or not output.exists():
            print(
                json.dumps(
                    {
                        "event": "social_science_proof_tts_start",
                        "provider": args.provider,
                        "scene": scene["id"],
                        "index": index,
                    }
                ),
                flush=True,
            )
            if args.provider == "kokoro":
                kokoro_audio(narration, output, voice, args.speed)
            else:
                ai4bharat_audio(narration, output, voice)
            retime_audio(output, args.audio_speed)
            append_silence(output, retrieval_pause + tail_pause)
        elif abs(args.retime_existing_factor - 1.0) > 0.001:
            retime_audio(output, args.retime_existing_factor)
        duration_seconds = probe_duration(output)
        chunks.append(
            {
                "id": scene["id"],
                "sceneIndex": index - 1,
                "outputPath": output.relative_to(ROOT).as_posix(),
                "narration": narration,
                "displayText": scene.get(f"display_text_{language}", narration),
                "wordCount": len(re.findall(r"\S+", narration)),
                "durationSeconds": duration_seconds,
                "retrievalPauseSeconds": retrieval_pause,
                "pauseAfterSeconds": tail_pause,
            }
        )
        print(
            json.dumps(
                {
                    "event": "social_science_proof_tts_complete",
                    "provider": args.provider,
                    "scene": scene["id"],
                    "durationSeconds": duration_seconds,
                }
            ),
            flush=True,
        )

    manifest = {
        "provider": f"{args.provider}-local",
        "voice": voice,
        "synthesisSpeed": args.speed if args.provider == "kokoro" else 1.0,
        "postAudioSpeed": args.audio_speed * args.retime_existing_factor,
        "language": language,
        "lessonPath": lesson_path.relative_to(ROOT).as_posix(),
        "sceneCount": len(chunks),
        "wordCount": sum(item["wordCount"] for item in chunks),
        "totalDurationSeconds": round(sum(item["durationSeconds"] for item in chunks), 3),
        "chunks": chunks,
    }
    manifest_path = output_dir / "manifest.json"
    manifest_path.write_text(
        json.dumps(manifest, indent=2, ensure_ascii=False) + "\n",
        encoding="utf-8",
    )
    print(
        json.dumps(
            {
                "event": "social_science_proof_manifest_ready",
                "manifest": str(manifest_path),
                "durationSeconds": manifest["totalDurationSeconds"],
                "sceneCount": manifest["sceneCount"],
            }
        )
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())