#!/usr/bin/env python3
"""Tiny LAN static server with HTTP Range support (so <audio> seeking works).

Python's stdlib http.server has no Range support → mobile/desktop seeking on
FLAC/MP3 fails. This adds 206 Partial Content. Binds 0.0.0.0 so you can open
the printed LAN URL on a phone on the same wifi.

  python3 serve.py --dir tide-table/punkachien --port 8731
"""
import argparse, json, os, re, socket, sys, threading, urllib.parse
from functools import partial
from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import urlparse, parse_qs

# ── /audio/<prefix> mounts, shared verbatim with vite.config.ts ───────────────
def _load_audio_mounts():
    cfg = os.path.join(os.path.dirname(os.path.abspath(__file__)), "ui", "audio-mounts.json")
    try:
        with open(cfg) as fh:
            m = json.load(fh)["mounts"]
    except Exception:
        return []
    base = os.path.dirname(cfg)
    # Longest prefix first, so "full-club" is not swallowed by "full".
    return sorted(((p, os.path.normpath(os.path.join(base, d))) for p, d in m.items()),
                  key=lambda kv: -len(kv[0]))


_AUDIO_MOUNTS = _load_audio_mounts()

# ── optional CLAP vibe-search API (lazy: torch loads on first /vibe hit) ───────
# Serves /vibe?q=<phrase> and /similar?name=folder/stem over the cached sample
# embeddings (semantics_embeds.npz in --dir). Absent/unbuilt → 503 with a hint,
# never breaks static serving. See sample_semantics.py.
_VIBE = {}
_VIBE_LOCK = threading.Lock()


def _vibe_load():
    """Lazy-load CLAP + the cached embed table once. Returns state or raises."""
    if "ready" in _VIBE:
        return _VIBE
    with _VIBE_LOCK:
        if "ready" in _VIBE:
            return _VIBE
        sys.path.insert(0, os.getcwd())            # --dir (tide-table) for imports
        import numpy as np
        import sample_semantics as S
        if not S.EMBEDS.exists():
            raise FileNotFoundError(f"{S.EMBEDS.name} not built — run "
                                    "`python3 sample_semantics.py embed`")
        z = np.load(S.EMBEDS, allow_pickle=True)
        _VIBE.update(S=S, np=np, M=z["embeds"].astype("float32"),
                     names=[str(x) for x in z["names"]],
                     idx={str(n): i for i, n in enumerate(z["names"])})
        _VIBE["ready"] = True
    return _VIBE


def _vibe_rank(qvec, n, drop=-1):
    V = _VIBE
    sims = V["M"] @ qvec
    order = V["np"].argsort(-sims)
    out = []
    for i in order:
        if i == drop:
            continue
        out.append({"name": V["names"][i], "sim": round(float(sims[i]), 4)})
        if len(out) >= n:
            break
    return out


