#!/usr/bin/env python3
"""
text_to_audio.py - High-Quality Text-to-Speech & Podcast Synthesizer for OpenClaw

Features:
- Primary TTS: Google Gemini 3.1 Flash TTS (natural, expressive, low cost).
- Automatic Fast Fallback: MiniMax Speech 02 HD (instant, 100% reliable studio quality).
- Parallel Synthesis: Multi-threaded chunk processing (ThreadPoolExecutor) for fast generation.
- Full Voice Normalization: Maps any voice name or role (Puck, Julius, Lara, Host, etc.) safely.
- Live Progress: Unbuffered stdout logging so OpenClaw process watchdogs never timeout.
- Format Guarantee: Standard monologue by default; multi-speaker dialog only when requested.
"""

import argparse
from concurrent.futures import ThreadPoolExecutor
import json
import os
import re
import shutil
import subprocess
import sys
import tempfile
import time
import urllib.error
import urllib.request
import wave
from pathlib import Path

OPENROUTER_SPEECH_URL = "https://openrouter.ai/api/v1/audio/speech"
DEFAULT_PRIMARY_MODEL = "google/gemini-3.1-flash-tts-preview"
DEFAULT_FALLBACK_MODEL = "minimax/speech-2.8-hd"

VOICE_CONFIG = {
    "de": {
        "gemini": {
            "solo": "Puck",
            "host": "Puck",
            "cohost": "Aoede",
            "deep": "Charon",
        },
        "minimax": {
            "solo": "presenter_male",
            "host": "presenter_male",
            "cohost": "presenter_female",
            "deep": "audiobook_male_1",
        },
    },
    "en": {
        "gemini": {
            "solo": "Zephyr",
            "host": "Zephyr",
            "cohost": "Aoede",
            "deep": "Charon",
        },
        "minimax": {
            "solo": "presenter_male",
            "host": "presenter_male",
            "cohost": "presenter_female",
            "deep": "audiobook_male_1",
        },
    },
}


def log_progress(msg: str) -> None:
    """Print immediate progress line with flush for OpenClaw process watchers."""
    print(f"[text-to-audio] {msg}", flush=True)


def find_ffmpeg() -> str:
    """Locate the ffmpeg binary across common container and host locations."""
    env_path = os.environ.get("FFMPEG_PATH")
    if env_path and os.path.exists(env_path):
        return env_path

    which_path = shutil.which("ffmpeg")
    if which_path:
        return which_path

    candidates = [
        "/home/node/.openclaw/bin/ffmpeg",
        "/opt/homebrew/bin/ffmpeg",
        "/usr/local/bin/ffmpeg",
        "/usr/bin/ffmpeg",
    ]
    for c in candidates:
        if os.path.exists(c) and os.access(c, os.X_OK):
            return c

    raise RuntimeError("ffmpeg binary not found. Please ensure ffmpeg is installed or set FFMPEG_PATH.")


def normalize_role(role_or_voice: str) -> str:
    """Normalize any role, speaker name, or legacy voice string into a canonical role key."""
    v = (role_or_voice or "").strip().lower()
    if any(x in v for x in ["cohost", "lara", "female", "expert", "aoede", "guest", "shaonv"]):
        return "cohost"
    if any(x in v for x in ["deep", "charon", "doku", "news", "audiobook"]):
        return "deep"
    return "solo"


def resolve_model_voice(role_or_voice: str, model: str, lang: str = "de") -> str:
    """Resolve a speaker role/name into the exact valid voice parameter for a specific TTS model."""
    lang_key = "en" if lang.startswith("en") else "de"
    is_gemini = "gemini" in model.lower()
    model_family = "gemini" if is_gemini else "minimax"

    # Direct passthrough if already a valid exact model voice
    if is_gemini and role_or_voice in ["Puck", "Aoede", "Charon", "Zephyr", "Kore", "Fenrir"]:
        return role_or_voice
    if not is_gemini and role_or_voice in [
        "presenter_male",
        "presenter_female",
        "audiobook_male_1",
        "audiobook_female_1",
        "male-qn-qingse",
        "female-shaonv",
    ]:
        return role_or_voice

    role = normalize_role(role_or_voice)
    return VOICE_CONFIG[lang_key][model_family][role]


