#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
import time
from pathlib import Path
from typing import Any


def _parse_args(argv: list[str]) -> argparse.Namespace:
    p = argparse.ArgumentParser(description="SongStack runner: render multiple Songcraft tracks via ACG + ACE-Step.")
    p.add_argument("--tracks", default="projects/2026-02-19_songstack/work/songstack_tracks.json", help="Path to tracks JSON.")
    p.add_argument("--project-dir", default="projects/2026-02-19_songstack", help="Project root (contains outputs/ and work/).")
    p.add_argument("--category", default="", help="Optional filter: pop|rock|synth")
    p.add_argument("--limit", type=int, default=0, help="Optional limit (0 = no limit).")
    p.add_argument("--min-bytes", type=int, default=1_000_000, help="Skip if output exists larger than this.")
    return p.parse_args(argv)


def _load_json(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text(encoding="utf-8"))


def _is_done(path: Path, min_bytes: int) -> bool:
    try:
        return path.is_file() and path.stat().st_size > min_bytes
    except FileNotFoundError:
        return False


def _notify(msg: str) -> None:
    try:
        subprocess.run(["notify", msg], check=False, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    except Exception:
        pass


def _notify_file(path: Path, caption: str) -> None:
    try:
        subprocess.run(["notify", "--file", str(path), caption], check=False, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    except Exception:
        pass


def _acg_run(resource_id: str, payload: dict[str, Any], *, repo_root: Path) -> tuple[int, str]:
    env = dict(os.environ)
    env["PYTHONPATH"] = str(repo_root) + (":" + env["PYTHONPATH"] if env.get("PYTHONPATH") else "")
    p = subprocess.run(
        ["python", "-m", "acg", "run", resource_id, "--json-in"],
        input=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        cwd=str(repo_root),
        env=env,
        check=False,
    )
    return int(p.returncode), p.stdout.decode("utf-8", errors="replace")


def main(argv: list[str] | None = None) -> int:
    args = _parse_args(argv or sys.argv[1:])
    # __file__ = projects/<id>/work/run_songstack.py
    # parents[0]=work, [1]=<id>, [2]=projects, [3]=repo root
    repo_root = Path(__file__).resolve().parents[3]
    project_dir = Path(args.project_dir).expanduser().resolve()
    tracks_path = Path(args.tracks).expanduser().resolve()

    cfg = _load_json(tracks_path)
    resource_id = str(cfg.get("engine") or "music.acestep15_local")
    audio_format_default = str(cfg.get("audio_format") or "mp3")
    tracks = cfg.get("tracks") or []
    if not isinstance(tracks, list) or not tracks:
        raise SystemExit("tracks.json has no tracks[]")

    out_root = (project_dir / "outputs").resolve()
    work_root = (project_dir / "work").resolve()
    songcraft_root = (work_root / "songcraft").resolve()
    log_path = (work_root / "generation_log.jsonl").resolve()
    out_root.mkdir(parents=True, exist_ok=True)
    songcraft_root.mkdir(parents=True, exist_ok=True)

    want_cat = (args.category or "").strip().lower()
    if want_cat and want_cat not in {"pop", "rock", "synth"}:
        raise SystemExit("--category must be pop|rock|synth")

    selected: list[dict[str, Any]] = []
    for t in tracks:
        if not isinstance(t, dict):
            continue
        cat = str(t.get("category") or "").strip().lower()
        if want_cat and cat != want_cat:
            continue
        selected.append(t)

    if args.limit and args.limit > 0:
        selected = selected[: int(args.limit)]

    msg0 = f"SongStack started: {len(selected)} track(s) ({want_cat or 'all'})"
    print(msg0, flush=True)
    _notify(msg0)

    for idx, t in enumerate(selected, start=1):
        category = str(t.get("category") or "misc").strip().lower() or "misc"
        stem = str(t["stem"]).strip()
        title = str(t.get("title") or stem).strip()
        caption = str(t["caption"])
        lyrics = str(t["lyrics"])

        out_dir = (out_root / category / "ACESTep").resolve()
        out_dir.mkdir(parents=True, exist_ok=True)
        out_path = (out_dir / f"{stem}.{str(t.get('audio_format') or audio_format_default)}").resolve()

        # Save Songcraft artifacts (inputs) for reproducibility.
        sc_dir = (songcraft_root / category / stem).resolve()
        sc_dir.mkdir(parents=True, exist_ok=True)
        (sc_dir / "caption.txt").write_text(caption.rstrip() + "\n", encoding="utf-8")
        (sc_dir / "lyrics.txt").write_text(lyrics.rstrip() + "\n", encoding="utf-8")
        spec = {
            "title": title,
            "category": category,
            "stem": stem,
            "caption": caption,
            "lyrics": lyrics,
        }
        for k in ("bpm", "keyscale", "timesignature"):
            if t.get(k) is not None and str(t.get(k)).strip():
                spec[k] = t.get(k)
        (sc_dir / "song_spec.json").write_text(json.dumps(spec, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")

        if _is_done(out_path, args.min_bytes):
            print(f"[skip] {idx}/{len(selected)} {category}/{stem}", flush=True)
            rec = {"ts": time.time(), "stem": stem, "title": title, "category": category, "status": "skipped", "out": str(out_path)}
            log_path.parent.mkdir(parents=True, exist_ok=True)
            with log_path.open("a", encoding="utf-8") as fp:
                fp.write(json.dumps(rec, ensure_ascii=False) + "\n")
            continue

        payload: dict[str, Any] = {
            "caption": caption,
            "lyrics": lyrics,
            "out": str(out_path),
            "audio_format": str(t.get("audio_format") or audio_format_default),
        }
        for k in ("bpm", "keyscale", "timesignature"):
            if t.get(k) is not None and str(t.get(k)).strip():
                payload[k] = t.get(k)

        started = time.time()
        msg = f"[{idx}/{len(selected)}] Generating: {category}/{title}"
        print(msg, flush=True)
        _notify(msg)
        code, raw = _acg_run(resource_id, payload, repo_root=repo_root)
        elapsed = time.time() - started
        ok = (code == 0) and out_path.exists() and out_path.stat().st_size > 10_000

        rec = {
            "ts": time.time(),
            "stem": stem,
            "title": title,
            "category": category,
            "status": "ok" if ok else "error",
            "elapsed_s": round(elapsed, 3),
            "out": str(out_path),
            "resource_id": resource_id,
            "exit_code": code,
        }
        log_path.parent.mkdir(parents=True, exist_ok=True)
        with log_path.open("a", encoding="utf-8") as fp:
            fp.write(json.dumps(rec, ensure_ascii=False) + "\n")

        if ok:
            print(f"[ok] {idx}/{len(selected)} {category}/{stem} -> {out_path}", flush=True)
            _notify_file(out_path, f"SongStack: {category} — {title}")
            continue

        # Failure: write log tail artifact and notify it.
        print(f"[error] {idx}/{len(selected)} {category}/{stem} (exit_code={code})", flush=True)
        tail_path = (work_root / f"{stem}.acg_tail.txt").resolve()
        tail = "\n".join((raw or "").splitlines()[-120:])
        tail_path.write_text(tail + "\n", encoding="utf-8")
        _notify_file(tail_path, f"SongStack FAILED: {category} — {title}")

    print("SongStack finished.", flush=True)
    _notify("SongStack finished.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())