#!/usr/bin/env python3
"""
Stable Audio Open connector — Gradio API wrapper.

Talks to the Stable Audio Open Gradio app at tower-of-joy:11500.
Uses the async Gradio API pattern: POST to submit, SSE stream for results.

Usage:
    python3 audio_connector.py generate "prompt text" [--duration 10] [--steps 100] [--cfg 7] [--output file.wav] [--post]
    python3 audio_connector.py health
"""

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 shutil
import subprocess
import time
import urllib.error
import urllib.request

CONFIG_PATH = os.path.join(os.path.dirname(__file__), "config.json")

def load_config():
    with open(CONFIG_PATH) as f:
        return json.load(f)

def get_base_url():
    config = load_config()
    return config.get("stable_audio_url", "http://tower-of-joy:11500")

def health():
    """Check if the Stable Audio API is reachable."""
    base = get_base_url()
    try:
        req = urllib.request.Request(f"{base}/config", method="GET")
        with urllib.request.urlopen(req, timeout=10) as resp:
            data = json.loads(resp.read())
            version = data.get("version", "unknown")
            # Extract component info for the generate endpoint
            api_names = []
            for dep in data.get("dependencies", []):
                name = dep.get("api_name", "")
                if name and not name.startswith("js_"):
                    api_names.append(name)
            print(json.dumps({
                "ok": True,
                "url": base,
                "gradio_version": version,
                "api_endpoints": api_names
            }, indent=2))
    except Exception as e:
        print(json.dumps({
            "ok": False,
            "url": base,
            "error": str(e)
        }, indent=2))
        sys.exit(1)

def post_process(wav_path, ogg_path=None, lufs=-16, quality=6, threshold=-50):
    """Run trim + normalize + convert on a WAV file via audio-post pipeline."""
    script_dir = os.path.dirname(os.path.abspath(__file__))
    post_script = os.path.join(script_dir, "audio_post.py")
    if ogg_path is None:
        ogg_path = os.path.splitext(wav_path)[0] + ".ogg"
    cmd = [
        sys.executable, post_script, "pipeline", wav_path,
        "--output", ogg_path,
        "--lufs", str(lufs),
        "--quality", str(quality),
        "--threshold", str(threshold),
    ]
    print(f"  Post-processing → {os.path.basename(ogg_path)}...", file=sys.stderr)
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        return {"ok": False, "error": f"Post-processing failed: {result.stderr.strip()}"}
    try:
        return json.loads(result.stdout)
    except json.JSONDecodeError:
        return {"ok": True, "output": ogg_path}


def _check_available(base):
    """Quick check if Stable Audio is reachable. Fail fast with a clear message."""
    try:
        req = urllib.request.Request(f"{base}/config", method="GET")
        urllib.request.urlopen(req, timeout=5)
    except Exception:
        print(json.dumps({
            "ok": False,
            "error": f"Stable Audio is not available at {base}. The service may be switched off to save system resources. Start it before generating audio."
        }, indent=2))
        sys.exit(1)


