#!/usr/bin/env python3
"""
Minimax TTS narration generator — deterministic CLI for hyper-animator.

Reads a narration script JSON, calls Minimax TTS API per scene, downloads
WAV files. Each scene gets its own audio clip for precise HTML sync.

Usage:
    # List available voices (system + cloned + generated)
    python3 scripts/tts-gen.py --list-voices

    # Generate narration audio
    python3 scripts/tts-gen.py \\
        --script narration.json \\
        --voice XiaoR_001 \\
        -o hyperframes-output/assets/
"""
import argparse
import json
import os
import pathlib
import subprocess
import sys
import tempfile
import time
import urllib.request
import re
import wave

SR = 32000
TTS_API_PATH = "/v1/t2a_v2"

# ── .env loading ────────────────────────────────────────────

def load_env():
    env_file = pathlib.Path.home() / ".claude" / "skills" / "hyper-animator" / ".env"
    if env_file.exists():
        for line in env_file.read_text().splitlines():
            line = line.strip()
            if line and not line.startswith("#") and "=" in line:
                k, v = line.split("=", 1)
                v = v.strip().strip('"').strip("'")
                if k.strip() not in os.environ:
                    os.environ[k.strip()] = v

# ── Voice listing ────────────────────────────────────────────

def list_voices(api_key, api_host):
    """Fetch and print available voices grouped by type."""
    body = json.dumps({"voice_type": "all"})
    tmp = tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False)
    tmp.write(body)
    tmp.close()

    try:
        proc = subprocess.run([
            "curl", "-s", "--http1.1", "--connect-timeout", "10", "--max-time", "30",
            "-X", "POST", f"https://{api_host}/v1/get_voice",
            "-H", f"Authorization: Bearer {api_key}",
            "-H", "Content-Type: application/json",
            "-d", f"@{tmp.name}"
        ], capture_output=True, text=True, timeout=35)
        os.unlink(tmp.name)

        resp = json.loads(proc.stdout.strip())
        base = resp.get("base_resp", {})
        if base.get("status_code") != 0:
            print(f"Voice list API error: {base.get('status_msg', 'unknown')}")
            return None

        # Group voices
        groups = {}
        for cat in ["system_voice", "voice_cloning", "voice_generation"]:
            voices = resp.get(cat, [])
            if voices:
                groups[cat] = voices

        return groups

    except Exception as e:
        print(f"Voice list error: {e}")
        os.unlink(tmp.name)
        return None

def print_voices(groups):
    """Print voice list grouped by type."""
    labels = {
        "system_voice": "System Voices",
        "voice_cloning": "Cloned Voices (your group)",
        "voice_generation": "AI-Generated Voices",
    }
    for cat, label in labels.items():
        voices = groups.get(cat, [])
        if not voices:
            continue
        print(f"\n── {label} ({len(voices)}) ──")
        for v in voices[:50 if cat != "system_voice" else 10]:
            vid = v.get("voice_id", "?")
            name = v.get("voice_name", "") or vid
            desc = ", ".join(v.get("description", [])[:2]) or ""
            print(f"  {vid}")
            if name != vid:
                print(f"    name: {name}")
            if desc:
                print(f"    desc: {desc[:80]}")

# ── Audio normalization ──────────────────────────────────────

def _normalize_wav(path, target_db=-3):
    """Normalize WAV to target dB peak."""
    import struct
    try:
        with wave.open(path, 'rb') as w:
            params = w.getparams()
            frames = w.readframes(params.nframes)
        samples = struct.unpack(f'<{params.nframes}h', frames)
        peak = max(abs(s) for s in samples) if samples else 1
        if peak == 0:
            return
        target_peak = int(32767 * (10 ** (target_db / 20)))
        scale = target_peak / peak
        normalized = [int(max(-32768, min(32767, s * scale))) for s in samples]
        with wave.open(path, 'wb') as w:
            w.setparams(params)
            w.writeframes(struct.pack(f'<{len(normalized)}h', *normalized))
    except Exception as e:
        print(f"  WARNING: normalize failed for {path}: {e}", file=sys.stderr)


