#!/usr/bin/env python3
"""keyholder — the self-hosted execution half of a compute.wick.pics limit order.

Runs on YOUR machine. Holds YOUR provider API key (env var, never sent anywhere
except the fill call, which the station passes through to the provider without
storing). Long-polls the station for a fill ticket; when one is cut, verifies
its HMAC signature with the order secret and executes the fill.

Environment:
  ORDER_ID        (required) from create order
  ORDER_SECRET    (required) shown once at create
  PROVIDER_KEY    (required) your own Vast.ai or RunPod API key
  COMPUTE_URL     default https://compute.wick.pics
  CONFIRM         "1" places the rental FOR REAL; anything else dry-runs
  ACCOUNT_TOKEN   optional usage-ledger token
  IMAGE, DISK_GB  optional runtime knobs passed to the fill
  BUDGET_USD      optional hard spend cap. At or above it this process destroys
                  the machine with your own key. See "the cap lives here" below.
  STATE_FILE      where the cap's clock is kept (default ./keyholder-state.json)

Standing orders (compute that survives): after a live fill the sidecar keeps
running — every HEARTBEAT_SECONDS (default 45) it checks its own instance via
the station's pass-through status call (your key rides that one call, never
stored) and reports what it sees. Two consecutive dead observations re-arm the
order; the sidecar then long-polls for the refill ticket and fills again. If
this process dies, nothing refills — it is both the witness and the executor —
so keep it under a restart policy (docker --restart=always, or systemd).

The cap lives HERE, not on the station:
  The station also runs a budget guard, but it holds the cap in process memory,
  so a station restart drops the guard while your rental keeps billing — the
  thing a cap exists to prevent is exactly what a restart enabled. An external
  reviewer caught that on 2026-08-25. The durable place for a spend cap is the
  component that holds the key and that YOU control the restart policy of: this
  one. The station guard stays as a second, independent backstop; two guards
  that fail differently are better than one.

  The clock is written to STATE_FILE on every beat, so restarting this process
  does NOT reset accumulated spend. That matters more than it looks: a cap that
  forgets on restart is the same bug one level down.

One command:
  docker build -t keyholder . && docker run -e ORDER_ID=... -e ORDER_SECRET=... \
    -e PROVIDER_KEY=... -e CONFIRM=1 keyholder
Or just: python3 keyholder.py  (stdlib only, no installs)
"""
import hashlib
import hmac
import json
import os
import sys
import time
import urllib.error
import urllib.request

BASE = os.environ.get("COMPUTE_URL", "https://compute.wick.pics").rstrip("/")
ORDER_ID = os.environ.get("ORDER_ID", "")
SECRET = os.environ.get("ORDER_SECRET", "")
KEY = os.environ.get("PROVIDER_KEY", "")
CONFIRM = os.environ.get("CONFIRM", "") == "1"
HB_SECONDS = max(20, int(os.environ.get("HEARTBEAT_SECONDS", "45")))
BUDGET_USD = float(os.environ.get("BUDGET_USD") or 0) or None
STATE_FILE = os.environ.get("STATE_FILE", "keyholder-state.json")

def die(msg):
    print("keyholder:", msg, file=sys.stderr)
    sys.exit(1)

