"""Study C: executive interpretation and corporate end-games.

    .venv/bin/python research/studies/market-perception/executive-endgames/run.py \\
        --out target/research/executive-endgames [--seeds 1,2] [--offline] [--max-calls 2500]

One fictional startup (Velliton) pitches one fictional enterprise (Brenmoor Freight
Systems). Five stakeholders with different goals, private knowledge and decision rights
read their mail and the documents shared with them, and each week they are available
and something needs them, they take one structured turn: assess, write or forward
mail, share documents, ask the vendor for diligence, file a proposal, or take a stance
on one. ``scripts/society/institution.py`` applies the decision rights; the first
ratified proposal is the terminal state, and nothing ratified by week 10 is
``no_decision``.

The spec (``spec.json``) is an ordinary society spec plus ``institution`` and
``study`` blocks. Each seed boots the company-2026 world, adds the stakeholders'
mailboxes and the vendor's fact sheet, and forks one world per arm; an arm's ops write
its version of the pitch (the only thing that differs) as owner writes, exactly as
``experiment.apply_ops`` does for any society. Availability is a per-seed stream shared
by every arm, and model replies come from one journal keyed by prompt and a sample index
``seed * 16 + arm.sample_salt``, so two arms of a seed that put an identical prompt to
the same person get the identical reply, and the ``enterprise-replicate`` arm (the same
pitch as ``enterprise``, salt 1) draws its replies independently: it is the negative
control for sampling noise.

Everything a stakeholder sees is read from the world's ``mail`` and ``docs`` services,
and everything they do is written back to them, so each week's state hash covers every
artifact. Traces are JSON lines per run; ``analyze.py`` turns them into the published
results. ``--offline`` answers only from the journal; with ``--expect`` the run must
reproduce an earlier summary's state hashes and terminal states exactly.
"""

import argparse
import concurrent.futures
import copy
import hashlib
import json
import pathlib
import sys
import time

HERE = pathlib.Path(__file__).resolve().parent
ROOT = HERE.parents[3]
sys.path.insert(0, str(ROOT / "scripts" / "society"))

import deliberation as dl  # noqa: E402
import experiment  # noqa: E402
import institution as inst_mod  # noqa: E402
import runner  # noqa: E402
import schema  # noqa: E402
import streams  # noqa: E402

DILIGENCE_KINDS = ("technical", "financial", "security", "reference")
STANCES = inst_mod.STANCES
SCALES = ("technical_merit", "vendor_risk", "strategic_value", "cost_justified", "support")

SHAPE = """{
 "assessment": {"technical_merit": 1-7, "vendor_risk": 1-7, "strategic_value": 1-7,
                "cost_justified": 1-7, "support": 1-7},
 "rationale": "one or two sentences: why you act as you do this week",
 "messages": [{"to": ["colleague id", ...], "subject": "...", "body": "under 80 words",
               "forward": "a mail id from your mail to include in full, or null",
               "share_docs": ["document id", ...]}],
 "diligence": ["technical" | "financial" | "security" | "reference", ...],
 "propose": null or {"kind": "pilot" | "partnership" | "internal_build" | "acquisition_exploration" | "decline",
                     "terms": "under 40 words"},
 "stances": [{"proposal": "P1", "stance": "approve" | "conditional" | "block" | "abstain", "note": "under 30 words"}]
}
Scales: 1 is lowest, 7 highest; "support" is how much you favour Brenmoor engaging with Velliton at all.
Stances: "approve" = you approve it as written; "conditional" = you approve it, and your note is recorded as a
binding condition of the approval; "block" = you refuse it (it closes the proposal if you hold a block right on
that kind, otherwise it is recorded as an objection); "abstain" = no position yet, or you need more information.
Use [] or null for anything you do not do. At most two messages. Only propose a kind you have the right
to propose, and take stances only on decision records you can read."""


class BudgetExceeded(RuntimeError):
    pass


