import argparse
import json
import subprocess
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path

# DeepSeek-V4-Flash-0731 via the oMLX server (repo: monster/deepseek_v4_flash_0731_mrp_mlx).
# The server is a resident service (launchd party.monster.deepseek-v4) holding ~150GB
# wired when warm; this operator is a thin HTTP client. It runs on the gpu lane so the
# queue serializes it against trellis/flux on this node — that serialization, not this
# script, is what stops the LLM and mesh gen fighting over Metal.
BASE = "http://127.0.0.1:8020"
MODEL = "DeepSeek-V4-Flash-0731-MXFP4-MLX"
LAUNCHD_LABEL = "party.monster.deepseek-v4"

ap = argparse.ArgumentParser()
ap.add_argument("--input", action="append", default=[])
ap.add_argument("--outdir", required=True)
ap.add_argument("--params", default="{}")
a = ap.parse_args()
p = json.loads(a.params)

prompt = (p.get("prompt") or "").strip()
if not prompt:
    print("ERROR: prompt is required")
    sys.exit(1)


def server_up() -> bool:
    try:
        with urllib.request.urlopen(f"{BASE}/health", timeout=5) as r:
            return r.status == 200
    except (urllib.error.URLError, OSError):
        return False


if not server_up():
    # Cold path: ask launchd to (re)start the service, then wait for readiness.
    # First token after a cold start also pays the ~30s weight load.
    print("oMLX server down — kickstarting launchd service", flush=True)
    uid = subprocess.run(["id", "-u"], capture_output=True, text=True).stdout.strip()
    subprocess.run(["launchctl", "kickstart", f"gui/{uid}/{LAUNCHD_LABEL}"],
                   capture_output=True)
    deadline = time.time() + 180
    while time.time() < deadline and not server_up():
        time.sleep(3)
    if not server_up():
        print("ERROR: oMLX server did not come up on :8020. Install the service with "
              "~/Documents/deepseek_v4_flash_0731_mrp_mlx/scripts/install_launchagent.sh")
        sys.exit(1)

messages = []
if (p.get("system") or "").strip():
    messages.append({"role": "system", "content": p["system"].strip()})
messages.append({"role": "user", "content": prompt})

body = json.dumps({
    "model": MODEL,
    "messages": messages,
    "max_tokens": int(p.get("max_tokens", 1024)),
    "temperature": float(p.get("temperature", 0.7)),
}).encode()

t0 = time.time()
req = urllib.request.Request(f"{BASE}/v1/chat/completions", data=body,
                             headers={"Content-Type": "application/json"})
try:
    # generous timeout: cold start pays ~30s model load before first token
    with urllib.request.urlopen(req, timeout=600) as r:
        d = json.loads(r.read())
except (urllib.error.URLError, OSError) as e:
    print(f"ERROR: request failed: {e}")
    sys.exit(1)

if "error" in d:
    print(f"ERROR: {d['error']}")
    sys.exit(1)

msg = d["choices"][0]["message"]
text = (msg.get("content") or "").strip()
usage = d.get("usage", {})
el = time.time() - t0
ct = usage.get("completion_tokens") or 0

out = Path(a.outdir) / "llm_output.txt"
out.write_text(text + "\n")
tps = f", {ct / el:.1f} tok/s" if ct and el > 0 else ""
print(f"done in {el:.1f}s ({ct} tokens{tps}), {len(text)} chars:\n{text[:400]}", flush=True)