#!/usr/bin/env python3
"""Batch audio generation from a manifest file.

Processes multiple assets sequentially: SAO generation or harmonic synthesis,
followed by post-processing (trim, normalize, convert to OGG).

Usage:
    python3 audio_batch.py manifest.json [--dry-run] [--only ID,ID,...] [--skip-existing]
"""

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

from common import ensure_venv  # noqa: E402

ensure_venv()

import json
import os
import subprocess
import wave

import numpy as np


def load_manifest(path):
    with open(path) as f:
        return json.load(f)


def resolve_paths(manifest, manifest_dir):
    """Resolve output_dir and gen_dir relative to the git root."""
    # Find git root by walking up from manifest_dir
    # Check for .git as file (worktree) or directory (regular repo)
    git_root = manifest_dir
    while git_root != "/":
        if os.path.exists(os.path.join(git_root, ".git")):
            break
        git_root = os.path.dirname(git_root)
    else:
        git_root = manifest_dir

    output_dir = os.path.join(git_root, manifest.get("output_dir", "client/assets/audio"))
    gen_dir = os.path.join(git_root, manifest.get("gen_dir", "client/assets/audio/gen"))
    os.makedirs(output_dir, exist_ok=True)
    os.makedirs(gen_dir, exist_ok=True)
    return output_dir, gen_dir, git_root


def get_default(manifest, asset, key):
    """Get a value from the asset, falling back to manifest defaults."""
    defaults = manifest.get("defaults", {})
    return asset.get(key, defaults.get(key))


def run_sao_generate(asset, manifest, gen_dir, output_dir, script_dir):
    """Generate audio via Stable Audio Open + post-processing."""
    filename = asset["filename"]
    base_name = os.path.splitext(filename)[0]
    wav_path = os.path.join(gen_dir, base_name + ".wav")
    ogg_path = os.path.join(output_dir, filename)

    prompt = asset["prompt"]
    duration = asset.get("duration", 10)
    steps = get_default(manifest, asset, "steps") or 100
    cfg = get_default(manifest, asset, "cfg") or 7
    timeout = get_default(manifest, asset, "timeout") or 600

    # Run audio-generate with --post
    cmd = [
        sys.executable, os.path.join(script_dir, "audio_connector.py"),
        "generate", prompt,
        "--duration", str(duration),
        "--steps", str(steps),
        "--cfg", str(cfg),
        "--output", wav_path,
        "--output-ogg", ogg_path,
        "--timeout", str(timeout),
    ]

    result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout + 60)
    if result.returncode != 0:
        stderr = result.stderr.strip()
        try:
            err = json.loads(result.stdout)
            return {"ok": False, "error": err.get("error", stderr)}
        except (json.JSONDecodeError, ValueError):
            return {"ok": False, "error": stderr or "generation failed"}

    try:
        return json.loads(result.stdout)
    except json.JSONDecodeError:
        return {"ok": False, "error": f"Unexpected output: {result.stdout[:200]}"}


def synthesize_harmonic(params, wav_path):
    """Synthesize audio from harmonic parameters."""
    sr = 44100
    duration = params["duration"]
    fundamental = params["fundamental"]
    harmonics = params.get("harmonics", [])
    attack_ms = params.get("attack_ms", 10)
    sustain_ratio = params.get("sustain_ratio", 0.2)
    decay = params.get("decay", "exponential")

    n = int(sr * duration)
    t = np.linspace(0, duration, n, endpoint=False)

    # Fundamental
    signal = np.sin(2 * np.pi * fundamental * t)

    # Add harmonics
    for h in harmonics:
        freq = h["freq"]
        db = h["db"]
        amplitude = 10 ** (db / 20)
        signal = signal + amplitude * np.sin(2 * np.pi * freq * t)

    # Envelope: attack + sustain + decay
    attack_s = attack_ms / 1000
    attack_env = np.minimum(t / attack_s, 1.0) if attack_s > 0 else np.ones(n)

    sustain_end = duration * sustain_ratio
    if decay == "exponential":
        # Decay rate: reach -60dB by end of duration
        decay_rate = 6.9 / (duration - sustain_end) if duration > sustain_end else 10
        decay_env = np.where(t < sustain_end, 1.0, np.exp(-decay_rate * (t - sustain_end)))
    else:
        # Linear decay
        decay_env = np.where(t < sustain_end, 1.0,
                             1.0 - (t - sustain_end) / (duration - sustain_end))

    envelope = attack_env * decay_env
    signal = signal * envelope

    # Normalize to peak
    peak = np.max(np.abs(signal))
    if peak > 0:
        signal = signal / peak * 0.9

    # Write WAV
    int_samples = np.clip(signal * 32767, -32767, 32767).astype(np.int16)
    with wave.open(wav_path, "w") as f:
        f.setnchannels(1)
        f.setsampwidth(2)
        f.setframerate(sr)
        f.writeframes(int_samples.tobytes())

    return wav_path