def _call_tts_api(api_key, api_host, body_dict, timeout=90):
    """Call Minimax TTS API via urllib — no curl, no tempfile, no encoding bug."""
    import urllib.error as urlerror
    api_url = f"https://{api_host}{TTS_API_PATH}"
    data = json.dumps(body_dict).encode('utf-8')
    req = urllib.request.Request(
        api_url, data=data,
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json; charset=utf-8"
        },
        method="POST"
    )
    try:
        resp = urllib.request.urlopen(req, timeout=timeout)
        return json.loads(resp.read().decode('utf-8'))
    except urlerror.HTTPError as e:
        body = e.read().decode('utf-8')
        return json.loads(body) if body else {"base_resp": {"status_code": e.code, "status_msg": str(e)}}
    except Exception as e:
        return {"base_resp": {"status_code": -1, "status_msg": str(e)}}


def _split_into_chunks(text, max_chars=80):
    """Split Chinese text at sentence boundaries."""
    parts = re.split(r'(?<=[。！？；\n])', text)
    chunks = []
    current = ""
    for part in parts:
        part = part.strip()
        if not part:
            continue
        if len(current) + len(part) <= max_chars:
            current += part
        else:
            if current:
                chunks.append(current)
            current = part
    if current:
        chunks.append(current)
    return chunks if chunks else [text]


def _concat_wavs(input_paths, output_path):
    """Concatenate multiple WAV files into one."""
    all_frames = []
    params = None
    for p in input_paths:
        with wave.open(p, 'rb') as w:
            if params is None:
                params = w.getparams()
            all_frames.append(w.readframes(w.getnframes()))
    with wave.open(output_path, 'wb') as w:
        w.setparams(params)
        for frames in all_frames:
            w.writeframes(frames)


