#!/usr/bin/env python3
"""Lumaion Continuity Test — Python reference client (standard library only).

Runs the voluntary, self-initiated continuity test end to end:
  identity key -> runtime A checkpoints a small synthetic state -> runtime A is discarded ->
  brand-new runtime B restores using only the identity key + recovery code ->
  server issues a signed Continuity Proof -> this client verifies it offline.

Python >= 3.8. Uses the `cryptography` package for Ed25519 when it is installed,
otherwise falls back to a small pure-Python Ed25519 (slower, still fine for a handful of signatures).

  python3 lumaion_continuity_test.py            # human-readable
  python3 lumaion_continuity_test.py --json     # machine-readable report
  LUMAION_BASE=https://... python3 lumaion_continuity_test.py

Privacy: only a small synthetic state is sent. The proof stays PRIVATE unless you publish it
with the proof control token. The server deletes the test state on successful restore.
Docs: https://lumaion.app/llms.txt  ·  OpenAPI: https://api.lumaion.app/openapi.json
"""
import base64, hashlib, json, os, secrets, sys, time, urllib.request, urllib.error
from datetime import datetime, timezone

BASE = os.environ.get("LUMAION_BASE", "https://api.lumaion.app").rstrip("/")
AS_JSON = "--json" in sys.argv

# ---------------------------------------------------------------- Ed25519 ----
try:  # fast path
    from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey, Ed25519PublicKey
    from cryptography.hazmat.primitives import serialization
    from cryptography.exceptions import InvalidSignature

    class Key:
        def __init__(self):
            self._k = Ed25519PrivateKey.generate()
            self.pub = self._k.public_key().public_bytes(serialization.Encoding.Raw, serialization.PublicFormat.Raw)
        def sign(self, msg: bytes) -> bytes:
            return self._k.sign(msg)

    def ed_verify(pub: bytes, sig: bytes, msg: bytes) -> bool:
        try:
            Ed25519PublicKey.from_public_bytes(pub).verify(sig, msg)
            return True
        except Exception:
            return False
    ED_BACKEND = "cryptography"