def run_synth(asset, manifest, gen_dir, output_dir, script_dir):
    """Synthesize audio from harmonic parameters + post-process."""
    filename = asset["filename"]
    base_name = os.path.splitext(filename)[0]
    wav_path = os.path.join(gen_dir, base_name + "_synth.wav")
    ogg_path = os.path.join(output_dir, filename)

    synth_params = asset.get("synth")
    if not synth_params:
        return {"ok": False, "error": "No synth parameters provided"}

    synth_type = synth_params.get("type", "harmonic")
    if synth_type != "harmonic":
        return {"ok": False, "error": f"Unknown synth type: {synth_type}"}

    try:
        synthesize_harmonic(synth_params, wav_path)
    except Exception as e:
        return {"ok": False, "error": f"Synthesis failed: {e}"}

    # Post-process: normalize + convert (skip trim for synth — no silence to trim)
    post_script = os.path.join(script_dir, "audio_post.py")
    lufs = get_default(manifest, asset, "lufs") or -16
    quality = get_default(manifest, asset, "quality") or 6

    # Normalize
    norm_path = os.path.join(gen_dir, base_name + "_norm.wav")
    cmd = [sys.executable, post_script, "normalize", wav_path, "--output", norm_path,
           "--lufs", str(lufs)]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        return {"ok": False, "error": f"Normalize failed: {result.stderr.strip()}"}

    # Convert to OGG
    cmd = [sys.executable, post_script, "convert", norm_path, "--output", ogg_path,
           "--quality", str(quality)]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        return {"ok": False, "error": f"Convert failed: {result.stderr.strip()}"}

    # Clean up intermediate
    try:
        os.remove(norm_path)
    except OSError:
        pass

    ogg_size = os.path.getsize(ogg_path)
    return {
        "ok": True,
        "file": wav_path,
        "ogg_file": ogg_path,
        "ogg_size_bytes": ogg_size,
        "synth_params": synth_params,
        "post_processed": True,
    }


def main():
    if len(sys.argv) < 2:
        print("Usage: audio_batch.py manifest.json [--dry-run] [--only ID,ID,...] [--skip-existing]",
              file=sys.stderr)
        sys.exit(1)

    manifest_path = sys.argv[1]
    dry_run = "--dry-run" in sys.argv
    skip_existing = "--skip-existing" in sys.argv

    only_ids = None
    for i, arg in enumerate(sys.argv):
        if arg == "--only" and i + 1 < len(sys.argv):
            only_ids = set(sys.argv[i + 1].split(","))

    manifest = load_manifest(manifest_path)
    manifest_dir = os.path.dirname(os.path.abspath(manifest_path))
    output_dir, gen_dir, git_root = resolve_paths(manifest, manifest_dir)
    script_dir = os.path.dirname(os.path.abspath(__file__))

    assets = manifest.get("assets", [])
    if only_ids:
        assets = [a for a in assets if a["id"] in only_ids]

    # Health check if any SAO assets
    sao_assets = [a for a in assets if a.get("method") == "sao"]
    if sao_assets and not dry_run:
        print(f"Checking SAO API health...", file=sys.stderr)
        health_cmd = [sys.executable, os.path.join(script_dir, "audio_connector.py"), "health"]
        result = subprocess.run(health_cmd, capture_output=True, text=True)
        if result.returncode != 0:
            print(json.dumps({"ok": False, "error": "SAO API health check failed",
                              "details": result.stdout.strip()}))
            sys.exit(1)
        print(f"  SAO API is up.", file=sys.stderr)

    total = len(assets)
    results = []
    success = 0
    failed = 0
    skipped = 0

    print(f"Processing {total} assets from {os.path.basename(manifest_path)}...", file=sys.stderr)
    if dry_run:
        print("  (dry run — no generation will occur)", file=sys.stderr)

    for i, asset in enumerate(assets, 1):
        asset_id = asset["id"]
        filename = asset["filename"]
        method = asset.get("method", "sao")

        print(f"\n[{i}/{total}] {asset_id}: {filename} ({method})", file=sys.stderr)

        if skip_existing:
            ogg_path = os.path.join(output_dir, filename)
            if os.path.exists(ogg_path):
                print(f"  Skipping — already exists", file=sys.stderr)
                results.append({"id": asset_id, "status": "skipped", "reason": "exists"})
                skipped += 1
                continue

        if dry_run:
            print(f"  Would generate: {filename}", file=sys.stderr)
            if method == "sao":
                print(f"  Prompt: {asset.get('prompt', '(none)')[:80]}...", file=sys.stderr)
            elif method == "synth":
                synth = asset.get("synth", {})
                print(f"  Synth: {synth.get('fundamental')}Hz, {synth.get('duration')}s",
                      file=sys.stderr)
            results.append({"id": asset_id, "status": "dry_run"})
            continue

        if method == "sao":
            result = run_sao_generate(asset, manifest, gen_dir, output_dir, script_dir)
        elif method == "synth":
            result = run_synth(asset, manifest, gen_dir, output_dir, script_dir)
        else:
            result = {"ok": False, "error": f"Unknown method: {method}"}

        result["id"] = asset_id
        if result.get("ok"):
            success += 1
            result["status"] = "success"
            print(f"  OK → {result.get('ogg_file', filename)}", file=sys.stderr)
        else:
            failed += 1
            result["status"] = "failed"
            print(f"  FAILED: {result.get('error', 'unknown')}", file=sys.stderr)

        results.append(result)

    # Summary
    summary = {
        "ok": failed == 0,
        "total": total,
        "success": success,
        "failed": failed,
        "skipped": skipped,
        "results": results,
    }
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()