# ── Main ─────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="Minimax TTS narration generator")
    parser.add_argument("--script", help="Path to narration JSON")
    parser.add_argument("--voice", default="", help="Voice ID (required — selected by user in Round 3)")
    parser.add_argument("--emotion", default="calm", help="Voice emotion: happy, sad, angry, fearful, disgusted, surprised, calm, fluent (default: calm)")
    parser.add_argument("--speed", type=float, default=1.0, help="Speech speed 0.5-2.0 (default: 1.0)")
    parser.add_argument("-o", "--output", help="Output directory for scene WAVs")
    parser.add_argument("--list-voices", action="store_true", help="List available voices and exit")
    args = parser.parse_args()

    load_env()
    api_key = os.environ.get("MINIMAX_API_KEY", "")
    api_host = os.environ.get("MINIMAX_API_HOST", "api.minimaxi.com")

    if not api_key:
        print("MINIMAX_API_KEY not configured. Exit.")
        sys.exit(1)

    # ── List voices mode ──
    if args.list_voices:
        groups = list_voices(api_key, api_host)
        if groups:
            print_voices(groups)
            total = sum(len(v) for v in groups.values())
            print(f"\nTotal: {total} voices available")
        sys.exit(0 if groups else 1)

    # ── Generate mode ──
    if not args.script or not args.output:
        print("Error: --script and -o are required for generation mode.")
        print("Use --list-voices to see available voices.")
        sys.exit(1)

    # Load script
    with open(args.script, 'r') as f:
        script = json.load(f)

    voice = args.voice or script.get("voice")
    if not voice:
        print("Error: voice is required. Use --voice <id> or set 'voice' in narration.json.", file=sys.stderr)
        sys.exit(1)
    if script.get("voice") and args.voice and script.get("voice") != args.voice:
        print(f"  WARNING: CLI --voice '{args.voice}' overrides narration.json voice '{script.get('voice')}'", file=sys.stderr)
    scenes = script.get("scenes", [])

    if not scenes:
        print("No scenes found in script. Exit.")
        sys.exit(1)

    os.makedirs(args.output, exist_ok=True)
    total = len(scenes)
    success = 0
    failed = 0

    for i, scene in enumerate(scenes):
        scene_num = scene.get("scene", i + 1)
        title = scene.get("title", f"scene-{scene_num}")
        text = scene.get("narration", "").strip()
        # Strip invisible/control characters that cause TTS gibberish
        text = ''.join(c for c in text if c.isprintable() or c in '\n\r\t')
        text = text.strip()
        # Insert pause tags at sentence boundaries — helps model handle long text
        text = re.sub(r'([。！？；\n])\s*', r'\1<#0.3#>', text)
        text = re.sub(r'<#0\.3#>$', '', text)  # Remove trailing pause tag

        if not text:
            print(f"  [{scene_num}/{total}] {title}: EMPTY — skipped")
            continue

        print(f"  [{scene_num}/{total}] {title}: \"{text[:50]}...\" ({len(text)} chars)")

        # Split into chunks for sentence-level TTS generation
        chunks = _split_into_chunks(text)

        chunk_paths = []
        chunk_failed = False
        for ci, chunk_text in enumerate(chunks):
            vs = {"voice_id": voice, "speed": args.speed, "vol": 1.0, "pitch": 0}
            em = scene.get("emotion", args.emotion)
            if em:
                vs["emotion"] = em

            resp = _call_tts_api(api_key, api_host, {
                "text": chunk_text,
                "stream": False,
                "model": "speech-2.8-hd",
                "voice_setting": vs,
                "audio_setting": {
                    "format": "wav",
                    "sample_rate": SR,
                    "bitrate": 128000,
                    "channel": 1
                },
                "language_boost": "auto",
                "output_format": "url"
            })

            base = resp.get("base_resp", {})
            status = base.get("status_code", -1)

            if status == 0:
                audio_url = resp["data"]["audio"]
                chunk_path = os.path.join(args.output, f"scene-{scene_num}-chunk-{ci}.wav")
                urllib.request.urlretrieve(audio_url, chunk_path)
                _normalize_wav(chunk_path)
                chunk_paths.append(chunk_path)
            elif status == 1002:
                print("    Rate limited, retrying in 3s...")
                time.sleep(3)
                vs = {"voice_id": voice, "speed": args.speed, "vol": 1.0, "pitch": 0}
                em = scene.get("emotion", args.emotion)
                if em:
                    vs["emotion"] = em
                resp = _call_tts_api(api_key, api_host, {
                    "text": chunk_text,
                    "stream": False,
                    "model": "speech-2.8-hd",
                    "voice_setting": vs,
                    "audio_setting": {
                        "format": "wav",
                        "sample_rate": SR,
                        "bitrate": 128000,
                        "channel": 1
                    },
                    "language_boost": "auto",
                    "output_format": "url"
                })
                if resp.get("base_resp", {}).get("status_code") == 0:
                    audio_url = resp["data"]["audio"]
                    chunk_path = os.path.join(args.output, f"scene-{scene_num}-chunk-{ci}.wav")
                    urllib.request.urlretrieve(audio_url, chunk_path)
                    _normalize_wav(chunk_path)
                    chunk_paths.append(chunk_path)
                else:
                    print(f"    Failed on retry for chunk {ci}")
                    chunk_failed = True
                    break
            else:
                msg = base.get("status_msg", "")
                print(f"    Status {status}: {msg} — chunk {ci} failed")
                chunk_failed = True
                break

        if chunk_failed:
            failed += 1
            for p in chunk_paths:
                try:
                    os.unlink(p)
                except:
                    pass
            continue

        # All chunks succeeded — concatenate into final scene audio
        out_path = os.path.join(args.output, f"scene-{scene_num}.wav")
        _concat_wavs(chunk_paths, out_path)

        # Calculate total duration from concatenated WAV
        with wave.open(out_path, 'rb') as w:
            dur_ms = w.getnframes() / w.getframerate() * 1000

        sub_path = os.path.join(args.output, f"scene-{scene_num}-timing.json")
        subtitles = scene.get("subtitles", [])
        # If subtitles have absolute timing (from narration.json), write them directly.
        # Otherwise, compute relative timing based on chunk durations.
        if subtitles:
            timing = []
            for sub in subtitles:
                timing.append({
                    "text": sub.get("text", ""),
                    "start": sub.get("start", 0),
                    "end": sub.get("end", 0)
                })
        else:
            timing = [{"text": text, "start": 0, "end": dur_ms / 1000}]
        with open(sub_path, 'w', encoding='utf-8') as sf:
            json.dump(timing, sf, indent=2, ensure_ascii=False)
        print(f"    OK: {dur_ms/1000:.1f}s, {len(chunk_paths)} chunk(s) -> {out_path}")

        # Clean up individual chunk files
        for p in chunk_paths:
            try:
                os.unlink(p)
            except:
                pass

        success += 1

    print(f"\nTTS complete: {success}/{total} scenes generated, {failed} failed")
    sys.exit(0 if failed == 0 else 1)

if __name__ == "__main__":
    main()