# session_manager.py — the GPU mutex. one warm model at a time on workpc.
# owns: model registry, llama-server subprocess lifecycle, session leases + reaper.
# workpc holds no durable state; this only guards the single local GPU.
#
# subprocess models (gemma4): this process starts/stops llama-server and health-waits.
# in-process models (siglip2, dots): no server — the lease just reserves the GPU; the
# worker loads the transformers model itself after /session/open returns (port=None).
# the guarantee that matters is the mutex: a second open() gets 409 until the first closes.
import time, uuid, threading, subprocess
import requests
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

app = FastAPI()

MODELS = {
    "gemma4": {
        "binary": "/mnt/D/AI/llama-cpp/llama.cpp/build/bin/llama-server",
        "model":  "/mnt/D/AI/gemma4/gemma-4-12B-it-qat-UD-Q4_K_XL.gguf",
        "mmproj": "/mnt/D/AI/gemma4/mmproj-F16.gguf",
        "port": 8090, "ctx": 32768,   # 12B fits 32k KV unquantized on 16GB ROCm; do NOT -ctk/-ctv quant (leaks on ROCm)
        # MTP speculative decoding: the ~430M gemma-4 draft head drafts up to 4 tokens the 12B verifies.
        # needs `--spec-type draft-mtp` to actually engage (without it llama-server loads the draft but
        # silently runs no speculation). ~42% accept on our JSON stages -> faster decode; ~0.9GB VRAM.
        "args_extra": ["--jinja", "--reasoning-budget", "0",
                       "-md", "/mnt/D/AI/gemma4/mtp-gemma-4-12b-it-BF16.gguf", "-ngld", "99",
                       "--spec-type", "draft-mtp", "--spec-draft-n-max", "4"],
    },
    # in-process models: the worker loads them and holds ~5GB resident. `worker` is the port whose
    # /unload frees that GPU memory on close -- the mutex alone can't reclaim it (leaked across runs,
    # starving the next model on open). subprocess models (gemma4) free via terminate instead.
    "siglip2": {"binary": "transformers", "model": "google/siglip2-so400m-patch16-384", "port": None, "worker": 8003},
    "dots":    {"binary": "transformers", "model": "rednote-hilab/dots.tts-base", "port": None, "worker": 8006},
}

SKIP_HEALTH = False  # self-check flips this; real runs health-wait the llama-server

_lock = threading.Lock()
# the single resident session, or None. {model, session_id, opened_at, port, ttl, last_beat, proc}
_active = None


class OpenReq(BaseModel):
    model: str
    ttl: int = 3600

class SessionReq(BaseModel):
    session_id: str


def _health_wait(port: int, timeout: int = 300):
    deadline = time.time() + timeout
    while time.time() < deadline:
        try:
            if requests.get(f"http://127.0.0.1:{port}/health", timeout=2).status_code == 200:
                return
        except requests.RequestException:
            pass
        time.sleep(1)
    raise TimeoutError(f"llama-server on :{port} never became healthy")


def _start_subprocess(cfg):
    """spawn llama-server for a subprocess-type model, return the Popen once healthy."""
    cmd = [cfg["binary"], "-m", cfg["model"], "--port", str(cfg["port"]),
           "-c", str(cfg.get("ctx", 4096)), "-ngl", "99"]
    if cfg.get("mmproj"):
        cmd += ["--mmproj", cfg["mmproj"]]
    cmd += cfg.get("args_extra", [])
    # keep llama-server's stderr — it prints the assert / CUDA error / context-overflow that kills it.
    # DEVNULL here is why crashes showed up as "no error in the logs". append so a respawn keeps history.
    log = open(f"/tmp/llama-server-{cfg['port']}.log", "a")
    proc = subprocess.Popen(cmd, stdout=log, stderr=subprocess.STDOUT)
    if not SKIP_HEALTH and cfg["port"]:
        _health_wait(cfg["port"])
    return proc


def _teardown(sess):
    proc = sess.get("proc")
    if proc:
        if proc.poll() is None:
            proc.terminate()
        try:
            proc.wait(timeout=15)          # wait() even on an already-dead proc, else it stays a zombie
        except subprocess.TimeoutExpired:
            proc.kill()
    # in-process models live in the worker, not here -- ask it to free the GPU (best-effort;
    # if the worker is down the VRAM is already gone with it).
    worker = sess.get("worker")
    if worker:
        try:
            requests.post(f"http://127.0.0.1:{worker}/unload", timeout=30)
        except requests.RequestException:
            pass


@app.post("/session/open")
def open_session(req: OpenReq):
    if req.model not in MODELS:
        raise HTTPException(400, f"unknown model {req.model}")
    global _active
    cfg = MODELS[req.model]
    session_id = uuid.uuid4().hex
    with _lock:
        if _active is not None:
            raise HTTPException(409, f"{_active['model']} resident (session {_active['session_id']})")
        # ponytail: claim the slot before the slow spawn/health-wait so /active and /close
        # for other sessions aren't blocked behind it; only the actual subprocess start runs unlocked.
        _active = {
            "model": req.model, "session_id": session_id, "opened_at": time.time(),
            "port": cfg.get("port"), "ttl": req.ttl, "last_beat": time.time(), "proc": None,
            "worker": cfg.get("worker"),
        }
    proc = _start_subprocess(cfg) if cfg["binary"] != "transformers" else None
    with _lock:
        if _active is not None and _active["session_id"] == session_id:
            _active["proc"] = proc
            return {"session_id": session_id, "port": cfg.get("port")}
    # The lease was closed or reaped while we were spawning (close saw proc=None and tore down
    # nothing). The server we just started is unreferenced and would hold its VRAM until someone
    # killed it by hand, and the next open would spawn a SECOND one on the same port.
    _teardown({"proc": proc})
    raise HTTPException(409, "session closed while the model was loading")


