#!/usr/bin/env python3
"""
text_to_audio.py - OpenRouter Text-to-Speech & Podcast Synthesizer

High-quality speech synthesis using Google Gemini 3.1 Flash TTS (default) or other OpenRouter TTS models.
Produces single-speaker audio dossiers by default, and multi-speaker dialogue only when explicitly requested.
"""

import argparse
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

DEFAULT_MODEL = "google/gemini-3.1-flash-tts-preview"
FALLBACK_MODEL = "microsoft/mai-voice-2"
OPENROUTER_SPEECH_URL = "https://openrouter.ai/api/v1/audio/speech"

# Pre-selected voices for German and English
DEFAULT_VOICES = {
    "de": {
        "solo": "Puck",         # Warm, lebendig, dynamisch (Standard für Solo Audio-Dossiers)
        "host": "Puck",         # Host im Dialog-Modus (nur bei expliziter "Podcast"-Anfrage)
        "cohost": "Aoede",      # Melodisch, klar, ausdrucksstark (Co-Host / Expertin)
        "expert": "Aoede",
        "male": "Puck",
        "female": "Aoede",
        "deep": "Charon",       # Ruhig, tief, autoritär (Nachrichten/Doku)
        "zephyr": "Zephyr",     # Klar, ausgewogen, natürlich
        "puck": "Puck",
        "aoede": "Aoede",
        "charon": "Charon",
        "kore": "Kore",
        "fenrir": "Fenrir",
    },
    "en": {
        "solo": "Zephyr",       # Clear, natural, balanced conversational male/neutral
        "host": "Zephyr",       # Host in dialog mode
        "cohost": "Aoede",      # Expressive, melodious, friendly female
        "expert": "Aoede",
        "male": "Zephyr",
        "female": "Aoede",
        "deep": "Charon",       # Deep, resonant, calm male
        "puck": "Puck",         # Energetic, warm, youthful
        "aoede": "Aoede",
        "zephyr": "Zephyr",
        "charon": "Charon",
        "kore": "Kore",
    },
}


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",
        "/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 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 synthesize_segment(
    text: str,
    voice: str,
    output_path: str,
    model: str = DEFAULT_MODEL,
    api_key: str | None = None,
    max_retries: int = 3,
) -> None:
    """
    Synthesizes speech using OpenRouter with automatic retries and format handling.
    - For Gemini models: requests 'pcm' and writes standard 24kHz WAV.
    - For other models (MiniMax, Microsoft, Grok): requests 'mp3' and writes MP3.
    """
    key = api_key or os.environ.get("OPENROUTER_API_KEY")
    if not key:
        raise ValueError("OPENROUTER_API_KEY environment variable is not set.")

    is_gemini = "gemini" in model.lower()
    response_format = "pcm" if is_gemini else "mp3"

    payload = {
        "model": model,
        "input": text,
        "voice": voice,
        "response_format": response_format,
    }
    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        OPENROUTER_SPEECH_URL,
        data=data,
        headers={
            "Authorization": f"Bearer {key}",
            "Content-Type": "application/json",
            "User-Agent": "OpenClaw-TextToAudio/2.0",
        },
        method="POST",
    )

    last_error = None
    for attempt in range(1, max_retries + 1):
        try:
            with urllib.request.urlopen(req, timeout=60) as resp:
                content = resp.read()
                if not content or len(content) < 100:
                    raise RuntimeError("Received empty or incomplete audio stream from TTS provider.")

                if is_gemini:
                    # Wrap raw 24kHz s16le PCM into standard WAV
                    with wave.open(output_path, "wb") as w:
                        w.setnchannels(1)
                        w.setsampwidth(2)       # 16-bit
                        w.setframerate(24000)   # 24kHz
                        w.writeframes(content)
                else:
                    with open(output_path, "wb") as f:
                        f.write(content)
                return  # Success!
        except urllib.error.HTTPError as e:
            body = e.read().decode("utf-8", errors="replace")
            last_error = f"HTTP {e.code}: {body}"
            # Retry on 5xx or transient 429
            if attempt < max_retries and e.code in (429, 500, 502, 503, 504):
                time.sleep(1.5 * attempt)
                continue
            break
        except Exception as e:
            last_error = str(e)
            if attempt < max_retries:
                time.sleep(1.5 * attempt)
                continue
            break

    raise RuntimeError(f"OpenRouter TTS failed after {max_retries} attempts: {last_error}")