def split_text_into_chunks(text: str, max_chars: int = 1200) -> list[str]:
    """Split text into natural sentence-level chunks to satisfy API limits and natural pacing."""
    text = text.strip()
    if not text:
        return []
    if len(text) <= max_chars:
        return [text]

    sentences = re.split(r"(?<=[.?!;:\n])\s+", text)
    chunks = []
    current = []
    current_len = 0

    for s in sentences:
        s = s.strip()
        if not s:
            continue
        if current_len + len(s) + 1 > max_chars and current:
            chunks.append(" ".join(current))
            current = [s]
            current_len = len(s)
        else:
            current.append(s)
            current_len += len(s) + 1

    if current:
        chunks.append(" ".join(current))

    return chunks


def _call_openrouter_speech(
    model: str,
    text: str,
    voice: str,
    response_format: str,
    api_key: str,
    timeout_sec: float,
) -> bytes:
    """Single HTTP call to OpenRouter /api/v1/audio/speech."""
    payload = {
        "model": model,
        "input": text,
        "voice": voice,
        "response_format": response_format,
    }
    req = urllib.request.Request(
        OPENROUTER_SPEECH_URL,
        data=json.dumps(payload).encode("utf-8"),
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
            "User-Agent": "OpenClaw-TextToAudio/2.1",
        },
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=timeout_sec) as resp:
        content = resp.read()
        if not content or len(content) < 100:
            raise RuntimeError("Received empty audio stream from TTS provider.")
        return content


def synthesize_chunk_to_wav(
    chunk_idx: int,
    total_chunks: int,
    text: str,
    speaker_or_voice: str,
    output_wav_path: str,
    ffmpeg_bin: str,
    preferred_model: str = DEFAULT_PRIMARY_MODEL,
    fallback_model: str = DEFAULT_FALLBACK_MODEL,
    lang: str = "de",
    api_key: str | None = None,
) -> None:
    """
    Synthesizes a text chunk to a standardized 24kHz mono 16-bit WAV file.
    Tries preferred_model first; on timeout/error, immediately falls back to fallback_model.
    """
    key = api_key or os.environ.get("OPENROUTER_API_KEY")
    if not key:
        raise ValueError("OPENROUTER_API_KEY environment variable is not set.")

    t0 = time.time()
    voice_pref = resolve_model_voice(speaker_or_voice, model=preferred_model, lang=lang)
    pref_is_gemini = "gemini" in preferred_model.lower()
    pref_format = "pcm" if pref_is_gemini else "mp3"

    log_progress(f"Chunk {chunk_idx + 1}/{total_chunks}: synthesizing ({len(text)} chars, {preferred_model} / {voice_pref})...")

    # Try preferred model (timeout: 16 seconds)
    try:
        content = _call_openrouter_speech(
            model=preferred_model,
            text=text,
            voice=voice_pref,
            response_format=pref_format,
            api_key=key,
            timeout_sec=16.0,
        )
        if pref_is_gemini:
            # Wrap raw 24kHz s16le PCM into standard WAV
            with wave.open(output_wav_path, "wb") as w:
                w.setnchannels(1)
                w.setsampwidth(2)
                w.setframerate(24000)
                w.writeframes(content)
        else:
            # Re-encode MP3 to uniform 24kHz mono WAV
            with tempfile.NamedTemporaryFile(suffix=".mp3", delete=False) as tf:
                tf.write(content)
                tmp_mp3 = tf.name
            try:
                subprocess.run(
                    [ffmpeg_bin, "-y", "-i", tmp_mp3, "-ar", "24000", "-ac", "1", output_wav_path],
                    stdout=subprocess.DEVNULL,
                    stderr=subprocess.DEVNULL,
                    check=True,
                )
            finally:
                if os.path.exists(tmp_mp3):
                    os.remove(tmp_mp3)

        dt = time.time() - t0
        log_progress(f"Chunk {chunk_idx + 1}/{total_chunks}: generated via {preferred_model} in {dt:.1f}s")
        return
    except Exception as err:
        log_progress(f"Chunk {chunk_idx + 1}/{total_chunks}: primary {preferred_model} failed ({err}). Falling back to {fallback_model}...")

    # Immediate fallback
    t1 = time.time()
    voice_fb = resolve_model_voice(speaker_or_voice, model=fallback_model, lang=lang)
    fb_is_gemini = "gemini" in fallback_model.lower()
    fb_format = "pcm" if fb_is_gemini else "mp3"

    try:
        content = _call_openrouter_speech(
            model=fallback_model,
            text=text,
            voice=voice_fb,
            response_format=fb_format,
            api_key=key,
            timeout_sec=15.0,
        )
        if fb_is_gemini:
            with wave.open(output_wav_path, "wb") as w:
                w.setnchannels(1)
                w.setsampwidth(2)
                w.setframerate(24000)
                w.writeframes(content)
        else:
            with tempfile.NamedTemporaryFile(suffix=".mp3", delete=False) as tf:
                tf.write(content)
                tmp_mp3 = tf.name
            try:
                subprocess.run(
                    [ffmpeg_bin, "-y", "-i", tmp_mp3, "-ar", "24000", "-ac", "1", output_wav_path],
                    stdout=subprocess.DEVNULL,
                    stderr=subprocess.DEVNULL,
                    check=True,
                )
            finally:
                if os.path.exists(tmp_mp3):
                    os.remove(tmp_mp3)

        dt = time.time() - t1
        log_progress(f"Chunk {chunk_idx + 1}/{total_chunks}: generated via fallback {fallback_model} in {dt:.1f}s")
        return
    except Exception as fb_err:
        raise RuntimeError(
            f"Chunk {chunk_idx + 1}/{total_chunks} synthesis failed on both primary ({err}) and fallback ({fb_err})"
        )


