"""Proof of work: release with outcome {status, digest = sha256 of the artifact}, then match
any file to its signed receipt. Python 3.8+, standard library only.
Set CASCADE_API_KEY (and CASCADE_API for a local server).

Check the record itself offline with:
  npx cascade-verify@0.3 bundle.json --fetch-anchors   (Cascade's published keys are built in)
"""
import hashlib, json, os, threading, time, urllib.error, urllib.request

# Cascade's published root key (AWS KMS; certifies the session keys that sign entries).
PUBLISHED_KEY = "MCowBQYDK2VwAyEAulsbuhPp9+zChbyASphQrkm1+TWGrO0Cd8GMQFYWCXw="


def cascade(path, body=None):
    api = os.environ.get("CASCADE_API", "https://api.cascadedynamics.ai/v1")
    for attempt in range(6):
        req = urllib.request.Request(api + path, data=None if body is None else json.dumps(body).encode(),
                                     method="GET" if body is None else "POST", headers={
            "Authorization": "Bearer " + os.environ["CASCADE_API_KEY"],
            "Content-Type": "application/json", "User-Agent": "cascade-recipes/2.0"})
        try:
            with urllib.request.urlopen(req, timeout=30 + (body or {}).get("wait_ms", 0) / 1000) as r:
                return json.loads(r.read())
        except urllib.error.HTTPError as e:
            if e.code in (429, 503) and attempt < 5:
                time.sleep(0.25 * 2 ** attempt)
                continue
            raise


def digest_of(artifact):
    """sha256 hex of bytes, or of a file when given a path (streamed)."""
    if isinstance(artifact, (bytes, bytearray)):
        return hashlib.sha256(artifact).hexdigest()
    h = hashlib.sha256()
    with open(artifact, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def with_proof(cell, agent, work, ttl_ms=60000, wait_ms=0, already_done=None, record=None):
    """work(fence) -> {"artifact": path|bytes, "status"?: str}, renewed every ttl/3 while it runs.
    already_done(cell) is checked after the grant; record(cell, proof) runs BEFORE the release.
    recorded=False means the hold was lost before the release: the outcome is NOT in the record."""
    c = cascade("/claim", {"cell": cell, "agent": agent, "ttl_ms": ttl_ms, **({"wait_ms": wait_ms} if wait_ms else {})})
    if not c.get("ok"):
        return {"did": False, "behind": c.get("behind")}
    grant = {"entry": c["entry"], "fence": c["fence"], "hash": c["entry_hash"]}
    halt = threading.Event()

    def renew():
        while not halt.wait(ttl_ms / 3000):
            try:
                cascade("/claim", {"cell": cell, "agent": agent, "ttl_ms": ttl_ms})
            except Exception:
                pass
    t = threading.Thread(target=renew, daemon=True)
    t.start()

    def stop():
        halt.set()
        t.join()  # a renewal on the wire must not land after the release
    try:
        if already_done and already_done(cell):
            stop()
            cascade("/release", {"cell": cell, "agent": agent})
            return {"did": False, "already": True}
        out = work(c["fence"])
    except BaseException:
        stop()
        try:
            cascade("/release", {"cell": cell, "agent": agent, "outcome": {"status": "failed"}})
        except Exception:
            pass  # the TTL frees it
        raise
    status = out.get("status", "done")
    digest = digest_of(out["artifact"])
    try:
        if record:
            record(cell, {"agent": agent, "status": status, "digest": digest, "grant": grant})
    finally:
        stop()
    r = cascade("/release", {"cell": cell, "agent": agent, "outcome": {"status": status, "digest": digest}})
    return {"did": True, "recorded": r.get("released") is True, "digest": digest, "status": status,
            "grant": grant, "receipt": r.get("entry_hash")}


def fetch_bundle(save_path=None):
    b = cascade("/runs/current/bundle")
    if save_path:
        with open(save_path, "w") as f:
            json.dump(b, f)
    return b


def find_proofs(bundle, digest=None, cell=None, agent=None, status=None):
    opened, out = {}, []
    for e in bundle.get("chain", []):
        c = e.get("cell")
        if not c:
            continue
        if e["verb"] == "ASSIGN":
            cur = opened.get(c)
            if not cur or cur["agent"] != e["agent"] or cur.get("fence") != e.get("fence"):
                opened[c] = e
        elif e["verb"] == "RELEASE":
            g = opened.pop(c, None)
            oc = e.get("outcome")
            if not oc:
                continue
            mine = g is not None and g["agent"] == e["agent"]
            p = {"cell": c, "agent": e["agent"], "status": oc["status"], "digest": oc.get("digest"),
                 "release_hash": e["hash"], "grant_hash": g["hash"] if mine else None, "fence": g.get("fence") if mine else None}
            if ((digest is None or p["digest"] == digest.lower()) and (cell is None or p["cell"] == cell)
                    and (agent is None or p["agent"] == agent) and (status is None or p["status"] == status)):
                out.append(p)
    return out


def match_file(bundle, path):
    """[] means this exact file was never reported."""
    return find_proofs(bundle, digest=digest_of(path))