def generate_silence_wav(duration_sec: float, output_path: str, sample_rate: int = 24000) -> None:
    """Generate a clean silence WAV file using Python's standard library wave module."""
    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,
    audio_files: list[str],
    output_file: str,
    title: str = "Audio Dossier",
    artist: str = "OpenClaw",
    album: str = "OpenClaw Audio Briefings",
) -> None:
    """Concatenate audio segments (WAV or MP3) and produce high-quality 128k MP3 with ID3 tags."""
    if not audio_files:
        raise ValueError("No audio segments to stitch.")

    if len(audio_files) == 1:
        cmd = [
            ffmpeg_bin,
            "-y",
            "-i",
            audio_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 single-file export failed: {res.stderr}")
        return

    # Multi-segment concat demuxer
    with tempfile.NamedTemporaryFile("w", suffix=".txt", delete=False) as list_file:
        for f in audio_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 resolve_voice(speaker_name: str, lang: str = "de", explicit_voice: str | None = None) -> str:
    """Resolve a speaker alias into a pre-selected voice name."""
    if explicit_voice:
        return explicit_voice

    lang_map = DEFAULT_VOICES.get(lang, DEFAULT_VOICES["de"])
    clean_name = speaker_name.lower().strip()

    if clean_name in lang_map:
        return lang_map[clean_name]

    # Check common female or cohost roles
    if any(k in clean_name for k in ["cohost", "female", "lara", "expert", "guest", "aoede"]):
        return lang_map.get("cohost", "Aoede")

    if any(k in clean_name for k in ["deep", "doku", "news", "charon"]):
        return lang_map.get("deep", "Charon")

    return lang_map.get("solo", "Puck")


def process_dialog(
    dialog_items: list[dict],
    ffmpeg_bin: str,
    lang: str,
    model: str,
    output_file: str,
    title: str,
    pause_duration: float = 0.40,
) -> None:
    """Generate audio for structured multi-speaker dialog (only when 'podcast' is requested)."""
    is_gemini = "gemini" in model.lower()
    ext = "wav" if is_gemini else "mp3"

    with tempfile.TemporaryDirectory(prefix="audio_dialog_") as tmpdir:
        silence_file = os.path.join(tmpdir, f"pause.{ext}")
        if is_gemini:
            generate_silence_wav(pause_duration, silence_file)
        else:
            subprocess.run([
                ffmpeg_bin, "-y", "-f", "lavfi", "-i", "anullsrc=r=24000:cl=mono",
                "-t", str(pause_duration), "-c:a", "libmp3lame", "-b:a", "64k", silence_file
            ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True)

        segment_files = []
        seg_idx = 0

        for item in dialog_items:
            speaker = item.get("speaker", "host")
            text = item.get("text", "").strip()
            if not text:
                continue

            voice = resolve_voice(speaker, lang=lang, explicit_voice=item.get("voice"))
            chunks = split_text_into_chunks(text)

            for chunk in chunks:
                chunk_file = os.path.join(tmpdir, f"seg_{seg_idx:04d}.{ext}")
                synthesize_segment(chunk, voice=voice, output_path=chunk_file, model=model)
                segment_files.append(chunk_file)
                seg_idx += 1

            # Natural pause between speaker turns
            segment_files.append(silence_file)

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

        stitch_audio_segments(
            ffmpeg_bin=ffmpeg_bin,
            audio_files=segment_files,
            output_file=output_file,
            title=title,
        )


def process_solo_text(
    text: str,
    ffmpeg_bin: str,
    voice: str,
    model: str,
    output_file: str,
    title: str,
) -> None:
    """Generate audio for a single speaker text (default for all audio dossiers)."""
    chunks = split_text_into_chunks(text)
    if not chunks:
        raise ValueError("No text provided to synthesize.")

    is_gemini = "gemini" in model.lower()
    ext = "wav" if is_gemini else "mp3"

    with tempfile.TemporaryDirectory(prefix="audio_solo_") as tmpdir:
        segment_files = []
        for i, chunk in enumerate(chunks):
            chunk_file = os.path.join(tmpdir, f"chunk_{i:04d}.{ext}")
            synthesize_segment(chunk, voice=voice, output_path=chunk_file, model=model)
            segment_files.append(chunk_file)

        stitch_audio_segments(
            ffmpeg_bin=ffmpeg_bin,
            audio_files=segment_files,
            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 (used only when 'podcast' style 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 override (Default DE: Puck, EN: Zephyr)")
    parser.add_argument("--co-voice", help="Secondary voice override for dialog (Default: Aoede)")
    parser.add_argument("--lang", default="de", choices=["de", "en"], help="Language preset (de/en)")
    parser.add_argument("--model", default=DEFAULT_MODEL, help=f"OpenRouter speech model (default: {DEFAULT_MODEL})")
    parser.add_argument("--pause", type=float, default=0.40, help="Pause in seconds between dialog turns")
    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)

    # Dialog mode only if explicit dialog file/json provided
    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,
            lang=args.lang,
            model=args.model,
            output_file=str(out_path),
            title=args.title,
            pause_duration=args.pause,
        )
    else:
        # Solo text mode (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.")

        primary_voice = args.voice or resolve_voice("solo", lang=args.lang)
        process_solo_text(
            text=raw_text,
            ffmpeg_bin=ffmpeg_bin,
            voice=primary_voice,
            model=args.model,
            output_file=str(out_path),
            title=args.title,
        )

    print(f"Generated audio: {out_path} ({out_path.stat().st_size} bytes)")
    if args.media_tag:
        print(f"MEDIA:{out_path}")


if __name__ == "__main__":
    main()