def generate_silence_wav(duration_sec: float, output_path: str, sample_rate: int = 24000) -> None:
    """Generate a clean silence WAV file."""
    num_frames = int(duration_sec * sample_rate)
    silence_bytes = b"\x00\x00" * num_frames
    with wave.open(output_path, "wb") as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(sample_rate)
        w.writeframes(silence_bytes)


def stitch_audio_segments(
    ffmpeg_bin: str,
    wav_files: list[str],
    output_file: str,
    title: str = "Audio Dossier",
    artist: str = "OpenClaw",
    album: str = "OpenClaw Audio Briefings",
) -> None:
    """Concatenate uniform WAV segments and produce high-quality 128k MP3 with metadata."""
    if not wav_files:
        raise ValueError("No audio segments to stitch.")

    log_progress(f"Stitching {len(wav_files)} segment(s) into MP3 with metadata...")

    if len(wav_files) == 1:
        cmd = [
            ffmpeg_bin,
            "-y",
            "-i",
            wav_files[0],
            "-metadata",
            f"title={title}",
            "-metadata",
            f"artist={artist}",
            "-metadata",
            f"album={album}",
            "-c:a",
            "libmp3lame",
            "-b:a",
            "128k",
            output_file,
        ]
        res = subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True)
        if res.returncode != 0:
            raise RuntimeError(f"FFmpeg export failed: {res.stderr}")
        return

    # Concat demuxer with single-quoted file paths
    with tempfile.NamedTemporaryFile("w", suffix=".txt", delete=False) as list_file:
        for f in wav_files:
            safe_path = f.replace("'", "'\\''")
            list_file.write(f"file '{safe_path}'\n")
        list_file_path = list_file.name

    try:
        cmd = [
            ffmpeg_bin,
            "-y",
            "-f",
            "concat",
            "-safe",
            "0",
            "-i",
            list_file_path,
            "-metadata",
            f"title={title}",
            "-metadata",
            f"artist={artist}",
            "-metadata",
            f"album={album}",
            "-c:a",
            "libmp3lame",
            "-b:a",
            "128k",
            output_file,
        ]
        res = subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True)
        if res.returncode != 0:
            raise RuntimeError(f"FFmpeg concatenation failed: {res.stderr}")
    finally:
        if os.path.exists(list_file_path):
            os.remove(list_file_path)