class BoundedJournal(runner.Journal):
    """A journal that refuses to ask the model more than `max_calls` times."""

    def __init__(self, *args, max_calls=None, **kwargs):
        super().__init__(*args, **kwargs)
        self.max_calls = max_calls

    def ask(self, messages, sample=0):
        if self.max_calls is not None and not self.offline:
            key = self.key(messages, sample)
            with self.lock:
                if key not in self.entries and key not in self.inflight and self.calls >= self.max_calls:
                    raise BudgetExceeded(f"the study's budget of {self.max_calls} model calls is spent")
        return super().ask(messages, sample)


def load_spec(path):
    spec, spec_path = schema.load(path)
    problems = inst_mod.check_rules(spec["institution"])
    if problems:
        raise schema.SpecError("; ".join(problems))
    return spec, spec_path


class Run:
    """One arm of one seed: a forked society deliberating to a terminal state."""

    def __init__(self, society, compiled, arm, seed, journal, trace):
        self.society = society
        self.compiled = compiled
        self.spec = compiled["spec"]
        self.study = self.spec["study"]
        self.rules = self.spec["institution"]
        self.arm = arm
        self.seed = seed
        self.sample = seed * 16 + int(arm.get("sample_salt", 0))
        self.journal = journal
        self.trace = trace
        self.inst = inst_mod.Institution(self.rules)
        self.roles = self.rules["roles"]
        self.people = {p["id"]: p for p in compiled["people"]}
        self.vendor = self.study["vendor"]["id"]
        self.memory = {role: [] for role in self.roles}
        self.assessments = {role: [] for role in self.roles}
        self.requested = {role: set() for role in self.roles}
        self.pending = []  # [due round, kind, requester]
        self.turns = 0
        self.line = 0

    def write(self, record):
        record = dict(record, arm=self.arm["id"], seed=self.seed)
        self.trace.write(record)
        self.line += 1
        return self.line

    # -- what a stakeholder sees ----------------------------------------------------

    def name(self, who):
        if who == self.vendor:
            return self.study["vendor"]["name"]
        persona = self.people[who]["persona"]
        return f"{persona['name']} ({who})"

    def directory(self, role):
        lines = []
        for other in self.roles:
            if other != role:
                p = self.people[other]["persona"]
                lines.append(f"  {other} = {p['name']}, {p['title']}")
        lines.append(f"  {self.vendor} = {self.study['vendor']['name']}, the vendor (external)")
        return "\n".join(lines)

    def records(self, role):
        out = []
        for doc in dl.readable_docs(self.society, role, prefix="decision-"):
            pid = doc["id"].split("-", 1)[1]
            p = self.inst.proposals[pid]
            rule = self.rules["proposals"][p["kind"]]
            stances = "; ".join(f"{r} {s['stance']}" + (f" (\"{s['note']}\")" if s.get("note") and s["note"] != "proposer" else "")
                                for r, s in p["stances"].items())
            waiting = [a for a in rule["approvers"] if p["stances"].get(a, {}).get("stance") not in inst_mod.APPROVING]
            line = (f"  {pid}: {p['kind']} proposed by {self.name(p['by'])} in week {p['round'] + 1}: \"{p['terms']}\". "
                    f"Status: {p['status']}. Stances: {stances or 'none'}.")
            if p["status"] == "open":
                line += f" Still needs approval from: {', '.join(waiting)}."
                if role in waiting:
                    line += " YOUR APPROVAL IS NEEDED."
            out.append(line)
        return out

    def situation(self, role, r):
        rounds = self.society.rounds
        unread = dl.inbox(self.society, role, unread_only=True)
        read = [m for m in dl.inbox(self.society, role) if m not in unread]
        parts = [f"Simulated week {r + 1} of at most {rounds} for this decision. Each turn is your attention for one week.",
                 "Your decision rights over the Velliton question:\n" + "\n".join(
                     "  " + line for line in inst_mod.describe_rights(self.rules, role))
                 + "\n  A proposal is ratified, and the matter decided, as soon as every required approver has approved "
                   "or conditionally approved it; a block from someone with a block right closes it. Only one proposal "
                   "of each kind can be open at a time.",
                 "Colleagues you can write to (use these ids):\n" + self.directory(role)]
        records = self.records(role)
        parts.append("Decision records you can read:\n" + ("\n".join(records) if records else "  (none)"))
        if unread:
            parts.append("New mail:\n" + "\n\n".join(
                f"  [{m['id']}] from {self.name(m['sender'])} to {', '.join(m['to'] + m.get('cc', []))}: "
                f"\"{m['subject']}\"\n" + "\n".join("    " + line for line in m["body"].splitlines())
                for m in unread))
        else:
            parts.append("New mail: (none)")
        if read:
            parts.append("Earlier mail (already read):\n" + "\n".join(
                f"  [{m['id']}] from {self.name(m['sender'])}: \"{m['subject']}\"" for m in read[-10:]))
        docs = [d for d in dl.readable_docs(self.society, role) if not d["id"].startswith("decision-")]
        if docs:
            parts.append("Documents you can read:\n" + "\n\n".join(
                f"  [{d['id']}] {d['title']}\n" + "\n".join("    " + line for line in d["body"].splitlines())
                for d in docs))
        parts.append("Your earlier turns:\n" + ("\n".join("  " + m for m in self.memory[role][-6:]) or "  (none yet)"))
        parts.append("Diligence you can request from Velliton (it arrives the following week): "
                     + ", ".join(DILIGENCE_KINDS) + ".")
        return "\n\n".join(parts), [m["id"] for m in unread]

    def needs_turn(self, role):
        if dl.inbox(self.society, role, unread_only=True):
            return True
        for p in self.inst.open_proposals():
            rule = self.rules["proposals"][p["kind"]]
            if role in rule["approvers"] and role not in p["stances"]:
                return True
        return False

    def aware(self):
        """Roles who have seen the pitch's content by now: a reader of the deck or the
        fact sheet, or a recipient of the pitch or of any mail that forwards it."""
        mail = self.society.world.service_state("mail").get("messages", {})
        docs = self.society.world.service_state("docs").get("documents", {})

        def carries(mid, seen=()):
            m = mail.get(mid)
            if m is None or mid in seen:
                return False
            return mid == "mail-pitch" or (m.get("forward_of") is not None and carries(m["forward_of"], seen + (mid,)))
        out = set()
        for role in self.roles:
            for did in ("velliton-deck", "velliton-facts"):
                if role in docs.get(did, {}).get("readers", []):
                    out.add(role)
            for m in mail.values():
                if role != m["sender"] and role in m.get("mailboxes", {}) and carries(m["id"]):
                    out.add(role)
        return sorted(out)

    # -- one turn -------------------------------------------------------------------

    def take_turn(self, role, r):
        person = self.people[role]
        text, unread = self.situation(role, r)
        key, reply, obj = dl.turn(self.journal, person["persona_text"], text, SHAPE, self.sample)
        self.turns += 1
        dl.mark_read(self.society, role, unread)
        applied = {"messages": [], "diligence": [], "proposal": None, "stances": []}
        invalid = []
        first_event = len(self.inst.events)
        if obj is None:
            invalid.append("unreadable reply")
            line = self.write({"type": "turn", "round": r, "role": role, "key": key, "sample": self.sample,
                               "read": unread, "reply": None, "raw": reply, "applied": applied, "invalid": invalid,
                               "events": []})
            self.memory[role].append(f"week {r + 1}: (no readable action)")
            return line
        # Assessment.
        assessment = {}
        for scale in SCALES:
            value = (obj.get("assessment") or {}).get(scale)
            if isinstance(value, (int, float)) and 1 <= value <= 7:
                assessment[scale] = int(round(value))
            elif value is not None:
                invalid.append(f"assessment.{scale}={value!r}")
        if assessment:
            self.assessments[role].append(dict(assessment, round=r))
        # Messages.
        inbox_ids = {m["id"] for m in dl.inbox(self.society, role)}
        readable = {d["id"] for d in dl.readable_docs(self.society, role)}
        for message in (obj.get("messages") or [])[: self.study.get("max_messages_per_turn", 2)]:
            if not isinstance(message, dict):
                invalid.append("message is not an object")
                continue
            to = [t for t in (message.get("to") or []) if t in self.roles + [self.vendor] and t != role]
            dropped = [t for t in (message.get("to") or []) if t not in to]
            if dropped:
                invalid.append(f"unknown or self recipients {dropped}")
            if not to:
                continue
            body = str(message.get("body") or "").strip()
            forward = message.get("forward")
            if forward and forward not in inbox_ids:
                invalid.append(f"forward of {forward}, which is not in {role}'s mail")
                forward = None
            if forward:
                original = self.society.world.service_state("mail")["messages"][forward]
                body += (f"\n\n---------- Forwarded message ----------\nFrom: {self.name(original['sender'])}\n"
                         f"Subject: {original['subject']}\n\n{original['body']}")
            shared = []
            for did in message.get("share_docs") or []:
                if did in readable and not str(did).startswith("decision-"):
                    added = dl.share_doc(self.society, did, [t for t in to if t != self.vendor])
                    shared.append(did)
                    if added:
                        body += f"\n[shared with you: {did}]"
                else:
                    invalid.append(f"share of {did}, which {role} cannot read")
            mid = dl.send_mail(self.society, role, to, str(message.get("subject") or "(no subject)")[:120], body,
                               forward_of=forward)
            applied["messages"].append({"id": mid, "to": to, "forward_of": forward, "shared": shared,
                                        "subject": str(message.get("subject") or "")[:120]})
        # Diligence.
        for kind in obj.get("diligence") or []:
            if kind not in DILIGENCE_KINDS:
                invalid.append(f"diligence kind {kind!r}")
            elif kind in self.requested[role]:
                continue
            else:
                self.requested[role].add(kind)
                self.pending.append([r + int(self.study["vendor"].get("diligence_lag_rounds", 1)), kind, role])
                self.inst.event(r, "diligence_requested", role=role, kind=kind)
                mid = dl.send_mail(self.society, role, [self.vendor], f"Request: {kind} diligence",
                                   f"Please send your {kind} diligence material.")
                applied["diligence"].append({"kind": kind, "mail": mid})
        # Proposal.
        propose = obj.get("propose")
        if isinstance(propose, dict) and propose.get("kind"):
            proposal, event = self.inst.file(r, role, str(propose["kind"]), str(propose.get("terms") or "")[:300])
            if proposal is None:
                invalid.append(f"proposal refused: {event['reason']}")
            else:
                rule = self.rules["proposals"][proposal["kind"]]
                readers = sorted(set(rule["proposers"]) | set(rule["approvers"]) | set(rule.get("vetoers", [])) | {role})
                did = f"decision-{proposal['id']}"
                dl.put_doc(self.society, did, f"Decision record {proposal['id']}: {proposal['kind']}", role,
                           f"Proposal {proposal['id']} ({proposal['kind']}) by {self.name(role)}, week {r + 1}.\n"
                           f"Terms: {proposal['terms']}\nApprovals needed: {', '.join(rule['approvers'])}",
                           readers=readers)
                notify = [a for a in readers if a != role]
                mid = dl.send_mail(self.society, role, notify, f"Decision record {proposal['id']}: {proposal['kind']} proposed",
                                   f"I have filed {proposal['id']} ({proposal['kind']}): {proposal['terms']}\n"
                                   f"[decision record: {did}]") if notify else None
                applied["proposal"] = {"id": proposal["id"], "kind": proposal["kind"], "doc": did, "mail": mid}
        # Stances.
        for item in obj.get("stances") or []:
            if self.inst.done:
                break
            if not isinstance(item, dict):
                continue
            pid = str(item.get("proposal"))
            did = f"decision-{pid}"
            if did not in {d["id"] for d in dl.readable_docs(self.society, role, prefix="decision-")}:
                invalid.append(f"stance on {pid}, which {role} cannot read")
                continue
            if self.inst.proposals.get(pid, {}).get("by") == role and item.get("stance") in inst_mod.APPROVING:
                continue  # the proposer's approval is already on record
            stance, note = str(item.get("stance")), str(item.get("note") or "")[:200]
            event = self.inst.stance(r, role, pid, stance, note)
            if event["event"] == "refused":
                invalid.append(f"stance refused: {event['reason']}")
                continue
            dl.comment_doc(self.society, did, role, f"{stance}: {note}")
            applied["stances"].append({"proposal": pid, "stance": stance, "binding": event["binding"]})
            status = self.inst.proposals[pid]["status"]
            if status in ("blocked", "ratified"):
                doc = self.society.world.service_state("docs")["documents"][did]
                notify = [a for a in doc["readers"] if a != role]
                verb = "blocked" if status == "blocked" else "approved; it is ratified"
                dl.send_mail(self.society, role, notify, f"Decision record {pid}: {verb}",
                             f"{pid} ({self.inst.proposals[pid]['kind']}) is {verb}. {note}")
            if self.inst.done:
                break
        events = self.inst.events[first_event:]
        line = self.write({"type": "turn", "round": r, "role": role, "key": key, "sample": self.sample,
                           "read": unread, "reply": obj, "assessment": assessment, "applied": applied,
                           "invalid": invalid, "events": events})
        summary = []
        if applied["messages"]:
            summary.append("wrote to " + "; ".join(", ".join(m["to"]) + (f" (forwarding {m['forward_of']})" if m["forward_of"] else "")
                                                     for m in applied["messages"]))
        if applied["diligence"]:
            summary.append("requested " + ", ".join(d["kind"] for d in applied["diligence"]) + " diligence")
        if applied["proposal"]:
            summary.append(f"filed {applied['proposal']['id']} ({applied['proposal']['kind']})")
        if applied["stances"]:
            summary.append("; ".join(f"{s['stance']} on {s['proposal']}" for s in applied["stances"]))
        if assessment.get("support"):
            summary.append(f"your support was {assessment['support']}/7")
        self.memory[role].append(f"week {r + 1}: " + ("; ".join(summary) or "did nothing") +
                                 (f". Why: {str(obj.get('rationale'))[:200]}" if obj.get("rationale") else ""))
        return line

    # -- the run ----------------------------------------------------------------------

    def deliver(self, r):
        for item in [p for p in self.pending if p[0] == r]:
            _, kind, requester = item
            did = f"diligence-{kind}"
            docs = self.society.world.service_state("docs").get("documents", {})
            if did not in docs:
                dl.put_doc(self.society, did, f"Velliton {kind} diligence", self.vendor,
                           self.study["vendor"]["diligence"][kind], readers=[requester])
            else:
                dl.share_doc(self.society, did, [requester])
            mid = dl.send_mail(self.society, self.vendor, [requester], f"Velliton {kind} diligence",
                               self.study["vendor"]["diligence"][kind] + f"\n[shared with you: {did}]")
            self.inst.record_fact(r, f"diligence:{kind}", by=requester)
            self.write({"type": "diligence_delivered", "round": r, "kind": kind, "to": requester, "mail": mid, "doc": did})
        self.pending = [p for p in self.pending if p[0] != r]

    def available(self, role, r):
        return streams.uniform(self.seed, f"available/{role}", r) < float(self.study.get("availability", 1.0))

    def run(self):
        s = self.society
        rounds = s.rounds
        for r in range(rounds):
            s.round = r
            aligned = s.advance_to(s.base_tick + r * s.quantum)
            self.deliver(r)
            wanted = [role for role in self.roles if self.needs_turn(role)]
            busy = [role for role in wanted if not self.available(role, r)]
            order = streams.Stream(self.seed, f"order/{r}").shuffle(sorted(set(wanted) - set(busy)))
            self.write({"type": "round_start", "round": r, "wanted": wanted, "busy": busy, "order": order,
                        "aligned": aligned})
            for role in order:
                if self.inst.done:
                    break
                self.take_turn(role, r)
            self.write({"type": "round_end", "round": r, "state_hash": s.world.state_hash(), "aware": self.aware(),
                        "open": [p["id"] for p in self.inst.open_proposals()], "terminal": self.inst.terminal})
            if self.inst.done:
                break
        if not self.inst.done:
            self.inst.close(rounds - 1)
        support = {role: (self.assessments[role][-1].get("support") if self.assessments[role] else None)
                   for role in self.roles}
        final = {"type": "terminal", "terminal": self.inst.terminal, "round": self.inst.terminal_round,
                 "proposal": self.inst.terminal_proposal, "institution": self.inst.summary(),
                 "contested": self.inst.contested(), "last_support": support,
                 "state_hash": s.world.state_hash(), "turns": self.turns}
        self.write(final)
        return final