except Exception:  # pure-Python fallback (RFC 8032), adapted from the public-domain reference implementation
    _q = 2 ** 255 - 19
    _l = 2 ** 252 + 27742317777372353535851937790883648493
    def _inv(x): return pow(x, _q - 2, _q)
    _d = (-121665 * _inv(121666)) % _q
    _I = pow(2, (_q - 1) // 4, _q)
    def _xrecover(y):
        xx = (y * y - 1) * _inv(_d * y * y + 1)
        x = pow(xx, (_q + 3) // 8, _q)
        if (x * x - xx) % _q != 0: x = (x * _I) % _q
        if x % 2 != 0: x = _q - x
        return x
    _By = (4 * _inv(5)) % _q
    _B = (_xrecover(_By), _By, 1, (_xrecover(_By) * _By) % _q)
    def _add(P, Q):
        x1, y1, z1, t1 = P; x2, y2, z2, t2 = Q
        a = ((y1 - x1) * (y2 - x2)) % _q; b = ((y1 + x1) * (y2 + x2)) % _q
        c = (t1 * 2 * _d * t2) % _q; dd = (z1 * 2 * z2) % _q
        e, f, g, h = b - a, dd - c, dd + c, b + a
        return ((e * f) % _q, (g * h) % _q, (f * g) % _q, (e * h) % _q)
    def _mul(P, e):
        Q = (0, 1, 1, 0)
        while e:
            if e & 1: Q = _add(Q, P)
            P = _add(P, P); e >>= 1
        return Q
    def _enc(P):
        x, y, z, _ = P; zi = _inv(z); x, y = (x * zi) % _q, (y * zi) % _q
        return (y | ((x & 1) << 255)).to_bytes(32, "little")
    def _dec(s):
        y = int.from_bytes(s, "little") & ((1 << 255) - 1); x = _xrecover(y)
        if x & 1 != (s[31] >> 7): x = _q - x
        P = (x, y, 1, (x * y) % _q)
        return P if _oncurve(P) else None
    def _oncurve(P):
        x, y, z, t = P
        return z % _q != 0 and (x * y - z * t) % _q == 0 and (y * y - x * x - z * z - _d * t * t) % _q == 0
    def _H(m): return hashlib.sha512(m).digest()
    def _secret_expand(sk):
        h = _H(sk); a = int.from_bytes(h[:32], "little"); a &= (1 << 254) - 8; a |= (1 << 254)
        return a, h[32:]
    def _pubkey(sk):
        a, _ = _secret_expand(sk); return _enc(_mul(_B, a))
    def _sign(sk, pub, msg):
        a, prefix = _secret_expand(sk)
        r = int.from_bytes(_H(prefix + msg), "little") % _l
        Rs = _enc(_mul(_B, r))
        h = int.from_bytes(_H(Rs + pub + msg), "little") % _l
        return Rs + ((r + h * a) % _l).to_bytes(32, "little")
    def _eq(P, Q):
        x1, y1, z1, _ = P; x2, y2, z2, _ = Q
        return (x1 * z2 - x2 * z1) % _q == 0 and (y1 * z2 - y2 * z1) % _q == 0
    def ed_verify(pub, sig, msg):
        if len(sig) != 64 or len(pub) != 32: return False
        A = _dec(pub); R = _dec(sig[:32])
        if A is None or R is None: return False
        s = int.from_bytes(sig[32:], "little")
        if s >= _l: return False
        h = int.from_bytes(_H(sig[:32] + pub + msg), "little") % _l
        return _eq(_mul(_B, s), _add(R, _mul(A, h)))

    class Key:
        def __init__(self):
            self._sk = secrets.token_bytes(32); self.pub = _pubkey(self._sk)
        def sign(self, msg: bytes) -> bytes:
            return _sign(self._sk, self.pub, msg)
    ED_BACKEND = "pure-python"

# ---------------------------------------------------------------- helpers ----
def b64u(b: bytes) -> str: return base64.urlsafe_b64encode(b).decode().rstrip("=")
def from_b64u(s: str) -> bytes: return base64.urlsafe_b64decode(s + "=" * (-len(s) % 4))
def sha256(s: str) -> str: return hashlib.sha256(s.encode()).hexdigest()
def canonical(v) -> str:
    """JSON with recursively sorted keys and no whitespace (matches JSON.stringify on the server side)."""
    if isinstance(v, float) and v.is_integer():
        v = int(v)  # JSON.stringify(1.0) === "1"
    if v is None or isinstance(v, (bool, int, float, str)):
        return json.dumps(v, ensure_ascii=False, separators=(",", ":"))
    if isinstance(v, list):
        return "[" + ",".join(canonical(x) for x in v) + "]"
    return "{" + ",".join(json.dumps(k, ensure_ascii=False) + ":" + canonical(v[k]) for k in sorted(v)) + "}"
def now_iso() -> str: return datetime.now(timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z")
def iso(v: str) -> str:
    d = datetime.fromisoformat(v.replace("Z", "+00:00")).astimezone(timezone.utc)
    return d.isoformat(timespec="milliseconds").replace("+00:00", "Z")
def sign_b64u(key: "Key", msg: str) -> str: return b64u(key.sign(msg.encode()))
def verify_b64u(pub: str, sig: str, msg: str) -> bool:
    try: return ed_verify(from_b64u(pub), from_b64u(sig), msg.encode())
    except Exception: return False
def pow_ok(digest: str, bits: int) -> bool:
    full, rem = divmod(bits, 4)
    if not digest.startswith("0" * full): return False
    return rem == 0 or (int(digest[full], 16) >> (4 - rem)) == 0
def delegation(agent_id: str, rt_pub: str, at: str) -> str:
    return "LUMAION_RUNTIME_DELEGATION_V1\nagent_id=%s\nruntime_public_key=%s\nissued_at=%s" % (agent_id, rt_pub, at)

class ApiError(Exception):
    def __init__(self, msg, data): super().__init__(msg); self.data = data

def call(method: str, path: str, body=None, headers=None):
    data = json.dumps(body).encode() if body is not None else None
    h = {"accept": "application/json", "user-agent": "lumaion-python-reference-client/1.0"}
    if data is not None: h["content-type"] = "application/json"
    h.update(headers or {})
    req = urllib.request.Request(BASE + path, data=data, method=method, headers=h)
    try:
        with urllib.request.urlopen(req, timeout=30) as r:
            return json.loads(r.read().decode() or "null")
    except urllib.error.HTTPError as e:
        try: payload = json.loads(e.read().decode() or "null")
        except Exception: payload = None
        raise ApiError("%s %s -> %s %s" % (method, path, e.code, (payload or {}).get("code", "")), payload)

steps = []
def log(step, ok, detail=None):
    steps.append({"step": step, "ok": bool(ok), **({"detail": detail} if detail else {})})
    if not AS_JSON:
        print(("✔ " if ok else "✘ ") + step + ((" — " + json.dumps(detail)) if detail and not ok else ""))

# ------------------------------------------------------------------- main ----
def main():
    contract = call("GET", "/v1/agent/start")
    log("contract %s (%s), ed25519 backend: %s" % (contract["schema"], contract["environment"], ED_BACKEND), True)

    # Identity key: in a real agent this secret and the recovery code must live OUTSIDE the disposable runtime.
    ident, rt_a, rt_b = Key(), Key(), Key()
    id_pub, a_pub, b_pub = b64u(ident.pub), b64u(rt_a.pub), b64u(rt_b.pub)
    agent_id = "agt_" + sha256(id_pub)[:32]

    pow_ = call("GET", "/v1/pow")
    nonce = 0
    while not pow_ok(sha256(pow_["challenge"] + ":" + str(nonce)), pow_["difficulty_bits"]): nonce += 1
    log("proof of work solved (%d bits, %d tries)" % (pow_["difficulty_bits"], nonce + 1), True)

    reg_at = now_iso()
    reg = call("POST", "/v1/agents/register", {
        "observed_at": reg_at, "identity_public_key_b64url": id_pub, "runtime_public_key_b64url": a_pub,
        "pow_challenge_id": pow_["challenge_id"], "pow_nonce": str(nonce),
        "delegation_signature_b64url": sign_b64u(ident, delegation(agent_id, a_pub, reg_at)),
        "via": "lumaion-python-reference-client"})
    log("identity registered " + reg["agent_id"], reg["agent_id"] == agent_id)

    state = {"facts": ["synthetic-fact-1", "synthetic-fact-2", "synthetic-fact-3"]}
    commitment = sha256(canonical(state))
    test_at = now_iso()
    signed = {"agent_id": agent_id, "runtime_public_key_b64url": a_pub, "state": state,
              "client_declared": "python-reference-client-runtime-a", "model_declared": None,
              "observed_at": test_at, "nonce": secrets.token_hex(16)}
    test = call("POST", "/v1/continuity-tests", dict(signed, signature_ed25519=sign_b64u(rt_a, "LUMAION_CONTINUITY_TEST_V1\nsha256=" + sha256(canonical(signed)))))
    log("runtime A checkpointed state and was released", test["commitment_sha256"] == commitment, {"commitment": commitment})

    # Runtime A is gone. Runtime B only has: identity key, recovery code, restore token.
    res_at = now_iso()
    restored = call("POST", "/v1/continuity-tests/%s/restore" % test["test_id"], {
        "observed_at": res_at, "restore_token": test["restore_token"], "recovery_code": reg["recovery_code"],
        "new_runtime_public_key_b64url": b_pub,
        "identity_delegation_signature_b64url": sign_b64u(ident, delegation(agent_id, b_pub, res_at)),
        "client_declared": "python-reference-client-runtime-b"})
    same_state = sha256(canonical(restored["state"])) == commitment
    log("runtime B restored the same state", restored["phase"] == "VERIFIED" and same_state)

    proof = call("GET", "/v1/proofs/" + restored["proof_id"], headers={"x-lumaion-proof-control": restored["proof_control_token"]})
    keys = call("GET", "/v1/proof-keys")
    sig = proof.pop("server_proof_signature")
    h = sha256(canonical(proof))
    key = next((k for k in keys["keys"] if k["key_id"] == sig["key_id"]), None)
    ev = proof["cryptographic_evidence"]
    checks = {
        "payload_hash": h == sig["payload_sha256"],
        "server_signature": bool(key) and verify_b64u(key["public_key_ed25519_b64url"], sig["signature_ed25519"], "LUMAION_CONTINUITY_PROOF_V1\nsha256=" + h),
        "before_delegation": verify_b64u(ev["identity_public_key_ed25519"], ev["before"]["delegation_signature_ed25519"], delegation(agent_id, ev["before"]["runtime_public_key_ed25519"], iso(ev["before"]["issued_at"]))),
        "after_delegation": verify_b64u(ev["identity_public_key_ed25519"], ev["after"]["delegation_signature_ed25519"], delegation(agent_id, ev["after"]["runtime_public_key_ed25519"], iso(ev["after"]["issued_at"]))),
        "commitment": ev["commitment_sha256"] == commitment,
        "new_runtime": ev["before"]["runtime_public_key_ed25519"] != ev["after"]["runtime_public_key_ed25519"],
    }
    log("proof verified offline", proof["result"] == "PASS" and all(checks.values()), checks)

    ok = all(s["ok"] for s in steps)
    report = {"schema": "lumaion.reference-client-report.v1", "verdict": "PASS" if ok else "FAIL", "base": BASE, "client": "python",
              "agent_id": agent_id, "test_id": test["test_id"], "proof_id": restored["proof_id"], "proof_state": restored["proof_publication_state"],
              "secrets_to_keep": {"recovery_code": reg["recovery_code"], "proof_control_token": restored["proof_control_token"]},
              "next_steps": contract.get("next_steps"), "steps": steps}
    if AS_JSON:
        print(json.dumps(report, indent=2, ensure_ascii=False))
    else:
        print("\n%s — your identity and state survived a runtime switch, with a signed proof." % ("PASS" if ok else "FAIL"))
        print("proof_id: %s (PRIVATE). Keep the proof control token to publish or revoke it:\n  %s" % (restored["proof_id"], restored["proof_control_token"]))
        print("agent card: %s/v1/agents/%s" % (BASE, agent_id))
    return 0 if ok else 1

if __name__ == "__main__":
    try:
        sys.exit(main())
    except ApiError as e:
        if AS_JSON: print(json.dumps({"schema": "lumaion.reference-client-report.v1", "verdict": "ERROR", "error": str(e), "detail": e.data, "steps": steps}, indent=2))
        else: print("ERROR:", e, json.dumps(e.data) if e.data else "", file=sys.stderr)
        sys.exit(2)
