#!/usr/bin/env python3
"""SPHERE weave lifecycle driver (auto-generated by LoomAI).

Drives the experiment through the local LoomAI backend's SPHERE API so it reuses
the running portal session/auth. Invoked by weave.sh as:

    experiment.py start   <experiment_name>
    experiment.py monitor <experiment_name>
    experiment.py stop    <experiment_name>
"""
import json, os, sys, time, urllib.request, urllib.error

BASE = os.environ.get("LOOMAI_API", "http://127.0.0.1:8000")
HERE = os.path.dirname(os.path.abspath(__file__))
CFG = json.load(open(os.path.join(HERE, "weave.json")))
SPH = CFG.get("sphere", {})
PROJECT = os.environ.get("SPHERE_PROJECT") or SPH.get("project") or None
RID = os.environ.get("REALIZATION_ID") or SPH.get("realization") or "r0"
MODEL_FILE = SPH.get("model", "model.py")
TAG = "weave"


def _api(method, path, body=None):
    url = BASE + path
    data = json.dumps(body).encode() if body is not None else None
    req = urllib.request.Request(url, data=data, method=method,
                                 headers={"Content-Type": "application/json"})
    try:
        with urllib.request.urlopen(req, timeout=120) as r:
            raw = r.read().decode()
            return json.loads(raw) if raw else {}
    except urllib.error.HTTPError as e:
        detail = e.read().decode()[:400]
        raise SystemExit(f"### ERROR: {method} {path} -> {e.code}: {detail}")


def _p(b):  # attach project to a body when configured
    return {**b, "project": PROJECT} if PROJECT else b


def start(exp):
    print(f"### PROGRESS: Recreating SPHERE experiment '{exp}' (realization '{RID}')")
    model = open(os.path.join(HERE, MODEL_FILE)).read()

    # 1. experiment (idempotent — ignore "already exists")
    try:
        _api("POST", "/api/sphere/experiments", _p({"experiment": exp,
             "description": CFG.get("description", "")}))
        print(f"### PROGRESS: created experiment {exp}")
    except SystemExit as e:
        print(f"### PROGRESS: experiment exists or create skipped ({e})")

    # 2. push model
    _api("POST", "/api/sphere/models", _p({"model": model, "experiment": exp, "tag": TAG}))
    print("### PROGRESS: pushed model")

    # 3. realize (poll to completion)
    _api("POST", "/api/sphere/realizations", _p({"experiment": exp, "realization": RID, "tag": TAG}))
    print("### PROGRESS: realize requested; embedding...")
    for _ in range(40):
        time.sleep(3)
        rs = _api("GET", f"/api/sphere/realizations?state=true"
                  + (f"&project={PROJECT}" if PROJECT else ""))
        cur = next((r for r in rs if r.get("eid") == exp and r.get("id") == RID), None)
        st = (cur or {}).get("state")
        if st == "complete":
            print("### PROGRESS: realization complete"); break
        if st == "failed":
            raise SystemExit(f"### ERROR: realization failed: {(cur or {}).get('error', '')}")
    else:
        raise SystemExit("### ERROR: realization did not complete in time")

    # 4. materialize (poll to Success)
    _api("POST", "/api/sphere/materializations", _p({"experiment": exp, "realization": RID}))
    print("### PROGRESS: materialize requested...")
    for _ in range(60):
        time.sleep(4)
        m = _api("GET", f"/api/sphere/materializations/{exp}/{RID}"
                 + (f"?project={PROJECT}" if PROJECT else ""))
        hs = (m.get("taskSummary") or {}).get("HighestStatus")
        if hs == "Success":
            print("### PROGRESS: materialization Success — experiment is up."); return
        if hs in ("Error", "Failure"):
            raise SystemExit(f"### ERROR: materialization {hs}")
    raise SystemExit("### ERROR: materialization did not finish in time")


def monitor(exp):
    rs = _api("GET", "/api/sphere/realizations?state=true"
              + (f"&project={PROJECT}" if PROJECT else ""))
    cur = next((r for r in rs if r.get("eid") == exp and r.get("id") == RID), None)
    if not cur:
        print("### ERROR: realization no longer present"); return 1
    if cur.get("state") == "failed":
        print(f"### ERROR: realization failed: {cur.get('error', '')}"); return 1
    print(f"### PROGRESS: {exp}/{RID} state={cur.get('state')}")
    return 0


def stop(exp):
    print(f"### PROGRESS: Tearing down {exp}/{RID}")
    q = f"?project={PROJECT}" if PROJECT else ""
    try:
        _api("DELETE", f"/api/sphere/materializations/{exp}/{RID}{q}")
        print("### PROGRESS: dematerialized")
    except SystemExit as e:
        print(f"### PROGRESS: dematerialize skipped ({e})")
    try:
        _api("DELETE", f"/api/sphere/realizations/{exp}/{RID}{q}")
        print("### PROGRESS: relinquished")
    except SystemExit as e:
        print(f"### PROGRESS: relinquish skipped ({e})")


if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("usage: experiment.py {start|monitor|stop} <experiment>"); sys.exit(2)
    cmd, exp = sys.argv[1], sys.argv[2]
    fn = {"start": start, "monitor": monitor, "stop": stop}.get(cmd)
    if not fn:
        print(f"unknown command: {cmd}"); sys.exit(2)
    sys.exit(fn(exp) or 0)