@app.post("/session/close")
def close_session(req: SessionReq):
    with _lock:
        global _active
        if _active is None or _active["session_id"] != req.session_id:
            raise HTTPException(404, "no such active session")
        _teardown(_active)
        _active = None
        return {"ok": True}


@app.get("/session/active")
def active():
    with _lock:
        if _active is None:
            return None
        return {k: _active[k] for k in ("model", "session_id", "opened_at", "port")}


@app.post("/session/heartbeat")
def heartbeat(req: SessionReq):
    with _lock:
        if _active is None or _active["session_id"] != req.session_id:
            raise HTTPException(404, "no such active session")
        _active["last_beat"] = time.time()
        return {"ok": True}


def _reap_once():
    with _lock:
        global _active
        if _active and time.time() - _active["last_beat"] > _active["ttl"]:
            _teardown(_active)
            _active = None
            return True
    return False


def _supervise_once():
    """Respawn a crashed subprocess model in place. llama-server can OOM/die mid-session; the worker
    talks to the port directly, so if we don't bring it back every remaining panel of the stage fails
    with 'connection refused'. Restarting under the SAME session keeps the lease + port valid so the
    orchestrator never has to re-open. Returns True if a respawn happened."""
    global _active
    with _lock:
        sess = _active
        if not sess:
            return False
        proc = sess.get("proc")
        if proc is None or proc.poll() is None:
            return False                   # in-process model, or subprocess still alive
        try:
            proc.wait(timeout=1)           # reap the zombie
        except subprocess.TimeoutExpired:
            pass
        # claim the respawn before releasing the lock: proc=None makes the next _supervise_once pass
        # return early, so the health wait can't be entered twice for one death.
        sess["proc"] = None
        session_id, model = sess["session_id"], sess["model"]
    # ponytail: spawn + health-wait UNLOCKED, like open_session. Holding _lock here blocked
    # /session/active, /close and /open for the full health wait (up to 300s).
    print(f"[supervise] gemma subprocess died (rc={proc.returncode}) session {session_id}; "
          f"respawning", flush=True)
    try:
        new_proc = _start_subprocess(MODELS[model])
    except Exception as e:
        print(f"[supervise] respawn failed: {e}; clearing session so the mutex frees", flush=True)
        with _lock:
            if _active is not None and _active["session_id"] == session_id:
                _active = None
        return False
    with _lock:
        if _active is not None and _active["session_id"] == session_id:
            _active["proc"] = new_proc
            _active["last_beat"] = time.time()  # don't count the downtime against the TTL reaper
            return True
    _teardown({"proc": new_proc})           # lease went away mid-respawn; don't orphan the server
    return False


def _reaper():
    while True:
        time.sleep(1)
        if _reap_once():
            continue
        _supervise_once()


@app.on_event("startup")
def _startup():
    threading.Thread(target=_reaper, daemon=True).start()


if __name__ == "__main__":
    # self-check: mutex (409 on second open), close releases, reaper force-closes a stale lease.
    SKIP_HEALTH = True
    MODELS["gemma4"] = {"binary": "/bin/sleep", "model": "5", "port": None, "args_extra": []}

    s1 = open_session(OpenReq(model="gemma4", ttl=3600))
    assert active()["model"] == "gemma4"
    try:
        open_session(OpenReq(model="siglip2"))
        assert False, "second open must 409"
    except HTTPException as e:
        assert e.status_code == 409
    close_session(SessionReq(session_id=s1["session_id"]))
    assert active() is None

    s2 = open_session(OpenReq(model="gemma4", ttl=2))
    _active["last_beat"] -= 5  # simulate a lease that missed its heartbeat window
    assert _reap_once() is True
    assert active() is None

    # supervisor: a crashed subprocess is respawned in place, keeping the same session_id + port.
    s4 = open_session(OpenReq(model="gemma4", ttl=3600))
    old_proc = _active["proc"]
    old_proc.terminate(); old_proc.wait()          # simulate the llama-server crashing mid-session
    assert _supervise_once() is True
    assert active()["session_id"] == s4["session_id"]      # lease survived
    assert _active["proc"] is not old_proc and _active["proc"].poll() is None  # fresh live proc
    close_session(SessionReq(session_id=s4["session_id"]))
    assert active() is None

    # a close landing DURING the spawn must not orphan the server: open reports 409 and tears it down.
    MODELS["gemma4"]["binary"] = "/bin/sleep"
    MODELS["gemma4"]["port"] = None
    real_start = _start_subprocess

    def racing_start(cfg):
        p = real_start(cfg)
        close_session(SessionReq(session_id=_active["session_id"]))   # close mid-load
        return p
    _start_subprocess = racing_start
    try:
        open_session(OpenReq(model="gemma4", ttl=3600))
        assert False, "open must 409 when its lease vanished mid-load"
    except HTTPException as e:
        assert e.status_code == 409
    _start_subprocess = real_start
    assert active() is None

    # in-process model: close must POST /unload to the worker so its resident VRAM is freed.
    calls = []
    requests.post = lambda url, **kw: calls.append(url) or type("R", (), {"status_code": 200})()
    s3 = open_session(OpenReq(model="siglip2", ttl=3600))
    close_session(SessionReq(session_id=s3["session_id"]))
    assert any("/unload" in u and ":8003" in u for u in calls), calls
    print("session_manager self-check ok")