Podcast Generation

import argparse
import base64
import json
import logging
import os
import uuid
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Literal, Optional

import requests

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


# Types
class ScriptLine:
    def __init__(self, speaker: Literal["male", "female"] = "male", paragraph: str = ""):
        self.speaker = speaker
        self.paragraph = paragraph


class Script:
    def __init__(self, locale: Literal["en", "zh"] = "en", lines: Optional[list[ScriptLine]] = None):
        self.locale = locale
        self.lines = lines or []

    @classmethod
    def from_dict(cls, data: dict) -> "Script":
        script = cls(locale=data.get("locale", "en"))
        for line in data.get("lines", []):
            script.lines.append(
                ScriptLine(
                    speaker=line.get("speaker", "male"),
                    paragraph=line.get("paragraph", ""),
                )
            )
        return script


def text_to_speech(text: str, voice_type: str) -> Optional[bytes]:
    """Convert text to speech using Volcengine TTS."""
    app_id = os.getenv("VOLCENGINE_TTS_APPID")
    access_token = os.getenv("VOLCENGINE_TTS_ACCESS_TOKEN")
    cluster = os.getenv("VOLCENGINE_TTS_CLUSTER", "volcano_tts")

    if not app_id or not access_token:
        raise ValueError(
            "VOLCENGINE_TTS_APPID and VOLCENGINE_TTS_ACCESS_TOKEN environment variables must be set"
        )

    url = "https://openspeech.bytedance.com/api/v1/tts"

    # Authentication: Bearer token with semicolon separator
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer;{access_token}",
    }

    payload = {
        "app": {
            "appid": app_id,
            "token": "access_token",  # literal string, not the actual token
            "cluster": cluster,
        },
        "user": {"uid": "podcast-generator"},
        "audio": {
            "voice_type": voice_type,
            "encoding": "mp3",
            "speed_ratio": 1.2,
        },
        "request": {
            "reqid": str(uuid.uuid4()),  # must be unique UUID
            "text": text,
            "text_type": "plain",
            "operation": "query",
        },
    }

    try:
        response = requests.post(url, json=payload, headers=headers)

        if response.status_code != 200:
            logger.error(f"TTS API error: {response.status_code} - {response.text}")
            return None

        result = response.json()
        if result.get("code") != 3000:
            logger.error(f"TTS error: {result.get('message')} (code: {result.get('code')})")
            return None

        audio_data = result.get("data")
        if audio_data:
            return base64.b64decode(audio_data)

    except Exception as e:
        logger.error(f"TTS error: {str(e)}")

    return None


def _process_line(args: tuple[int, ScriptLine, int]) -> tuple[int, Optional[bytes]]:
    """Process a single script line for TTS. Returns (index, audio_bytes)."""
    i, line, total = args

    # Select voice based on speaker gender
    if line.speaker == "male":
        voice_type = "zh_male_yangguangqingnian_moon_bigtts"  # Male voice
    else:
        voice_type = "zh_female_sajiaonvyou_moon_bigtts"  # Female voice

    logger.info(f"Processing line {i + 1}/{total} ({line.speaker})")
    audio = text_to_speech(line.paragraph, voice_type)

    if not audio:
        logger.warning(f"Failed to generate audio for line {i + 1}")

    return (i, audio)


def tts_node(script: Script, max_workers: int = 4) -> list[bytes]:
    """Convert script lines to audio chunks using TTS with multi-threading."""
    logger.info(f"Converting script to audio using {max_workers} workers...")

    total = len(script.lines)
    
    # Handle empty script case
    if total == 0:
        raise ValueError("Script contains no lines to process")

    # Validate required environment variables be