"""
Practice: kill the agent, then resume it. Lab 13's crash test (06_crash_resume.py), shrunk to plain
Python: a shop that keeps its own records, an agent that works through fixed steps, a journal that
saves progress after every step, and a kill right after the shop commits a write.

Fill in the TODOs, then run the tests:
    pytest test_durable_practice.py        (locally)
or press "Run tests" on the chapter page.
"""
import copy


class Killed(Exception):
    """Stands in for SIGKILL: the agent stops here, and nothing after this line runs."""


WRITES = {"issue_refund", "cancel_order"}


class Shop:
    """A separate system with its own records: killing the agent doesn't undo what it committed.
    orders: {order_id: {"customer": email, "price": unit price, "qty": units, "status": "delivered" | "cancelled"}}"""

    def __init__(self, orders):
        self.orders = copy.deepcopy(orders)
        self.refunds = []  # (order_id, amount), in the order they were committed
        self.keys = {}     # idempotency key -> the result the first time
        self.calls = 0     # every tool call that reached the shop

    def _run(self, tool, args):
        if tool == "find_customer":
            return {"customer": args["email"]}
        if tool == "list_orders":
            return {"orders": [oid for oid, o in self.orders.items() if o["customer"] == args["customer"]]}
        o = self.orders[args["order_id"]]
        if tool == "get_refunds":
            return {"refunds": [a for oid, a in self.refunds if oid == args["order_id"]]}
        if tool == "get_order":
            return {"order_id": args["order_id"], **o}
        if tool == "issue_refund":
            left = o["price"] * o["qty"] - sum(a for oid, a in self.refunds if oid == args["order_id"])
            if args["amount"] > left + 1e-9:
                return {"error": f"refund over the refundable amount left (${left:.2f})"}
            self.refunds.append((args["order_id"], args["amount"]))
            return {"refunded": args["amount"]}
        if tool == "cancel_order":
            if o["status"] == "cancelled":
                return {"error": "order can't be cancelled: already cancelled"}
            o["status"] = "cancelled"
            return {"cancelled": args["order_id"]}
        raise ValueError(tool)

    def call(self, tool, args, key=None):
        """Run one tool. A write with a key the shop has already seen is not run again: return the
        result it got the first time. A new key: run it, and remember its result under the key."""
        self.calls += 1
        if key is not None and key in self.keys:
            return self.keys[key]
        result = self._run(tool, args)
        if key is not None and tool in WRITES:
            self.keys[key] = result
        return result


def step_key(run_id, step):
    """An idempotency key that stays the same when the step is repeated after a crash:
    the run id and the step number, joined by a colon, e.g. "T07:4"."""
    return f"{run_id}:{step}"


def run(steps, shop, journal=None, run_id="run", keys=False, kill_after_commit=None, key_fn=step_key):
    """Run the agent's steps (a list of (tool, args)) against the shop and return every result.

    journal: a dict saved between runs, or None (no journal: every start begins at step 0).
    With a journal, skip the steps whose results are already in journal["results"], and append each
    new result to it right after the step (that's the save).
    keys: send key_fn(run_id, i) with each write; without keys, send None.
    kill_after_commit: a step number. When that step is a write, raise Killed right after the shop
    returns, before the result is saved: the gap between a write and saving it."""
    if journal is None:
        journal = {}
    results = journal.setdefault("results", [])
    for i, (tool, args) in enumerate(steps):
        if i < len(results):
            continue
        key = key_fn(run_id, i) if keys and tool in WRITES else None
        result = shop.call(tool, args, key)
        if i == kill_after_commit and tool in WRITES:
            raise Killed(i)
        results.append(result)
    return results


def answer(results):
    """What the agent tells the customer: the model's last step, simplified."""
    for r in results:
        if "error" in r:
            return "The system blocked it, so I've passed this to a person."
    done = [r for r in results if "refunded" in r or "cancelled" in r]
    return "Done: " + ", ".join(f"refunded ${r['refunded']:.2f}" if "refunded" in r else "order cancelled" for r in done)


def kill_and_resume(steps, orders, journal_on, keys, kill_at, key_fn=step_key, run_id="T07"):
    """One kill-and-resume test: start the task, kill it right after step kill_at commits, start it
    again, and return (shop, results). Without a journal, the restart begins from scratch."""
    shop = Shop(orders)
    journal = {} if journal_on else None
    try:
        run(steps, shop, journal, run_id, keys, kill_after_commit=kill_at, key_fn=key_fn)
    except Killed:
        pass
    results = run(steps, shop, journal, run_id, keys, key_fn=key_fn)
    return shop, results


def expected_gap_hits(runs, gap_ms, kills_per_run, window=0.9):
    """How many random kills would land in the gap after a write. Each kill lands uniformly over
    `window` of the run, so for one run the chance is writes * gap / (window * run seconds).
    runs: a list of (writes, seconds). Sum kills_per_run * that chance over the runs."""
    return sum(kills_per_run * w * gap_ms / 1000 / (window * s) for w, s in runs)