def generate(prompt, duration=10.0, steps=100, cfg=7.0, output=None, timeout=600,
             post=False, output_ogg=None):
    """
    Generate audio from a text prompt.

    Args:
        prompt: Text description of the audio to generate
        duration: Duration in seconds (0-47, default 10)
        steps: Number of diffusion steps (default 100, lower = faster but lower quality)
        cfg: Classifier-free guidance scale (default 7)
        output: Output file path (default: auto-named in current directory)
        timeout: Maximum wait time in seconds (default 600 = 10 minutes)
        post: If True, run trim+normalize+convert after generation
        output_ogg: OGG output path when post=True (default: same basename .ogg)
    """
    base = get_base_url()
    _check_available(base)
    api_url = f"{base}/gradio_api/call/generate_audio"

    if output is None:
        # Auto-name: sanitize prompt to a filename
        safe = "".join(c if c.isalnum() or c in "-_ " else "" for c in prompt[:40])
        safe = safe.strip().replace(" ", "_").lower()
        output = f"{safe}_{int(duration)}s.wav"

    # Step 1: Submit the generation request
    payload = json.dumps({"data": [prompt, duration, steps, cfg]})
    req = urllib.request.Request(
        api_url,
        data=payload.encode(),
        headers={"Content-Type": "application/json"},
        method="POST"
    )

    print(f"Submitting generation request...", file=sys.stderr)
    print(f"  Prompt: {prompt}", file=sys.stderr)
    print(f"  Duration: {duration}s, Steps: {steps}, CFG: {cfg}", file=sys.stderr)

    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            result = json.loads(resp.read())
            event_id = result.get("event_id")
            if not event_id:
                print(json.dumps({"ok": False, "error": "No event_id returned", "response": result}, indent=2))
                sys.exit(1)
    except Exception as e:
        print(json.dumps({"ok": False, "error": f"Submit failed: {e}"}, indent=2))
        sys.exit(1)

    print(f"  Event ID: {event_id}", file=sys.stderr)
    print(f"  Waiting for generation (timeout: {timeout}s)...", file=sys.stderr)

    # Step 2: Poll the SSE stream for results
    stream_url = f"{api_url}/{event_id}"
    start_time = time.time()
    result_data = None
    last_status = None

    while time.time() - start_time < timeout:
        try:
            req = urllib.request.Request(stream_url, method="GET")
            with urllib.request.urlopen(req, timeout=timeout) as resp:
                # Read SSE events
                current_event = None
                for line_bytes in resp:
                    line = line_bytes.decode("utf-8").strip()

                    if line.startswith("event: "):
                        current_event = line[7:]
                    elif line.startswith("data: ") and current_event:
                        data_str = line[6:]

                        if current_event == "heartbeat":
                            elapsed = int(time.time() - start_time)
                            if elapsed % 30 == 0 and elapsed > 0:
                                print(f"  Still generating... ({elapsed}s elapsed)", file=sys.stderr)
                            continue

                        if current_event == "error":
                            error_msg = data_str
                            try:
                                error_msg = json.loads(data_str)
                            except (json.JSONDecodeError, TypeError):
                                pass
                            print(json.dumps({"ok": False, "error": "Generation failed", "details": error_msg}, indent=2))
                            sys.exit(1)

                        if current_event == "complete":
                            try:
                                result_data = json.loads(data_str)
                            except json.JSONDecodeError:
                                result_data = data_str
                            break

                        if current_event == "progress":
                            try:
                                progress = json.loads(data_str)
                                # Gradio progress events vary; log what we get
                                status = str(progress)[:80]
                                if status != last_status:
                                    print(f"  Progress: {status}", file=sys.stderr)
                                    last_status = status
                            except (json.JSONDecodeError, TypeError):
                                pass
                            continue

                if result_data is not None:
                    break

        except urllib.error.URLError as e:
            # Connection dropped — retry after brief pause
            time.sleep(2)
            continue
        except Exception as e:
            print(json.dumps({"ok": False, "error": f"Stream error: {e}"}, indent=2))
            sys.exit(1)

    if result_data is None:
        print(json.dumps({"ok": False, "error": f"Generation timed out after {timeout}s"}, indent=2))
        sys.exit(1)

    # Step 3: Download the audio file
    # Gradio returns file info in the data array
    elapsed = round(time.time() - start_time, 1)
    print(f"  Generation complete ({elapsed}s)", file=sys.stderr)

    try:
        # result_data is typically [{"path": "...", "url": "...", ...}] or similar
        if isinstance(result_data, list) and len(result_data) > 0:
            audio_info = result_data[0]
        elif isinstance(result_data, dict) and "data" in result_data:
            audio_info = result_data["data"][0] if result_data["data"] else None
        else:
            audio_info = result_data

        # Extract the file URL
        file_url = None
        if isinstance(audio_info, dict):
            file_url = audio_info.get("url") or audio_info.get("path")
        elif isinstance(audio_info, str):
            file_url = audio_info

        if not file_url:
            print(json.dumps({
                "ok": False,
                "error": "Could not extract audio URL from response",
                "response": result_data
            }, indent=2))
            sys.exit(1)

        # Handle relative URLs
        if file_url.startswith("/"):
            file_url = f"{base}{file_url}"
        elif not file_url.startswith("http"):
            file_url = f"{base}/file={file_url}"

        # Download the file
        print(f"  Downloading to {output}...", file=sys.stderr)
        req = urllib.request.Request(file_url, method="GET")
        with urllib.request.urlopen(req, timeout=60) as resp:
            with open(output, "wb") as f:
                shutil.copyfileobj(resp, f)

        file_size = os.path.getsize(output)
        result_json = {
            "ok": True,
            "file": output,
            "size_bytes": file_size,
            "duration_requested": duration,
            "steps": steps,
            "cfg": cfg,
            "prompt": prompt,
            "generation_time_s": elapsed
        }

        if post:
            post_result = post_process(output, ogg_path=output_ogg)
            if not post_result.get("ok"):
                result_json["post_processed"] = False
                result_json["post_error"] = post_result.get("error", "unknown")
            else:
                result_json["post_processed"] = True
                result_json["ogg_file"] = post_result.get("output", output_ogg)
                ogg_size = os.path.getsize(result_json["ogg_file"])
                result_json["ogg_size_bytes"] = ogg_size

        print(json.dumps(result_json, indent=2))

    except Exception as e:
        print(json.dumps({
            "ok": False,
            "error": f"Download failed: {e}",
            "response": result_data
        }, indent=2))
        sys.exit(1)