def process_solo_text(
    text: str,
    ffmpeg_bin: str,
    voice_or_role: str,
    primary_model: str,
    fallback_model: str,
    lang: str,
    output_file: str,
    title: str,
    max_workers: int = 4,
) -> None:
    """Generate audio for a single-speaker text in parallel (Default for all audio dossiers)."""
    chunks = split_text_into_chunks(text)
    if not chunks:
        raise ValueError("No text provided to synthesize.")

    log_progress(f"Solo text synthesis started: {len(chunks)} chunk(s), ~{len(text.split())} words, lang={lang}")

    with tempfile.TemporaryDirectory(prefix="audio_solo_") as tmpdir:
        wav_tasks = []
        for i, chunk in enumerate(chunks):
            chunk_wav = os.path.join(tmpdir, f"chunk_{i:04d}.wav")
            wav_tasks.append((i, len(chunks), chunk, voice_or_role, chunk_wav))

        with ThreadPoolExecutor(max_workers=min(len(chunks), max_workers)) as pool:
            futures = [
                pool.submit(
                    synthesize_chunk_to_wav,
                    idx,
                    tot,
                    chk,
                    v,
                    path,
                    ffmpeg_bin,
                    primary_model,
                    fallback_model,
                    lang,
                )
                for idx, tot, chk, v, path in wav_tasks
            ]
            for fut in futures:
                fut.result()

        wav_files = [task[4] for task in wav_tasks]
        stitch_audio_segments(
            ffmpeg_bin=ffmpeg_bin,
            wav_files=wav_files,
            output_file=output_file,
            title=title,
        )


def process_dialog(
    dialog_items: list[dict],
    ffmpeg_bin: str,
    primary_model: str,
    fallback_model: str,
    lang: str,
    output_file: str,
    title: str,
    pause_duration: float = 0.35,
    max_workers: int = 4,
) -> None:
    """Generate audio for multi-speaker dialog turns in parallel (used only when 'podcast' requested)."""
    if not dialog_items:
        raise ValueError("Dialog list is empty.")

    log_progress(f"Dialog synthesis started: {len(dialog_items)} turn(s), lang={lang}")

    with tempfile.TemporaryDirectory(prefix="audio_dialog_") as tmpdir:
        silence_file = os.path.join(tmpdir, "pause.wav")
        generate_silence_wav(pause_duration, silence_file)

        # Flatten dialog items into chunks if any turn exceeds max_chars
        flattened_items = []
        for d in dialog_items:
            speaker = d.get("speaker", "host")
            text = d.get("text", "").strip()
            if not text:
                continue
            chunks = split_text_into_chunks(text)
            for c in chunks:
                flattened_items.append((speaker, c))

        wav_tasks = []
        for i, (speaker, chunk_text) in enumerate(flattened_items):
            turn_wav = os.path.join(tmpdir, f"turn_{i:04d}.wav")
            wav_tasks.append((i, len(flattened_items), chunk_text, speaker, turn_wav))

        with ThreadPoolExecutor(max_workers=min(len(flattened_items), max_workers)) as pool:
            futures = [
                pool.submit(
                    synthesize_chunk_to_wav,
                    idx,
                    tot,
                    chk,
                    spk,
                    path,
                    ffmpeg_bin,
                    primary_model,
                    fallback_model,
                    lang,
                )
                for idx, tot, chk, spk, path in wav_tasks
            ]
            for fut in futures:
                fut.result()

        # Interleave silence between dialog turns
        stitched_wavs = []
        for task in wav_tasks:
            stitched_wavs.append(task[4])
            stitched_wavs.append(silence_file)

        if stitched_wavs and stitched_wavs[-1] == silence_file:
            stitched_wavs.pop()

        stitch_audio_segments(
            ffmpeg_bin=ffmpeg_bin,
            wav_files=stitched_wavs,
            output_file=output_file,
            title=title,
        )


