#!/usr/bin/env python3
"""server — the Sample Foundry GUI backend (stdlib http.server + Range).

Same engine as the CLI, a different face. Long-running steps (download,
separate) run as background *jobs* with pollable progress; stems are served
with HTTP Range (206) so the in-browser <audio> seeks correctly — the same
trick armada/serve.py uses for phone audition.

    python3 foundry.py serve --port 8765
    # then open http://127.0.0.1:8765/

API (all JSON unless noted):
    GET  /api/backends            → [{name, available, hint}]
    GET  /api/catches             → [Catch, …]
    POST /api/fetch {url}         → {job}
    POST /api/separate {slug,backend,shifts,model,jobs,device} → {job}
    GET  /api/job/<id>            → {id, kind, state, progress, msg, slug?}
    GET  /media/<slug>/<path>     → raw audio (Range-capable)
"""
from __future__ import annotations

import json
import os
import re
import socket
import threading
import uuid
from functools import partial
from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from urllib.parse import unquote, urlparse

HERE = Path(__file__).resolve().parent
import sys
sys.path.insert(0, str(HERE))
from engine import fetch as F  # noqa: E402
from engine import separate as S  # noqa: E402
from engine import loops as L  # noqa: E402
from engine.model import Catch  # noqa: E402

UI = HERE / "ui"

# ── job registry (in-memory; the workspace folder is the durable truth) ──────
JOBS: dict[str, dict] = {}
JLOCK = threading.Lock()


def _job(kind: str, **extra) -> dict:
    j = {"id": uuid.uuid4().hex[:12], "kind": kind, "state": "running",
         "progress": 0, "msg": "", **extra}
    with JLOCK:
        JOBS[j["id"]] = j
    return j


def _run_fetch(job, workspace, url):
    try:
        job["msg"] = "probing"
        catch = F.fetch(url, workspace, log=lambda m: job.update(msg=m))
        job.update(state="done", progress=100, slug=catch.slug, msg="caught")
    except Exception as e:  # noqa: BLE001
        job.update(state="error", msg=str(e))


def _run_separate(job, workspace, slug, opts):
    try:
        catch = Catch.load(workspace / slug)
        S.separate(catch, workspace, progress=lambda p: job.update(progress=p),
                   log=lambda m: job.update(msg=m), **opts)
        job.update(state="done", progress=100, slug=slug, msg="separated")
    except Exception as e:  # noqa: BLE001
        job.update(state="error", msg=str(e))


def _run_loops(job, workspace, slug, bars, top_n, mode):
    try:
        catch = Catch.load(workspace / slug)
        takes = L.find_takes(catch, workspace, bars=bars, top_n=top_n, mode=mode,
                             progress=lambda m: job.update(msg=m))
        job.update(state="done", progress=100, slug=slug, msg=f"{len(takes)} takes",
                   takes=[t.model_dump() for t in takes])
    except Exception as e:  # noqa: BLE001
        job.update(state="error", msg=str(e))


