#!/usr/bin/env python3
"""foundry — the Sample Foundry driver (URL → audio → stems).

The Sample Foundry turns a track URL into separated stems ready for sampling.
One engine, two faces: this CLI and the GUI (`foundry serve`) both call the
same `engine/` steps. Run under system python3 (pydantic); the heavy lifting
shells out to dedicated venvs (demucs today, Roformer next).

    python3 foundry.py fetch <url>             # download bestaudio → a Catch
    python3 foundry.py sep <slug|url>          # separate a catch into stems
    python3 foundry.py catch <url>             # fetch + sep (the full pipeline)
    python3 foundry.py list                    # list catches in the workspace
    python3 foundry.py backends                # show separation backends
    python3 foundry.py link <kit>              # link Samples/<kit> into Dirt-Samples
    python3 foundry.py serve [--port 8765]     # launch the web GUI

Workspace defaults to $FOUNDRY_HOME or ~/Downloads/foundry; override with
--workspace. engine1 = demucs/htdemucs. engine2 (Roformer) slots in behind the
same --backend flag — see README.
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

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

from engine import fetch as F  # noqa: E402
from engine import separate as S  # noqa: E402
from engine.model import Catch  # noqa: E402


def _resolve_catch(ref: str, workspace: Path) -> Catch:
    """A ref is either an existing slug dir or a URL to fetch first."""
    d = workspace / ref
    if (d / "catch.json").exists():
        return Catch.load(d)
    if "://" in ref:
        return F.fetch(ref, workspace)
    raise SystemExit(f"no catch '{ref}' in {workspace} (and not a URL)")


def cmd_fetch(a):
    F.fetch(a.url, a.workspace)


def cmd_sep(a):
    ws = a.workspace or F.default_workspace()
    catch = _resolve_catch(a.ref, ws)
    S.separate(
        catch, ws, backend=a.backend,
        model=a.model, shifts=a.shifts, jobs=a.jobs, device=a.device,
        progress=lambda p: print(f"\r  {a.backend}: {p:3d}%", end="", flush=True),
    )
    print()


def cmd_catch(a):
    ws = a.workspace or F.default_workspace()
    catch = F.fetch(a.url, ws)
    S.separate(
        catch, ws, backend=a.backend,
        model=a.model, shifts=a.shifts, jobs=a.jobs, device=a.device,
        progress=lambda p: print(f"\r  {a.backend}: {p:3d}%", end="", flush=True),
    )
    print()


def cmd_list(a):
    ws = a.workspace or F.default_workspace()
    if not ws.exists():
        print(f"(empty: {ws})")
        return
    for d in sorted(p for p in ws.iterdir() if (p / "catch.json").exists()):
        c = Catch.load(d)
        mark = "▣" if c.separated else "▢"
        meta = f"{c.backend}/{c.model}" if c.separated else "not separated"
        print(f"{mark} {c.slug:42.42}  {meta}")


def cmd_backends(a):
    for name, be in S.REGISTRY.items():
        ok, hint = be.available()
        print(f"{'✓' if ok else '✗'} {name:10} {hint}")


def cmd_link(a):
    from engine import publish as P
    r = P.link_kit(a.kit)
    print(f"{r.action}: {r.link}  →  {r.src}")
    if r.action == "skipped-real-dir":
        raise SystemExit(f"⚠ {r.link} is a real directory, not a symlink — left untouched")


def cmd_demos(a):
    """List the sampling-classic demos, or precompute them (fetch [+ separate])."""
    from engine import demos as D
    items = D.load()
    if not (a.catch or a.all):
        for d in items:
            v = "✓" if d.get("verified") else "≈"   # ≈ = ytsearch-resolved
            print(f"{v} {d['slug']:24.24} {d['title']:46.46} [{','.join(d['tags'])}]")
        print(f"\n{len(items)} demos · `foundry demos --catch <slug>` or `--all` to precompute"
              f" (add --sep to separate too).  ✓ verified url · ≈ ytsearch top-hit")
        return
    ws = a.workspace or F.default_workspace()
    targets = items if a.all else [D.by_slug(a.catch)]
    if targets == [None]:
        raise SystemExit(f"no demo {a.catch!r} — run `foundry demos` to list")
    for d in targets:
        print(f"\n▶ {d['slug']} — {d['title']}")
        catch = F.fetch(d["source"], ws)
        if a.sep:
            S.separate(catch, ws, backend=a.backend, model=a.model, shifts=a.shifts,
                       jobs=a.jobs, device=a.device,
                       progress=lambda p: print(f"\r  {a.backend}: {p:3d}%", end="", flush=True))
            print()


def cmd_serve(a):
    import server
    server.run(workspace=a.workspace or F.default_workspace(), port=a.port)


def main(argv=None):
    p = argparse.ArgumentParser(prog="foundry", description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--workspace", type=Path, default=None,
                   help="catch workspace (default $FOUNDRY_HOME or ~/Downloads/foundry)")
    sub = p.add_subparsers(dest="cmd", required=True)

    def add_sep_opts(sp):
        sp.add_argument("--backend", default="demucs")
        sp.add_argument("--model", default=None, help="backend model (default per backend)")
        sp.add_argument("--shifts", type=int, default=4, help="demucs equivariant shifts")
        sp.add_argument("--jobs", type=int, default=None)
        sp.add_argument("--device", default=None, help="cuda | cpu")

    sp = sub.add_parser("fetch"); sp.add_argument("url"); sp.set_defaults(fn=cmd_fetch)
    sp = sub.add_parser("sep"); sp.add_argument("ref"); add_sep_opts(sp); sp.set_defaults(fn=cmd_sep)
    sp = sub.add_parser("catch"); sp.add_argument("url"); add_sep_opts(sp); sp.set_defaults(fn=cmd_catch)
    sub.add_parser("list").set_defaults(fn=cmd_list)
    sub.add_parser("backends").set_defaults(fn=cmd_backends)
    sp = sub.add_parser("link"); sp.add_argument("kit"); sp.set_defaults(fn=cmd_link)
    sp = sub.add_parser("demos", help="list/precompute sampling-classic demo sources")
    sp.add_argument("--catch", metavar="SLUG", help="fetch one demo by slug")
    sp.add_argument("--all", action="store_true", help="fetch every demo")
    sp.add_argument("--sep", action="store_true", help="also separate (with --backend etc.)")
    add_sep_opts(sp); sp.set_defaults(fn=cmd_demos)
    sp = sub.add_parser("serve"); sp.add_argument("--port", type=int, default=8765); sp.set_defaults(fn=cmd_serve)

    a = p.parse_args(argv)
    a.fn(a)


if __name__ == "__main__":
    main()