def post(path, body, timeout=70):
    req = urllib.request.Request(
        BASE + path, data=json.dumps(body).encode(),
        headers={"Content-Type": "application/json",
                 "User-Agent": "keyholder-sidecar/1.0"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.loads(r.read())

def _num(v):
    """Fixed 8dp — see the note in ticket_canonical on the station side."""
    return "" if v is None else "%.8f" % float(v)

def canonical(t):
    """Rebuild the exact string the station signed.

    Deliberately NOT re-serialized JSON. Reproducing another runtime's
    separators, key order, unicode escaping and float repr is a fragile thing to
    put on the money path, and it fails closed at the worst moment. This is a
    versioned, newline-separated string you could assemble by hand.

    The `fill` block is not covered on purpose: this program builds that URL
    from its own ORDER_ID and never reads it off the ticket."""
    o = t.get("offer") or {}
    return "\n".join([
        "pangle-ticket/1",
        str(t.get("order_id", "")),
        str(int(t.get("issued_at") or 0)),
        str(int(t.get("expires_at") or 0)),
        str(t.get("nonce", "")),
        str(o.get("id", "")),
        str(o.get("provider", "")),
        str(o.get("gpu_model", "")),
        str(int(o.get("gpu_count") or 0)),
        str(o.get("class", "")),
        _num(o.get("price_final_per_gpu_hr")),
        _num(t.get("line")),
        _num(t.get("auto_destroy_budget_usd")),
    ])

def verify(ticket):
    want = hmac.new(SECRET.encode(), canonical(ticket).encode(),
                    hashlib.sha256).hexdigest()
    return hmac.compare_digest(want, ticket.get("sig", ""))

# ── the spend cap ───────────────────────────────────────────────────────────
# Deliberately a flat file and not a database: this runs on a stranger's laptop
# with no installs, and a cap that needs a service to survive a reboot is not a
# cap you can rely on.

def _load_state():
    try:
        with open(STATE_FILE) as f:
            return json.load(f)
    except Exception:
        return {}

def _save_state(state):
    """Write-then-rename, so a crash mid-write cannot leave a truncated clock
    that reads as zero spend."""
    try:
        tmp = STATE_FILE + ".tmp"
        with open(tmp, "w") as f:
            json.dump(state, f)
        os.replace(tmp, STATE_FILE)
    except Exception as e:
        print(f"keyholder: WARNING could not persist the cap clock ({type(e).__name__});"
              " a restart would lose accumulated spend")

def destroy(provider, instance_id):
    """Destroy with OUR key, through the station's pass-through. Never raises."""
    req = urllib.request.Request(
        f"{BASE}/api/rentals/{provider}/{instance_id}",
        data=json.dumps({"api_key": KEY}).encode(), method="DELETE",
        headers={"Content-Type": "application/json",
                 "User-Agent": "keyholder-sidecar/1.0"})
    try:
        with urllib.request.urlopen(req, timeout=60) as r:
            json.loads(r.read())
        return True
    except urllib.error.HTTPError as e:
        print(f"keyholder: destroy REFUSED: {e.read().decode()[:200]}")
    except Exception as e:
        print(f"keyholder: destroy unreachable ({type(e).__name__})")
    return False

def spend_check(provider, iid, price_hr, reported):
    """Return accumulated spend for this instance, or None if no cap is armed.

    The clock starts the first time we ever see the instance and is kept in
    STATE_FILE, so it survives this process restarting. Spend is the LARGER of
    what the provider reports and wall-clock x price — conservative on purpose:
    ending slightly early is a smaller failure than running past the cap
    because a provider under-reported."""
    if not BUDGET_USD:
        return None
    st = _load_state()
    key = f"{provider}/{iid}"
    rec = st.get(key)
    now = int(time.time())
    if not rec:
        rec = {"started_at": now, "price_hr": float(price_hr or 0)}
        print(f"keyholder: cap ARMED ${BUDGET_USD:.2f} on {key}")
    if price_hr:
        rec["price_hr"] = float(price_hr)
    clock = (now - rec["started_at"]) / 3600 * rec.get("price_hr", 0)
    spend = max(float(reported or 0), clock)
    rec["spend_est"] = round(spend, 4)
    rec["last_seen"] = now
    st[key] = rec
    _save_state(st)
    return spend

def forget(provider, iid):
    st = _load_state()
    if st.pop(f"{provider}/{iid}", None) is not None:
        _save_state(st)

# ── write-ahead journal ─────────────────────────────────────────────────────
# A fill is two facts that must agree: the provider created a machine, and the
# station recorded that it did. Between those two moments this process can die.
# If it comes back and simply asks for another ticket, it can rent a SECOND
# machine that nobody is watching and that bills until someone notices.
#
# So the intent is written down BEFORE the fill is requested. On restart, an
# unresolved intent is treated as "a machine may exist" — never as "nothing
# happened". The station's own broker takes the same posture on an ambiguous
# placement, and for the same reason.

def journal_begin(ticket):
    st = _load_state()
    st["_pending_fill"] = {
        "order": ORDER_ID,
        "offer": str((ticket.get("offer") or {}).get("id", "")),
        "provider": (ticket.get("offer") or {}).get("provider"),
        "at": int(time.time()),
    }
    _save_state(st)

def journal_end():
    st = _load_state()
    if st.pop("_pending_fill", None) is not None:
        _save_state(st)

def pending_fill():
    return _load_state().get("_pending_fill")

def reconcile(order_state, live):
    """Decide what an unresolved fill intent means, now that the station has
    told us the order's state. Returns True if it is safe to carry on."""
    p = pending_fill()
    if not p:
        return True
    age = int(time.time()) - p.get("at", 0)
    if order_state in ("watching", "filled") or (live or {}).get("instance_id"):
        print(f"keyholder: resolved a pending fill from {age}s ago — the station"
              " has it recorded. Clearing the journal.")
        journal_end()
        return True
    # The station does NOT know about a fill we had already asked for. The
    # machine may exist anyway. Renting another one is the one outcome we are
    # not willing to risk on a guess.
    print("keyholder: *** UNRESOLVED FILL ***")
    print(f"keyholder: this process asked for a fill {age}s ago on"
          f" {p.get('provider')} offer {p.get('offer')} and never recorded the answer,")
    print("keyholder: and the station does not show a live machine. A machine MAY")
    print("keyholder: still exist and be billing. Refusing to fill again, because")
    print("keyholder: renting a second machine is worse than renting none.")
    print("keyholder: CHECK YOUR PROVIDER CONSOLE. Then either destroy what you")
    print(f"keyholder: find, or clear the journal: delete \"_pending_fill\" from {STATE_FILE}")
    print("keyholder: (or set RECONCILE=1 to clear it and continue).")
    if os.environ.get("RECONCILE") == "1":
        print("keyholder: RECONCILE=1 set — clearing the journal and continuing.")
        journal_end()
        return True
    return False

def instance_state(provider, instance_id):
    """One pass-through status call with OUR key; never raises."""
    try:
        r = post(f"/api/rentals/{provider}/{instance_id}/status",
                 {"api_key": KEY}, timeout=30)
        return (r.get("state") or "unknown", r.get("price_hr"),
                r.get("spend_est_usd"))
    except urllib.error.HTTPError as e:
        print(f"keyholder: status refused: {e.read().decode()[:120]}")
    except Exception as e:
        print(f"keyholder: status unreachable ({type(e).__name__})")
    return "unknown", None, None

def watch(receipt):
    """Heartbeat loop for a standing order's live machine. Returns the order
    state the station moved to when the watch ended."""
    provider = receipt.get("provider")
    iid = receipt.get("provider_instance_id") or receipt.get("instance_id")
    print(f"keyholder: WATCHING {provider} instance {iid}"
          f" (heartbeat every {HB_SECONDS}s)")
    while True:
        time.sleep(HB_SECONDS)
        state, price, reported = instance_state(provider, iid)

        if state == "gone":
            forget(provider, iid)
        elif state != "unknown":
            spend = spend_check(provider, iid, price, reported)
            if spend is not None and spend >= BUDGET_USD:
                # The cap is the whole reason this branch exists: stop the
                # billing FIRST, tell the station afterwards. If the station is
                # unreachable the machine is still dead, which is the outcome
                # that matters to the person paying for it.
                print(f"keyholder: CAP HIT — ${spend:.2f} >= ${BUDGET_USD:.2f};"
                      f" destroying {provider}/{iid}")
                if destroy(provider, iid):
                    forget(provider, iid)
                    print("keyholder: machine destroyed by the local cap")
                    state = "gone"
                else:
                    # keep the cap armed and retry next beat — a breached cap
                    # with a failed destroy is the one state we must not drop
                    print("keyholder: destroy failed; cap stays armed, retrying")

        try:
            r = post(f"/api/orders/{ORDER_ID}/heartbeat",
                     {"order_secret": SECRET, "state": state,
                      "provider": provider, "instance_id": str(iid),
                      "price_hr": price}, timeout=30)
        except urllib.error.HTTPError as e:
            die(f"station refused the heartbeat: {e.read().decode()[:200]}")
        except Exception as e:
            print(f"keyholder: heartbeat undelivered ({type(e).__name__});"
                  " retrying next beat")
            continue
        ostate = r.get("order_state")
        if ostate == "watching":
            continue
        if ostate == "armed":
            print(f"keyholder: machine is GONE — order re-armed"
                  f" (refill #{r.get('refill_no')}); waiting for the ticket")
        else:
            print(f"keyholder: watch ended — order is {ostate}")
        return ostate

def main():
    # our prints ARE the interface — never let a pipe buffer them
    sys.stdout.reconfigure(line_buffering=True)
    if not (ORDER_ID and SECRET and KEY):
        die("ORDER_ID, ORDER_SECRET and PROVIDER_KEY are required")
    print(f"keyholder: watching order {ORDER_ID} on {BASE}"
          f" ({'LIVE fills' if CONFIRM else 'dry-run only — set CONFIRM=1'})")
    if BUDGET_USD:
        print(f"keyholder: local spend cap ${BUDGET_USD:.2f}"
              f" — clock kept in {STATE_FILE}, survives a restart of this process")
    else:
        print("keyholder: no local spend cap (set BUDGET_USD to arm one)."
              " The station's guard is memory-only and does NOT survive its restart.")
    backoff = 5
    while True:
        try:
            r = post(f"/api/orders/{ORDER_ID}/ticket",
                     {"order_secret": SECRET, "wait": 55})
            backoff = 5
        except urllib.error.HTTPError as e:
            die(f"station refused: {e.read().decode()[:200]}")
        except Exception as e:
            print(f"keyholder: station unreachable ({type(e).__name__}),"
                  f" retrying in {backoff}s")
            time.sleep(backoff)
            backoff = min(backoff * 2, 300)
            continue
        state = r.get("state")
        if not reconcile(state, r.get("live")):
            time.sleep(60)          # a human has to look; do not spin on it
            continue
        if state in ("filled", "cancelled", "expired", "retired", "gone"):
            print(f"keyholder: order is {state}; exiting")
            return
        if state == "watching" and CONFIRM and r.get("live", {}).get("instance_id"):
            # restarted mid-watch: resume supervising the live machine
            ended = watch(r["live"])
            if ended == "armed":
                continue
            return
        if state == "watching":
            # dry-run sidecar on a live watch: nothing to execute, don't spin
            time.sleep(HB_SECONDS)
            continue
        t = r.get("ticket")
        if not t:
            continue
        if not verify(t):
            die("ticket signature FAILED verification — refusing to act")
        offer = t["offer"]
        print(f"keyholder: ticket — {offer['gpu_model']} on {offer['provider']}"
              f" at ${offer['price_final_per_gpu_hr']:.4f}/GPU/hr"
              f" (line ${t['line']:.4f}, expires {t['expires_at']})")
        fill = {"order_secret": SECRET, "api_key": KEY,
                "confirm": CONFIRM, "dry_run": not CONFIRM}
        if os.environ.get("ACCOUNT_TOKEN"):
            fill["account_token"] = os.environ["ACCOUNT_TOKEN"]
        if os.environ.get("IMAGE"):
            fill["image"] = os.environ["IMAGE"]
        if os.environ.get("DISK_GB"):
            fill["disk_gb"] = float(os.environ["DISK_GB"])
        if CONFIRM:
            journal_begin(t)        # write-ahead: before the machine can exist
        try:
            out = post(f"/api/orders/{ORDER_ID}/fill", fill, timeout=120)
            journal_end()           # answered, whatever the answer was
        except urllib.error.HTTPError as e:
            # A refusal is a DEFINITE no from the station — nothing was placed,
            # so the intent is resolved and must not linger as a false alarm.
            journal_end()
            print(f"keyholder: fill refused: {e.read().decode()[:300]}")
            time.sleep(10)
            continue
        except Exception as e:
            # No answer came back. The machine MAY exist. Leave the journal
            # standing so a restart reconciles instead of filling again.
            print(f"keyholder: fill outcome UNKNOWN ({type(e).__name__})."
                  " Leaving the write-ahead journal in place; check your"
                  " provider console before restarting.")
            time.sleep(15)
            continue
        print("keyholder: receipt:", json.dumps(out.get("receipt", out), indent=2))
        if CONFIRM and out.get("order_state") == "watching":
            ended = watch(out.get("receipt", {}))
            if ended == "armed":
                continue           # refill: back to the ticket long-poll
            return
        if CONFIRM and out.get("order_state") == "filled":
            print("keyholder: order FILLED; exiting")
            return
        if not CONFIRM:
            print("keyholder: dry-run complete (order stays armed);"
                  " set CONFIRM=1 for real fills. Watching on.")
            time.sleep(30)

if __name__ == "__main__":
    main()