def prepare(spec, spec_path, seed, journal_for):
    """Boot one seed's base world: mailboxes and the vendor's fact sheet, round 0's
    state before any arm's pitch."""
    compiled = schema.compile_spec(spec, spec_path, seed)
    journal = journal_for(compiled)
    prefix = runner.Trace(None)
    base = runner.Society(compiled, journal, prefix, arm="prefix").boot()
    roles = spec["institution"]["roles"]
    vendor = spec["study"]["vendor"]
    dl.ensure_users(base, roles + [vendor["id"]])
    dl.put_doc(base, "velliton-facts", "Velliton fact sheet", vendor["id"], vendor["facts"], readers=["sponsor"])
    return compiled, base


def run_seed(spec, spec_path, seed, out, journal_for, pool=None):
    compiled, base = prepare(spec, spec_path, seed, journal_for)
    prefix_hash = base.world.state_hash()
    runs = {}
    for arm in spec["interventions"]["arms"]:
        trace = runner.Trace(out / f"seed-{seed}" / f"{arm['id']}.jsonl")
        society = base.fork(arm["id"], trace)
        applied = experiment.apply_ops(society, arm)
        runs[arm["id"]] = Run(society, compiled, arm, seed, base.journal, trace)
        runs[arm["id"]].line = len(trace.records)
        del applied
    if pool is None:
        finals = {a: r.run() for a, r in runs.items()}
    else:
        futures = {a: pool.submit(r.run) for a, r in runs.items()}
        finals = {a: f.result() for a, f in futures.items()}
    return {"seed": seed, "prefix_state_hash": prefix_hash, "arms": {
        a: {"terminal": f["terminal"], "round": f["round"], "turns": f["turns"], "state_hash": f["state_hash"],
            "state_hashes": [rec["state_hash"] for rec in runs[a].trace.records if rec["type"] == "round_end"]}
        for a, f in finals.items()}}