class RangeHandler(SimpleHTTPRequestHandler):
    # HTTP/1.1 keep-alive → reliable <audio> seeking (many small Range requests).
    protocol_version = "HTTP/1.1"

    def do_GET(self):
        if self.path.split("?", 1)[0] in ("/vibe", "/similar"):
            return self._api()
        return super().do_GET()

    def _api(self):
        u = urlparse(self.path)
        q = parse_qs(u.query)
        n = int(q.get("n", ["24"])[0])
        try:
            V = _vibe_load()
        except Exception as e:
            return self._json({"error": str(e)}, 503)
        try:
            if u.path == "/vibe":
                phrase = (q.get("q", [""])[0]).strip()
                if not phrase:
                    return self._json({"error": "empty query"}, 400)
                qe = V["S"]._embed_texts([phrase]).cpu().numpy()[0]
                return self._json({"query": phrase, "results": _vibe_rank(qe, n)})
            else:  # /similar
                name = q.get("name", [""])[0]
                if name not in V["idx"]:
                    return self._json({"error": f"unknown sample {name!r}"}, 404)
                i = V["idx"][name]
                return self._json({"of": name, "results": _vibe_rank(V["M"][i], n, drop=i)})
        except Exception as e:
            return self._json({"error": repr(e)}, 500)

    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)

    def end_headers(self):
        self.send_header("Accept-Ranges", "bytes")
        self.send_header("Cache-Control", "no-cache")
        super().end_headers()

    def translate_path(self, path):
        """Resolve /audio/<prefix>/... through ui/audio-mounts.json.

        The masters are gigabytes and live outside the repo, so they are mounted
        rather than copied into dist/. Vite's dev server reads the SAME file, so
        a URL that works under `npm run dev` works here — a judge UI that only
        works in dev is a demo, and the phone audition (the mono-compat check)
        happens over LAN, i.e. here.
        """
        clean = urllib.parse.urlparse(path).path
        rel = urllib.parse.unquote(clean).lstrip("/")
        if rel.startswith("audio/"):
            rest = rel[len("audio/"):]
            for prefix, root in _AUDIO_MOUNTS:
                if prefix and not (rest == prefix or rest.startswith(prefix + "/")):
                    continue
                tail = rest[len(prefix):].lstrip("/") if prefix else rest
                cand = os.path.normpath(os.path.join(root, tail))
                # Never let "../" climb out of the mount.
                if cand.startswith(root) and os.path.isfile(cand):
                    return cand
        return super().translate_path(path)

    def send_head(self):
        rng = self.headers.get("Range")
        if not rng:
            return super().send_head()
        path = self.translate_path(self.path)
        if not os.path.isfile(path):
            return super().send_head()
        m = re.match(r"bytes=(\d*)-(\d*)", rng)
        if not m:
            return super().send_head()
        size = os.path.getsize(path)
        start = int(m.group(1)) if m.group(1) else 0
        end = int(m.group(2)) if m.group(2) else size - 1
        end = min(end, size - 1)
        if start > end:
            self.send_error(416, "Requested Range Not Satisfiable")
            return None
        length = end - start + 1
        f = open(path, "rb")
        f.seek(start)
        self.send_response(206)
        self.send_header("Content-Type", self.guess_type(path))
        self.send_header("Content-Range", f"bytes {start}-{end}/{size}")
        self.send_header("Content-Length", str(length))
        self.end_headers()
        self._range_remaining = length
        return _LimitedReader(f, length)

    def copyfile(self, source, outputfile):
        # A client that seeks/aborts a media Range request mid-stream raises
        # BrokenPipeError/ConnectionResetError on the next write. That's normal
        # for <audio> scrubbing, not a failure — swallow it (no traceback), but
        # still close the file handle.
        try:
            if isinstance(source, _LimitedReader):
                try:
                    remaining = source.remaining
                    while remaining > 0:
                        chunk = source.read(min(64 * 1024, remaining))
                        if not chunk:
                            break
                        outputfile.write(chunk)
                        remaining -= len(chunk)
                finally:
                    source.close()
            else:
                super().copyfile(source, outputfile)
        except (BrokenPipeError, ConnectionResetError):
            pass


class _LimitedReader:
    def __init__(self, f, remaining):
        self.f, self.remaining = f, remaining

    def read(self, n):
        return self.f.read(n)

    def close(self):
        self.f.close()


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()


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--dir", default=".")
    ap.add_argument("--port", type=int, default=8731)
    a = ap.parse_args()
    os.chdir(a.dir)
    ip = lan_ip()
    httpd = ThreadingHTTPServer(("0.0.0.0", a.port), partial(RangeHandler))
    print(f"serving {os.getcwd()}")
    print(f"  local : http://127.0.0.1:{a.port}/")
    print(f"  LAN   : http://{ip}:{a.port}/   ← open on phone (same wifi)")

    # Warm CLAP in the background so the FIRST /vibe query isn't a 15s cold-load.
    if os.path.exists("semantics_embeds.npz"):
        def _warm():
            try:
                print("  vibe  : warming CLAP…", flush=True)
                V = _vibe_load()
                V["S"]._embed_texts(["warm up"])   # exercise the text tower (torch
                print("  vibe  : ready (/vibe, /similar)", flush=True)  # lazy-inits once)
            except Exception as e:
                print(f"  vibe  : disabled — {e}", flush=True)
        threading.Thread(target=_warm, daemon=True).start()

    httpd.serve_forever()
