"""Fenced writes: the store rejects a late write from a holder that already lost its cell.
Python 3.8+, standard library only. Set CASCADE_API_KEY (and CASCADE_API for a local server).

Every grant carries `fence`, a number that only grows per cell (renewals keep it).
Send it with each write; the store keeps the highest fence it accepted per cell and
refuses a lower one. Same three store-side checks as fenced-writes.mjs.
"""
import json, os, time, urllib.error, urllib.request


def cascade(path, body):
    api = os.environ.get("CASCADE_API", "https://api.cascadedynamics.ai/v1")
    for attempt in range(6):
        req = urllib.request.Request(api + path, data=json.dumps(body).encode(), method="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=15 + body.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 claim_fenced(cell, agent, ttl_ms=30000, wait_ms=0):
    """{"fence", "receipt"} when granted, {"held": True, "behind"} when not."""
    r = cascade("/claim", {"cell": cell, "agent": agent, "ttl_ms": ttl_ms, **({"wait_ms": wait_ms} if wait_ms else {})})
    return {"fence": r["fence"], "receipt": r["entry_hash"]} if r.get("ok") else {"held": True, "behind": r.get("behind")}


class StaleFence(Exception):
    def __init__(self, cell, fence, highest):
        super().__init__(f"stale fence {fence} for {cell}: the store has already accepted {highest}")
        self.cell, self.fence, self.highest = cell, fence, highest


# ── Postgres (psycopg / psycopg2, %s placeholders) ──────────────────────────
PG_SCHEMA = "CREATE TABLE IF NOT EXISTS cascade_fences (cell text PRIMARY KEY, fence bigint NOT NULL)"
PG_RAISE_FENCE = """
INSERT INTO cascade_fences (cell, fence) VALUES (%s, %s)
ON CONFLICT (cell) DO UPDATE SET fence = EXCLUDED.fence
  WHERE cascade_fences.fence <= EXCLUDED.fence
RETURNING fence"""


def pg_fenced_write(conn, cell, fence, write):
    """Raise the fence and run write(cursor) in ONE transaction. Raises StaleFence."""
    cur = conn.cursor()
    try:
        cur.execute(PG_RAISE_FENCE, (cell, fence))
        if cur.fetchone() is None:
            cur.execute("SELECT fence FROM cascade_fences WHERE cell = %s", (cell,))
            row = cur.fetchone()
            raise StaleFence(cell, fence, row and row[0])
        result = write(cur)
        conn.commit()
        return result
    except BaseException:
        conn.rollback()
        raise


# ── S3 (conditional PUT) ────────────────────────────────────────────────────
def s3_fenced_put(head, put, key, fence, body, tries=5):
    """head(key) -> {"etag", "fence"} | None; put(key, body, fence=, if_match=, if_none_match=)
    -> {"ok": True} | {"precondition": True} (HTTP 412)."""
    for i in range(tries):
        cur = head(key)
        if cur and int(cur["fence"]) > fence:
            raise StaleFence(key, fence, int(cur["fence"]))
        r = put(key, body, fence=fence, **({"if_match": cur["etag"]} if cur else {"if_none_match": "*"}))
        if r.get("ok"):
            return r
        if not r.get("precondition"):
            raise RuntimeError("s3 put failed")
        time.sleep(0.02 * 2 ** i)
    raise RuntimeError(f"s3_fenced_put: still contended after {tries} tries")


# ── Redis ───────────────────────────────────────────────────────────────────
REDIS_FENCED_SET = """
local cur = tonumber(redis.call('HGET', KEYS[1], 'fence') or '0')
local f = tonumber(ARGV[1])
if f < cur then return {0, cur} end
redis.call('HSET', KEYS[1], 'fence', ARGV[1], 'value', ARGV[2])
return {1, f}"""


def redis_fenced_set(r, key, fence, value):
    """r is a redis-py client."""
    ok, cur = r.eval(REDIS_FENCED_SET, 1, key, fence, value)
    if int(ok) != 1:
        raise StaleFence(key, fence, int(cur))
    return {"fence": int(cur)}