def main():
    if len(sys.argv) < 2:
        print("Usage:")
        print("  audio_connector.py health")
        print("  audio_connector.py generate 'prompt' [--duration N] [--steps N] [--cfg N] [--output file.wav] [--timeout N]")
        sys.exit(1)

    cmd = sys.argv[1]

    if cmd == "health":
        health()
    elif cmd == "generate":
        if len(sys.argv) < 3:
            print("Error: prompt required", file=sys.stderr)
            print("Usage: audio_connector.py generate 'prompt' [--duration N] [--steps N] [--cfg N] [--output file.wav]", file=sys.stderr)
            sys.exit(1)

        prompt = sys.argv[2]
        duration = 10.0
        steps = 100
        cfg = 7.0
        output = None
        timeout = 600
        post = False
        output_ogg = None

        # Parse optional args
        i = 3
        while i < len(sys.argv):
            if sys.argv[i] == "--duration" and i + 1 < len(sys.argv):
                duration = float(sys.argv[i + 1])
                i += 2
            elif sys.argv[i] == "--steps" and i + 1 < len(sys.argv):
                steps = int(sys.argv[i + 1])
                i += 2
            elif sys.argv[i] == "--cfg" and i + 1 < len(sys.argv):
                cfg = float(sys.argv[i + 1])
                i += 2
            elif sys.argv[i] == "--output" and i + 1 < len(sys.argv):
                output = sys.argv[i + 1]
                i += 2
            elif sys.argv[i] == "--timeout" and i + 1 < len(sys.argv):
                timeout = int(sys.argv[i + 1])
                i += 2
            elif sys.argv[i] == "--post":
                post = True
                i += 1
            elif sys.argv[i] == "--output-ogg" and i + 1 < len(sys.argv):
                output_ogg = sys.argv[i + 1]
                post = True  # --output-ogg implies --post
                i += 2
            else:
                print(f"Unknown argument: {sys.argv[i]}", file=sys.stderr)
                sys.exit(1)

        generate(prompt, duration=duration, steps=steps, cfg=cfg, output=output,
                 timeout=timeout, post=post, output_ogg=output_ogg)
    else:
        print(f"Unknown command: {cmd}", file=sys.stderr)
        sys.exit(1)

if __name__ == "__main__":
    main()