def main():
    parser = argparse.ArgumentParser(description="OpenClaw Text-to-Audio Dossier Generator")
    parser.add_argument("--text", help="Plaintext to synthesize (default solo mode)")
    parser.add_argument("--file", help="Path to plaintext file (default solo mode)")
    parser.add_argument("--dialog-json", help="JSON string of dialog turns (only if 'podcast' requested)")
    parser.add_argument("--dialog-file", help="Path to JSON file with dialog turns")
    parser.add_argument("--output", "-o", help="Target output MP3 path")
    parser.add_argument("--title", default="Audio Dossier", help="Audio title for ID3 metadata")
    parser.add_argument("--voice", help="Voice or speaker role (e.g. Puck, Zephyr, Aoede, host, solo)")
    parser.add_argument("--lang", default="de", choices=["de", "en"], help="Language preset (de/en)")
    parser.add_argument("--model", default=DEFAULT_PRIMARY_MODEL, help=f"Primary model (default: {DEFAULT_PRIMARY_MODEL})")
    parser.add_argument("--fallback-model", default=DEFAULT_FALLBACK_MODEL, help=f"Fallback model (default: {DEFAULT_FALLBACK_MODEL})")
    parser.add_argument("--pause", type=float, default=0.35, help="Pause in seconds between dialog turns")
    parser.add_argument("--workers", type=int, default=4, help="Max concurrent synthesis threads")
    parser.add_argument("--media-tag", action="store_true", default=True, help="Print MEDIA:<path> tag to stdout")

    args = parser.parse_args()

    ffmpeg_bin = find_ffmpeg()

    if args.output:
        out_path = Path(args.output).expanduser().resolve()
    else:
        ts = int(time.time())
        out_path = Path(f"/tmp/audio-dossier-{ts}.mp3")

    out_path.parent.mkdir(parents=True, exist_ok=True)

    t_start = time.time()

    dialog_data = None
    if args.dialog_file:
        with open(args.dialog_file, "r", encoding="utf-8") as f:
            dialog_data = json.load(f)
    elif args.dialog_json:
        dialog_data = json.loads(args.dialog_json)

    if dialog_data:
        if not isinstance(dialog_data, list):
            raise ValueError("Dialog input must be a JSON array of objects with 'speaker' and 'text'.")
        process_dialog(
            dialog_items=dialog_data,
            ffmpeg_bin=ffmpeg_bin,
            primary_model=args.model,
            fallback_model=args.fallback_model,
            lang=args.lang,
            output_file=str(out_path),
            title=args.title,
            pause_duration=args.pause,
            max_workers=args.workers,
        )
    else:
        # Solo monologue mode (Strict DEFAULT)
        raw_text = ""
        if args.file:
            with open(args.file, "r", encoding="utf-8") as f:
                raw_text = f.read()
        elif args.text:
            raw_text = args.text
        elif not sys.stdin.isatty():
            raw_text = sys.stdin.read()
        else:
            parser.error("Must provide either --text, --file, --dialog-json, --dialog-file, or pipe text to stdin.")

        chosen_voice = args.voice or "solo"
        process_solo_text(
            text=raw_text,
            ffmpeg_bin=ffmpeg_bin,
            voice_or_role=chosen_voice,
            primary_model=args.model,
            fallback_model=args.fallback_model,
            lang=args.lang,
            output_file=str(out_path),
            title=args.title,
            max_workers=args.workers,
        )

    total_time = time.time() - t_start
    log_progress(f"Finished in {total_time:.1f}s: {out_path} ({out_path.stat().st_size} bytes)")

    if args.media_tag:
        print(f"MEDIA:{out_path}", flush=True)


if __name__ == "__main__":
    main()