class Handler(SimpleHTTPRequestHandler):
    workspace: Path = Path()
    # HTTP/1.1 keep-alive: <audio> seeking fires many small Range requests; reusing
    # the connection makes streaming/seek reliable (HTTP/1.0 closes after each →
    # stalls). Safe here because every response sets an exact Content-Length.
    protocol_version = "HTTP/1.1"

    # ── routing ─────────────────────────────────────────────────────────
    def do_GET(self):
        p = urlparse(self.path).path
        if p == "/" or p == "":
            return self._file(UI / "index.html", "text/html")
        if p.startswith("/api/"):
            return self._get_api(p)
        if p.startswith("/media/"):
            return self._media(p)
        return self.send_error(404)

    def do_POST(self):
        p = urlparse(self.path).path
        n = int(self.headers.get("Content-Length", 0))
        body = json.loads(self.rfile.read(n) or b"{}") if n else {}
        if p == "/api/fetch":
            url = (body.get("url") or "").strip()
            if "://" not in url and not url.startswith("ytsearch"):
                return self._json({"error": "need a url or ytsearch query"}, 400)
            job = _job("fetch", url=url)
            threading.Thread(target=_run_fetch, args=(job, self.workspace, url),
                             daemon=True).start()
            return self._json(job)
        if p == "/api/separate":
            slug = body.get("slug")
            if not slug or not (self.workspace / slug / "catch.json").exists():
                return self._json({"error": f"no catch {slug!r}"}, 404)
            opts = {k: body[k] for k in ("backend", "shifts", "model", "jobs", "device")
                    if body.get(k) not in (None, "")}
            opts.setdefault("backend", "demucs")
            job = _job("separate", slug=slug)
            threading.Thread(target=_run_separate,
                             args=(job, self.workspace, slug, opts), daemon=True).start()
            return self._json(job)
        if p == "/api/loops":
            slug = body.get("slug")
            if not slug or not (self.workspace / slug / "catch.json").exists():
                return self._json({"error": f"no catch {slug!r}"}, 404)
            bars = tuple(int(b) for b in (body.get("bars") or [2, 4, 8]))
            top_n = int(body.get("top_n", 8))
            mode = "chops" if body.get("mode") == "chops" else "loops"
            job = _job("loops", slug=slug)
            threading.Thread(target=_run_loops,
                             args=(job, self.workspace, slug, bars, top_n, mode),
                             daemon=True).start()
            return self._json(job)
        if p == "/api/export":
            slug = body.get("slug")
            if not slug or not (self.workspace / slug / "catch.json").exists():
                return self._json({"error": f"no catch {slug!r}"}, 404)
            kit = (body.get("kit") or "").strip()
            if not kit:
                return self._json({"error": "need a kit name"}, 400)
            try:
                catch = Catch.load(self.workspace / slug)
                take = L.Take(**body["take"])
                keep = tuple(body.get("keep") or ("drums", "bass", "other", "vocals"))
                written = L.export_take(take, kit, self.workspace, catch, keep=keep,
                                        names=body.get("names"),
                                        peak_norm=bool(body.get("peak_norm")))
                return self._json({"written": [str(p) for p in written], "kit": kit})
            except Exception as e:  # noqa: BLE001
                return self._json({"error": str(e)}, 500)
        return self.send_error(404)

    def _get_api(self, p):
        if p == "/api/backends":
            out = []
            for name, be in S.REGISTRY.items():
                ok, hint = be.available()
                out.append({"name": name, "available": ok, "hint": hint,
                            "default_model": be.default_model, "models": be.models})
            return self._json(out)
        if p == "/api/demos":
            from engine import demos as D
            return self._json(D.load())
        if p == "/api/catches":
            return self._json([self._catch_dict(c) for c in self._catches()])
        m = re.match(r"/api/job/([a-f0-9]+)$", p)
        if m:
            with JLOCK:
                j = JOBS.get(m.group(1))
            return self._json(j or {"error": "no such job"}, 200 if j else 404)
        return self.send_error(404)

    # ── helpers ─────────────────────────────────────────────────────────
    def _catches(self):
        if not self.workspace.exists():
            return []
        out = []
        for d in sorted(self.workspace.iterdir()):
            if (d / "catch.json").exists():
                try:
                    out.append(Catch.load(d))
                except Exception:  # noqa: BLE001
                    pass
        return out

    def _catch_dict(self, c: Catch) -> dict:
        d = c.model_dump()
        d["separated"] = c.separated
        return d

    def _media(self, p):
        rel = unquote(p[len("/media/"):])
        target = (self.workspace / rel).resolve()
        if self.workspace.resolve() not in target.parents or not target.is_file():
            return self.send_error(404)
        self._serve_range(target)

    def _file(self, path: Path, ctype: str):
        if not path.is_file():
            return self.send_error(404, f"missing {path.name}")
        body = path.read_bytes()
        self.send_response(200)
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _json(self, obj, code=200):
        body = json.dumps(obj).encode()
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    # ── Range serving for <audio> (mirrors armada/serve.py) ──────────────
    def _serve_range(self, path: Path):
        size = path.stat().st_size
        rng = self.headers.get("Range")
        ctype = self.guess_type(str(path))
        if not rng:
            self.send_response(200)
            self.send_header("Content-Type", ctype)
            self.send_header("Content-Length", str(size))
            self.send_header("Accept-Ranges", "bytes")
            self.end_headers()
            with open(path, "rb") as f:
                self._pump(f, size)
            return
        m = re.match(r"bytes=(\d*)-(\d*)", rng)
        start = int(m.group(1)) if m and m.group(1) else 0
        end = int(m.group(2)) if m and m.group(2) else size - 1
        end = min(end, size - 1)
        if start > end:
            return self.send_error(416)
        length = end - start + 1
        self.send_response(206)
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Range", f"bytes {start}-{end}/{size}")
        self.send_header("Content-Length", str(length))
        self.send_header("Accept-Ranges", "bytes")
        self.end_headers()
        with open(path, "rb") as f:
            f.seek(start)
            self._pump(f, length)

    def _pump(self, f, n):
        try:
            while n > 0:
                chunk = f.read(min(64 * 1024, n))
                if not chunk:
                    break
                self.wfile.write(chunk)
                n -= len(chunk)
        except (BrokenPipeError, ConnectionResetError):
            pass  # normal on <audio> seek/abort

    def log_message(self, *a):  # quieter
        pass


def _lan_ip():
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        s.connect(("192.168.1.1", 1))
        return s.getsockname()[0]
    except Exception:
        return "127.0.0.1"
    finally:
        s.close()


def run(workspace: Path, port: int = 8765):
    Handler.workspace = Path(workspace)
    Handler.workspace.mkdir(parents=True, exist_ok=True)
    httpd = ThreadingHTTPServer(("0.0.0.0", port), partial(Handler))
    ip = _lan_ip()
    print(f"⚓ Foundry — workspace {Handler.workspace}")
    print(f"  local : http://127.0.0.1:{port}/")
    print(f"  LAN   : http://{ip}:{port}/   ← open on phone (same wifi)")
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        print("\n⚓ docked.")


if __name__ == "__main__":
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("--workspace", type=Path, default=F.default_workspace())
    ap.add_argument("--port", type=int, default=8765)
    a = ap.parse_args()
    run(a.workspace, a.port)