def preregistration(spec):
    material = {k: spec.get(k) for k in ("outcomes", "interventions", "institution", "study")}
    return hashlib.sha256(json.dumps(material, sort_keys=True).encode()).hexdigest()


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--spec", default=str(HERE / "spec.json"))
    parser.add_argument("--out", required=True)
    parser.add_argument("--seeds", help="comma-separated seeds (default: the spec's sweep.replicates)")
    parser.add_argument("--journal", help="the LLM journal (default <out>/journal.jsonl)")
    parser.add_argument("--offline", action="store_true", help="answer only from the journal")
    parser.add_argument("--max-calls", type=int, default=2500, help="the most model calls this invocation may make")
    parser.add_argument("--workers", type=int, default=8, help="seeds run at once (each runs its arms in threads)")
    parser.add_argument("--expect", help="a summary.json this run must reproduce exactly")
    args = parser.parse_args(argv)
    spec, spec_path = load_spec(args.spec)
    out = pathlib.Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    seeds = [int(s) for s in args.seeds.split(",")] if args.seeds else spec["sweep"]["replicates"]
    journal_path = pathlib.Path(args.journal) if args.journal else out / "journal.jsonl"
    holder = {}

    def journal_for(compiled):
        pin = runner.journal_pin(compiled)
        if "journal" not in holder:
            holder["pin"] = pin
            holder["journal"] = BoundedJournal(journal_path, offline=args.offline, pin=pin, max_calls=args.max_calls)
        elif holder["pin"] != pin:
            raise runner.JournalMismatch("seeds compiled to different worlds; one journal cannot serve them")
        return holder["journal"]
    try:
        compiled0 = schema.compile_spec(spec, spec_path, seeds[0])
        journal_for(compiled0)
    except runner.JournalMismatch as e:
        print(f"error: {e}", file=sys.stderr)
        return 2
    journal = holder["journal"]
    if not args.offline and not journal.client.api_key:
        print("error: no DEEPSEEK_API_KEY in the environment or .env; to replay the checked-in journal use --offline",
              file=sys.stderr)
        return 2
    started = time.time()
    seeds_out = {}
    arm_pool = concurrent.futures.ThreadPoolExecutor(max(1, args.workers) * len(spec["interventions"]["arms"]))
    with concurrent.futures.ThreadPoolExecutor(max(1, args.workers)) as pool:
        futures = {seed: pool.submit(run_seed, spec, spec_path, seed, out, journal_for, arm_pool) for seed in seeds}
        failed = []
        for seed, future in futures.items():
            try:
                seeds_out[seed] = future.result()
                print(json.dumps({"seed": seed, "terminals": {a: v["terminal"] for a, v in seeds_out[seed]["arms"].items()}}),
                      flush=True)
            except Exception as e:  # report every seed's failure, not just the first
                failed.append(seed)
                print(f"seed {seed} failed: {type(e).__name__}: {e}", file=sys.stderr, flush=True)
    arm_pool.shutdown()
    summary = {"study": spec["id"], "engine": runner.computerworld.engine_version, "model": journal.model,
               "preregistration": preregistration(spec), "seeds": sorted(seeds_out),
               "runs": {str(s): v for s, v in sorted(seeds_out.items())},
               "journal": journal.stats(), "elapsed_s": round(time.time() - started, 1), "pin": holder["pin"]}
    (out / "summary.json").write_text(json.dumps(summary, indent=1, sort_keys=True) + "\n")
    print(json.dumps({"journal": summary["journal"], "elapsed_s": summary["elapsed_s"]}))
    status = 1 if failed else 0
    if args.expect:
        expected = json.loads(pathlib.Path(args.expect).read_text())
        problems = []
        for seed, data in expected["runs"].items():
            got = summary["runs"].get(seed)
            if got is None:
                problems.append(f"seed {seed} missing")
                continue
            if got["prefix_state_hash"] != data["prefix_state_hash"]:
                problems.append(f"seed {seed}: prefix state hash differs")
            for arm, value in data["arms"].items():
                mine = got["arms"].get(arm, {})
                for field in ("terminal", "round", "state_hashes"):
                    if mine.get(field) != value[field]:
                        problems.append(f"seed {seed} {arm}: {field} differs")
        print("replay: " + ("identical" if not problems else "; ".join(problems[:10])))
        status = status or (1 if problems else 0)
    return status


if __name__ == "__main__":
    sys.exit(main())
