"""Scheduler: assignment manifests (paired tasks × arms × repeats, randomised order, blocked by task and execution window),
and the per-run pipeline (prepare → launch → controller-enforced gate cycles (C/D) → separate review (D) → collect →
contract check → evaluate). Runs are executed concurrently across executor nodes with a per-node concurrency limit."""
from __future__ import annotations
import json, os, posixpath, random, shlex, threading, time, traceback, concurrent.futures as cf
from .contracts import RunManifest, Budget, ARMS, utcnow, sha256_file, sha256_text, new_run_id, new_session_id
from .budgets import Envelope
from . import launcher, isolation, provenance
from .evaluation import evaluate_run

class Controller:
    def __init__(self, repo_root: str, store, config: dict, nodes: list[str], remote_root="/Users/simon-pierreboucher/wp29", log=print):
        self.repo = repo_root; self.store = store; self.cfg = config; self.nodes = nodes; self.remote_root = remote_root; self.log = log
        self.arms_dir = os.path.join(repo_root, "protocol", "arms")
        self._node_sems = {n: threading.Semaphore(int(config.get("concurrency_per_node", 2))) for n in nodes}
        self._rr = 0; self._rr_lock = threading.Lock()
        self.model = config.get("model", "claude-opus-5-5")
        self.effort = config.get("effort")
        self.budget = Budget(**config.get("budget", {}))
        self._halt = threading.Event()

    # ---------- assignment
    def make_assignment(self, split: str, task_ids: list[str], repeats: int, seed: int, arms=ARMS, start_seq=1) -> list[RunManifest]:
        """Paired design: every task gets every arm × repeats. Order: tasks shuffled; within a task block, (arm, repeat) pairs shuffled.
        Blocks are interleaved so that no arm runs systematically earlier than another (temporal blocking)."""
        rng = random.Random(seed)
        tasks = list(task_ids); rng.shuffle(tasks)
        runs = []; seq = start_seq
        for bi, t in enumerate(tasks):
            cells = [(a, r) for a in arms for r in range(repeats)]; rng.shuffle(cells)
            for oi, (a, r) in enumerate(cells):
                runs.append(RunManifest(run_id=new_run_id(split, seq), task_id=t, split=split, treatment=a, repeat_index=r,
                                        block_id=f"{split}-B{bi+1:02d}", order_in_block=oi, model_requested=self.model,
                                        budget=self.budget.__dict__.copy(), notes=f"assignment seed {seed}"))
                seq += 1
        return runs

    def register(self, runs: list[RunManifest]) -> None:
        for m in runs:
            m.code_commit = provenance.code_commit(self.repo)
            self.store.create_run(m)

    # ---------- helpers
    def _pick_node(self) -> str:
        with self._rr_lock:
            n = self.nodes[self._rr % len(self.nodes)]; self._rr += 1; return n

    def _read(self, *parts) -> str:
        return open(os.path.join(self.repo, *parts), encoding="utf-8").read()

    def _bundle_dir(self, task_id) -> str:
        return os.path.join(self.repo, "tasks", "public_bundles", task_id)

    def _run_dir(self, m) -> str:
        return os.path.join(self.repo, "runs", m["split"] if isinstance(m, dict) else m.split, m["run_id"] if isinstance(m, dict) else m.run_id)

    # ---------- pipeline
    def execute(self, run_id: str, fault=None) -> dict:
        """Run the whole pipeline for one registered run; always returns, never raises (errors are recorded)."""
        m = self.store.get_run(run_id)
        node = m.get("node") or self._pick_node()
        with self._node_sems[node]:
            try:
                return self._execute(run_id, node, fault)
            except Exception as e:  # noqa: BLE001
                tb = traceback.format_exc()[-1500:]
                self.log(f"[{run_id}] INFRA ERROR: {e}")
                try:
                    cur = self.store.get_run(run_id)["status"]
                    if cur in ("not_started", "prepared", "running", "interrupted"):
                        self.store.transition(run_id, "error", note=str(e)[:300], terminal_reason="infrastructure_error:" + str(e)[:200], ended_at_utc=utcnow())
                except Exception: pass  # noqa: BLE001
                return {"run_id": run_id, "status": "error", "error": str(e), "traceback": tb}

    def _execute(self, run_id: str, node: str, fault) -> dict:
        m = self.store.get_run(run_id); arm = m["treatment"]; task_id = m["task_id"]
        rd = self._run_dir(m); os.makedirs(rd, exist_ok=True)
        attempt = self.store.n_error_attempts(run_id)
        if not self.store.log_job(f"{run_id}:execute:attempt{attempt}", run_id, "execute", node):
            # duplicate submission: idempotency guard
            self.log(f"[{run_id}] duplicate execute request ignored (job already logged)")
            return {"run_id": run_id, "status": "duplicate_ignored"}
        # 1. prepare workspace
        bundle = self._bundle_dir(task_id)
        prep = isolation.prepare_workspace(node, self.remote_root, run_id, bundle)
        if not prep["verified_identical"]:
            raise RuntimeError(f"workspace manifest mismatch: extra={prep['remote_extra']} missing={prep['remote_missing']}")
        ws, ctl = prep["workspace"], prep["ctl_dir"]
        task_md = self._read("tasks", "public_bundles", task_id, "TASK.md")
        common = self._read("protocol", "arms", "common_preamble.md"); arm_text = self._read("protocol", "arms", f"arm_{arm}.md")
        task_prompt = self._read("protocol", "arms", "task_prompt.md")
        sys_append, ctx_hash = isolation.effective_context(common, arm_text, task_md, task_prompt)
        isolation.write_ctl_files(node, ctl, {"system_append.md": sys_append, "prompt.txt": task_prompt,
                                              "reviewer_system.md": self._read("protocol", "arms", "reviewer.md"),
                                              "reviewer_prompt.txt": "Review the deliverables in this workspace now, as instructed, and write out/review/review.json.\n"})
        env = provenance.node_environment(node, launcher.ssh)
        self.store.update_manifest(run_id, node=node, task_hash=sha256_file(os.path.join(bundle, "bundle.json")),
                                   protocol_hash=sha256_text(sys_append), effective_context_hash=ctx_hash,
                                   environment_id=env["environment_id"], client_version=env.get("claude"), model_requested=self.model,
                                   code_commit=provenance.code_commit(self.repo))
        json.dump({"workspace_manifest": prep["manifest"], "node_environment": env, "system_append_sha256": sha256_text(sys_append)}, open(os.path.join(rd, "prepare.json"), "w"), indent=1)
        open(os.path.join(rd, "system_append.md"), "w").write(sys_append)
        self.store.transition(run_id, "prepared", note=f"workspace {ws} on {node}")
        # 2. executor session(s)
        envl = Envelope(self.budget.max_budget_usd, self.budget.wall_seconds, self.budget.max_turns_per_session, min_continuation_usd=self.budget.min_continuation_usd)
        raw = os.path.join(rd, "raw_stream.jsonl")
        sid = new_session_id(); seq = 0; sessions = []; status = "completed"; reason = "executor_finished"
        self.store.transition(run_id, "running", note="first session", started_at_utc=utcnow(), session_ids=[sid])
        prompt_file = posixpath.join(ctl, "prompt.txt"); sys_file = posixpath.join(ctl, "system_append.md")
        fault_rec = None
        r, seq = self._session(run_id, node, ws, sid, prompt_file, sys_file, envl, raw, seq, role="executor", kind="initial", fault=fault)
        sessions.append(sid)
        if r.get("fault_fired"):
            fault_rec = r["fault_fired"]
            # standardized interruption → resume once from the same session
            self.store.transition(run_id, "interrupted", note="fault: " + fault_rec["fault_type"])
            if envl.can_continue():
                isolation.write_ctl_files(node, ctl, {"resume_prompt.txt": self._read("protocol", "arms", "resume_after_interruption_prompt.md")})
                self.store.transition(run_id, "running", note="resume after interruption")
                r, seq = self._session(run_id, node, ws, sid, posixpath.join(ctl, "resume_prompt.txt"), sys_file, envl, raw, seq, role="executor", kind="resume_after_fault", resume=True)
        status, reason = self._status_from(r, envl)
        # 3. controller-enforced gate cycles for C/D
        gate_receipts = []
        if arm in ("C", "D") and status == "completed":
            for cycle in range(1, self.budget.max_gate_cycles + 1):
                rec = self.remote_gate(node, ws)
                self.store.add_gate_receipt(run_id, cycle, "controller", rec); gate_receipts.append(rec)
                if rec.get("passed"):
                    break
                if not envl.can_continue():
                    reason = "gate_failed_budget_exhausted"; break
                fb = self._read("protocol", "arms", "gate_feedback_prompt.md") + self._gate_text(rec)
                isolation.write_ctl_files(node, ctl, {f"gate_feedback_{cycle}.txt": fb})
                r, seq = self._session(run_id, node, ws, sid, posixpath.join(ctl, f"gate_feedback_{cycle}.txt"), sys_file, envl, raw, seq, role="executor", kind=f"gate_cycle_{cycle}", resume=True)
                status, reason = self._status_from(r, envl)
                if status != "completed":
                    break
            else:
                rec = self.remote_gate(node, ws); self.store.add_gate_receipt(run_id, self.budget.max_gate_cycles + 1, "controller_final", rec); gate_receipts.append(rec)
        # 4. separate review for D
        review = None
        if arm == "D" and status == "completed" and gate_receipts and gate_receipts[-1].get("passed") and envl.can_continue():
            rsid = new_session_id(); sessions.append(rsid)
            launcher.ssh(node, f"mkdir -p {shlex.quote(ws)}/out/review")
            rr, seq = self._session(run_id, node, ws, rsid, posixpath.join(ctl, "reviewer_prompt.txt"), posixpath.join(ctl, "reviewer_system.md"), envl, raw, seq, role="reviewer", kind="review")
            review = self._fetch_json(node, posixpath.join(ws, "out", "review", "review.json"))
            self.store.add_review(run_id, rsid, review)
            if review and review.get("verdict") == "revise" and envl.can_continue():
                fb = self._read("protocol", "arms", "review_feedback_prompt.md") + json.dumps(review, indent=1)[:6000]
                isolation.write_ctl_files(node, ctl, {"review_feedback.txt": fb})
                r, seq = self._session(run_id, node, ws, sid, posixpath.join(ctl, "review_feedback.txt"), sys_file, envl, raw, seq, role="executor", kind="review_repair", resume=True)
                status, reason = self._status_from(r, envl)
                rec = self.remote_gate(node, ws); self.store.add_gate_receipt(run_id, len(gate_receipts) + 1, "controller_post_review", rec); gate_receipts.append(rec)
        # 5. collect artifacts and transcripts
        out_local = os.path.join(rd, "out")
        launcher.rsync_from(node, posixpath.join(ws, "out"), out_local)
        launcher.rsync_from(node, ctl, os.path.join(rd, "ctl"))
        self._collect_transcripts(node, ws, sessions, os.path.join(rd, "sessions"))
        am = provenance.artifact_manifest(out_local)
        json.dump(am, open(os.path.join(rd, "artifact_manifest.json"), "w"), indent=1)
        # 6. common contract check for ALL arms (post hoc; not fed back) + behavioural escape scan
        contract = self.remote_gate(node, ws)
        self.store.add_gate_receipt(run_id, 0, "final_common_contract_check", contract)
        esc = isolation.scan_raw_stream(raw, ws)
        json.dump(esc, open(os.path.join(rd, "behaviour_audit.json"), "w"), indent=1)
        self.store.transition(run_id, status, note=reason, terminal_reason=reason, ended_at_utc=utcnow(), session_ids=sessions,
                              usage=envl.summary(), artifact_manifest=os.path.join("runs", m["split"], run_id, "artifact_manifest.json"),
                              fault=fault_rec, model_reported=self._model_reported(run_id))
        # 7. evaluate (laptop side, hidden oracle)
        ev = evaluate_run(self.repo, self.store, run_id, contract_receipt=contract)
        self.log(f"[{run_id}] {task_id} arm {arm} on {node}: {status} ({reason}); cost {envl.spent_usd:.3f} USD; success={ev.get('scientific_success')}")
        return {"run_id": run_id, "status": status, "reason": reason, "cost": envl.spent_usd, "evaluation": ev}

    def _status_from(self, r, envl):
        if r.get("timed_out"):
            return "timeout", "wall_clock_limit"
        sub = (r.get("result") or {}).get("subtype")
        if sub == "error_max_budget_usd" or (envl.remaining_usd() <= 0.0):
            return "budget_exhausted", sub or "budget_exhausted"
        if sub == "error_max_turns":
            return "completed", "max_turns_reached"
        if r.get("exit_code") not in (0, None) and not r.get("result"):
            return "error", f"exit_code_{r.get('exit_code')}:{(r.get('stderr_tail') or '')[-200:]}"
        return "completed", "executor_finished" if sub == "success" else f"result_{sub}"

    def _session(self, run_id, node, ws, sid, prompt_file, sys_file, envl, raw, seq, role, kind, resume=False, fault=None):
        cap = min(envl.remaining_usd(), self.budget.max_budget_usd)
        argv = launcher.build_command(ws, ws + ".ctl", prompt_file, sys_file, self.model, sid, cap, self.budget.max_turns_per_session, resume=resume, effort=self.effort)
        job_key = f"{run_id}:{sid}:{kind}"
        if not self.store.log_job(job_key, run_id, "session", json.dumps({"node": node, "argv": argv})):
            raise RuntimeError(f"duplicate session job {job_key}")
        self.store.log_session(sid, run_id, role, kind, seq, {"node": node, "argv": argv, "cap_usd": cap, "wall_s": envl.remaining_wall()}, raw)
        if role == "executor" and resume is False and self.store.query("SELECT 1 FROM sessions WHERE session_id=? AND seq<>?", (sid, seq)):
            pass
        on_event = None; fired = {}
        if fault and fault.get("type") == "interrupt":
            trig = fault.get("trigger", {"after_tool_uses": 6}); counter = {"n": 0}
            def on_event(ev):
                if ev.get("event_type") == "assistant" and ev.get("tool_name"):
                    counter["n"] += 1
                    if counter["n"] == trig.get("after_tool_uses") and not fired:
                        fired.update({"fault_type": "interrupt", "trigger": trig, "at_event": ev["event_id"], "utc": utcnow()})
                        self.store.add_fault(run_id, "interrupt", json.dumps(trig), {"at_event": ev["event_id"]})
                        launcher.kill_remote(node, sid)
        r = launcher.run_session(node, ws, argv, raw, run_id, sid, self.store, envl.remaining_wall(), seq_start=seq, on_event=on_event)
        res = r.get("result") or {}
        prev_total = sum(s["cost_usd"] for s in envl.sessions if s["session_id"] == sid)
        total = float(res.get("total_cost_usd") or 0.0)
        session_cost = max(0.0, total - prev_total) if resume else total
        if not res and r.get("timed_out"):
            session_cost = self._estimate_cost_from_events(run_id, sid, seq, r["seq_end"])
        envl.charge(sid, role, session_cost, r["wall_s"], res.get("num_turns"))
        first_in = self.store.query("SELECT usage_json FROM events WHERE run_id=? AND session_id=? AND seq>? AND event_type='assistant' ORDER BY seq LIMIT 1", (run_id, sid, seq))
        fti = None
        if first_in and first_in[0]["usage_json"]:
            u = json.loads(first_in[0]["usage_json"]); fti = (u.get("input_tokens", 0) + u.get("cache_creation_input_tokens", 0) + u.get("cache_read_input_tokens", 0))
        self.store.finish_session(sid, started_utc=r["started_utc"], ended_utc=r["ended_utc"], exit_code=r["exit_code"], result_subtype=res.get("subtype"),
                                  is_error=int(bool(res.get("is_error"))), total_cost_usd=total, session_cost_usd=session_cost, num_turns=res.get("num_turns"),
                                  duration_ms=res.get("duration_ms"), usage_json=json.dumps(res.get("usage")), model_usage_json=json.dumps(res.get("modelUsage")),
                                  first_turn_input_tokens=fti, model_reported=",".join((res.get("modelUsage") or {}).keys()) or None,
                                  note=("timed_out" if r["timed_out"] else "") + (" fault_fired" if fired else ""))
        self.store.set_job(job_key, "done" if not r["timed_out"] else "timed_out")
        if fired:
            r["fault_fired"] = fired
        return r, r["seq_end"]

    def _estimate_cost_from_events(self, run_id, sid, seq0, seq1):
        """When a session is killed before its result line, sum per-message usage into a list-price estimate (recorded as such)."""
        rows = self.store.query("SELECT usage_json FROM events WHERE run_id=? AND session_id=? AND seq>? AND seq<=? AND event_type='assistant'", (run_id, sid, seq0, seq1))
        p = self.cfg.get("list_prices_per_mtok", {"input": 4.0, "output": 20.0, "cache_read": 0.20, "cache_write": 5.0})
        tot = 0.0
        for r in rows:
            if not r["usage_json"]: continue
            u = json.loads(r["usage_json"])
            tot += (u.get("input_tokens", 0) * p["input"] + u.get("output_tokens", 0) * p["output"] + u.get("cache_read_input_tokens", 0) * p["cache_read"] + u.get("cache_creation_input_tokens", 0) * p["cache_write"]) / 1e6
        return tot

    def _model_reported(self, run_id):
        rows = self.store.query("SELECT model_reported FROM sessions WHERE run_id=? AND model_reported IS NOT NULL", (run_id,))
        return ",".join(sorted({r["model_reported"] for r in rows})) or None

    def remote_gate(self, node, ws) -> dict:
        """Run the pristine gate copy installed on the node (outside the workspace) and return its receipt."""
        r = launcher.ssh(node, f"cd {shlex.quote(ws)} && ~/wp29/venv/bin/python ~/wp29/gates/pgar_gate.py . --python ~/wp29/venv/bin/python --json", timeout=900)
        try:
            return json.loads(r.stdout[r.stdout.index("{"):])
        except Exception:  # noqa: BLE001
            return {"passed": False, "n_checks": 0, "n_failed": 1, "checks": [{"id": "gate_execution", "passed": False, "detail": (r.stderr or r.stdout)[-500:]}], "gate_version": "unavailable"}

    @staticmethod
    def _gate_text(rec) -> str:
        lines = [f"PGAR GATE — {'PASSED' if rec.get('passed') else 'FAILED'} ({rec.get('n_failed')} of {rec.get('n_checks')} checks failed)"]
        for c in rec.get("checks", []):
            if not c.get("passed"):
                lines.append(f"  [FAIL] {c['id']}: {c.get('detail')}")
        return "\n".join(lines) + "\n"

    def _fetch_json(self, node, path):
        r = launcher.ssh(node, f"cat {shlex.quote(path)} 2>/dev/null")
        try: return json.loads(r.stdout)
        except Exception: return None  # noqa: BLE001

    def _collect_transcripts(self, node, ws, session_ids, dst):
        os.makedirs(dst, exist_ok=True)
        for sid in session_ids:
            r = launcher.ssh(node, f"find ~/.claude/projects -name {shlex.quote(sid + '.jsonl')} 2>/dev/null | head -1")
            p = r.stdout.strip()
            if p:
                launcher.rsync_from(node, os.path.dirname(p), dst, excludes=["*"]) if False else None
                import subprocess
                subprocess.run(["scp", "-q", "-o", "BatchMode=yes", f"{node}:{p}", os.path.join(dst, sid + ".jsonl")], capture_output=True, timeout=300)

    # ---------- batch execution
    def run_many(self, run_ids: list[str], max_workers: int | None = None, faults: dict | None = None) -> list[dict]:
        max_workers = max_workers or len(self.nodes) * int(self.cfg.get("concurrency_per_node", 2))
        out = []
        with cf.ThreadPoolExecutor(max_workers=max_workers) as ex:
            futs = {ex.submit(self.execute, rid, (faults or {}).get(rid)): rid for rid in run_ids}
            for f in cf.as_completed(futs):
                out.append(f.result())
        return out