#!/usr/bin/env python3
"""pi-task — run a headless pi subtask from an immutable SPEC, then verify the result.

Why this exists: `fork` hands the child getHeader()+getBranch(), i.e. the whole
untrimmed parent branch, so in a long session the child continues the PARENT's
narrative instead of doing the task (measured 2026-09-06, mbp-m1-2020: 4/4
dispatches ignored their brief; one filed a diary entry as the parent). This
runner inverts that default: context is EXPLICIT and default-EMPTY, the child is
a fresh isolated session, and its answer must parse as a declared envelope or the
run FAILS.

Honest limits, stated so nobody mistakes this for a sandbox:
  * pi has no tool allow/deny list, so --no-extensions removes EXTENSIONS
    (hence the palace write path) but NOT core read/write/edit/bash.
    "read_only" is therefore VERIFIED POST-HOC by a boundary diff, not enforced.
  * cost is checked after the fact; pi has no spend ceiling to pass in.

Usage:
    pi-task run SPEC.json [--dry-run] [--audit-dir DIR]
    pi-task schema
"""

from __future__ import annotations

import argparse, hashlib, json, os, re, shutil, subprocess, sys, time
from datetime import datetime, timezone
from pathlib import Path

SETTINGS = Path.home() / ".pi/agent/settings.json"
AUDIT_ROOT = Path(os.environ.get("PI_TASK_AUDIT_DIR", Path.home() / ".pi/agent/pi-task"))
REQUIRED_SPEC = ("id", "goal", "deliverable", "effort")
ENVELOPE_KEYS = ("status", "deliverable", "evidence", "unsure", "did_not_do")
FENCE = re.compile(r"```json\s*(\{.*?\})\s*```", re.DOTALL)


def die(msg: str, code: int = 2):
    print(f"pi-task: FAIL: {msg}", file=sys.stderr)
    sys.exit(code)


def sh(args, cwd=None) -> str:
    p = subprocess.run(args, cwd=cwd, capture_output=True, text=True)
    return p.stdout.strip()


# ---------------------------------------------------------------- boundary diff
def _git_porcelain(r: str) -> dict:
    """--ignored is NOT optional: plain `git status --porcelain` omits ignored
    files, so a child writing .env, credentials or build artifacts into a root
    reads back as CLEAN. Measured 2026-09-07: an ignored secret.txt produced zero
    porcelain lines, and `!! secret.txt` with --ignored.
    Huge repos (node_modules) can emit megabytes, so always keep a sha256 and the
    text only when it is small enough to show a human."""
    txt = sh(["git", "-C", r, "status", "--porcelain", "--ignored"])
    return {"sha": hashlib.sha256(txt.encode()).hexdigest()[:16],
            "lines": len(txt.splitlines()),
            "text": txt if len(txt) <= 8000 else None}


def boundary(roots) -> dict:
    """Cheap, checkable state of each root: git HEAD + porcelain, else file manifest."""
    snap = {}
    for r in roots:
        root = Path(r)
        if not root.exists():
            snap[r] = {"kind": "missing"}
        elif (root / ".git").exists():
            snap[r] = {"kind": "git",
                       "head": sh(["git", "-C", r, "rev-parse", "HEAD"]),
                       "porcelain": _git_porcelain(r)}
        else:
            man = {}
            for f in sorted(root.rglob("*")):
                if f.is_file():
                    man[str(f)] = hashlib.sha256(f.read_bytes()).hexdigest()[:16]
            snap[r] = {"kind": "manifest", "files": man}
    return snap


def boundary_delta(before: dict, after: dict) -> list:
    out = []
    for r in before:
        b, a = before[r], after.get(r, {})
        if b == a:
            continue
        if b.get("kind") != a.get("kind"):
            out.append(f"{r}: root kind changed {b.get('kind')} -> {a.get('kind')}")
            continue
        if b.get("kind") == "git":
            if b.get("head") != a.get("head"):
                out.append(f"{r}: HEAD {b.get('head','?')[:8]} -> {a.get('head','?')[:8]}")
            pb, pa = b.get("porcelain") or {}, a.get("porcelain") or {}
            if pb.get("sha") != pa.get("sha"):
                if pa.get("text") is not None:
                    out.append(f"{r}: working tree changed (incl. ignored):\n{pa['text']}")
                else:
                    out.append(f"{r}: working tree changed (incl. ignored): "
                               f"{pb.get('lines')} -> {pa.get('lines')} entries, "
                               f"sha {pb.get('sha')} -> {pa.get('sha')} (text too large to show)")
        else:
            bf, af = b.get("files", {}), a.get("files", {})
            for p in sorted(set(bf) | set(af)):
                if bf.get(p) != af.get(p):
                    out.append(f"{r}: {'added' if p not in bf else 'removed' if p not in af else 'modified'} {p}")
    return out


# -------------------------------------------------------------------- the brief
def writable_roots(spec: dict) -> list:
    """Which roots the child is allowed to change.

    `roots` is the WATCHED set; `write_allowed` is the CHANGEABLE subset. Keeping
    them separate is what makes a violation detectable at all: if the two are the
    same list, a write task can never trip the diff, and the only way to provoke
    one is to order the child to defy its own authority line — which a
    well-behaved child simply refuses (measured 2026-09-07, test T2). The real
    hazard is not defiance anyway; it is an INCIDENTAL write, e.g. a verification
    command that drops __pycache__ into a repo the child was only meant to read.
    """
    if spec.get("read_only", True):
        return []
    aw = spec.get("write_allowed")
    return list(aw) if aw else list(spec.get("roots") or [])
def build_prompt(spec: dict) -> str:
    ctx = spec.get("context") or {}
    roots = spec.get("roots") or []
    allowed = writable_roots(spec)
    lines = [
        "You are a TASK WORKER invoked by another agent. You are NOT continuing a",
        "conversation and you have NO shared history: there is no 'earlier', no",
        "'as we discussed', and no session to wind down. Do the task below and",
        "return the envelope. Nothing else.",
        "",
        f"## GOAL\n{spec['goal']}",
        f"\n## DELIVERABLE\n{spec['deliverable']}",
    ]
    if spec.get("read_only", True):
        lines += ["\n## AUTHORITY\nREAD-ONLY. Do not create, edit or delete any file. Do not run any",
                  "command that mutates state (no git commit/push, no installs). A boundary",
                  "diff runs after you exit and a violation fails the whole task."]
    else:
        lines += ["\n## AUTHORITY\nYou MAY create and modify files under:"]
        lines += [f"  - {a}" for a in allowed]
        watched_only = [r for r in roots if r not in allowed]
        if watched_only:
            lines += ["You may READ these, but must NOT change anything under them:"]
            lines += [f"  - {r}" for r in watched_only]
        lines += ["A boundary diff runs after you exit over every path above. ANY change",
                  "outside the writable set fails the task — including one made incidentally",
                  "by a command you ran rather than by an edit you intended."]
    if ctx.get("facts"):
        lines += ["\n## VERIFIED CONTEXT (asserted by the caller; you need not re-derive it)"]
        lines += [f"- {f}" for f in ctx["facts"]]
    if ctx.get("files"):
        lines += ["\n## RELEVANT PATHS (read them yourself; they are not pasted here)"]
        lines += [f"- {f}" for f in ctx["files"]]
    if ctx.get("commands"):
        lines += ["\n## SUGGESTED COMMANDS"]
        lines += [f"- {c}" for c in ctx["commands"]]
    lines += [
        "",
        "## REQUIRED OUTPUT — a single fenced json block, LAST thing you emit",
        "If it does not parse against this shape the task is recorded as FAILED,",
        "however good the prose was:",
        "```json",
        json.dumps({
            "status": "ok | partial | failed",
            "deliverable": "<the answer asked for, as text>",
            "evidence": [{"claim": "<one claim>", "pointer": "<file:line | exact command | url>"}],
            "unsure": ["<anything you could not establish — empty list only if truly none>"],
            "did_not_do": ["<anything in scope you skipped>"],
        }, indent=2),
        "```",
        "Rules for evidence: every load-bearing claim needs a pointer a third party",
        "can re-run or open. Do not invent quantities ('all N files') you did not count.",
        "If you could not do the task, status=failed with the reason is a CORRECT answer.",
    ]
    return "\n".join(lines)


# ------------------------------------------------------- envelope validation
def validate_envelope(text: str):
    """(envelope|None, problems[]). Deterministic and model-free so `selftest`
    can exercise it two-sidedly — a validator that has only ever returned PASS
    has not been shown to discriminate."""
    problems, env = [], None
    blocks = FENCE.findall(text or "")
    if not blocks:
        problems.append("no fenced json envelope in the final assistant message")
        return None, problems
    try:
        env = json.loads(blocks[-1])
    except json.JSONDecodeError as e:
        problems.append(f"envelope is not valid json: {e}")
        return None, problems
    if not isinstance(env, dict):
        problems.append("envelope is not a json object")
        return None, problems
    for k in ENVELOPE_KEYS:
        if k not in env:
            problems.append(f"envelope missing key '{k}'")
    if env.get("status") not in ("ok", "partial", "failed"):
        problems.append(f"envelope status not ok|partial|failed: {env.get('status')!r}")
    ev_list = env.get("evidence")
    if not isinstance(ev_list, list) or not ev_list:
        problems.append("envelope evidence is empty — no claim is checkable")
    else:
        for i, e in enumerate(ev_list):
            if not isinstance(e, dict) or not e.get("pointer"):
                problems.append(f"evidence[{i}] has no pointer")
    return env, problems


def cmd_selftest(_args):
    """Two-sided: known-good must pass, each known-bad must fail, and the
    boundary detector must both see a change and see no change. Aborts loudly
    if any pair fails to discriminate."""
    good = json.dumps({"status": "ok", "deliverable": "d",
                       "evidence": [{"claim": "c", "pointer": "f.py:1"}],
                       "unsure": [], "did_not_do": []})
    fixtures = [
        ("POSITIVE valid envelope", f"prose\n```json\n{good}\n```", True),
        ("no fence at all", "just fluent prose, no envelope", False),
        ("fence but broken json", "```json\n{not json,}\n```", False),
        ("missing key did_not_do", "```json\n" + json.dumps(
            {"status": "ok", "deliverable": "d",
             "evidence": [{"claim": "c", "pointer": "p"}], "unsure": []}) + "\n```", False),
        ("bad status value", "```json\n" + json.dumps(
            {"status": "done", "deliverable": "d",
             "evidence": [{"claim": "c", "pointer": "p"}],
             "unsure": [], "did_not_do": []}) + "\n```", False),
        ("empty evidence", "```json\n" + json.dumps(
            {"status": "ok", "deliverable": "d", "evidence": [],
             "unsure": [], "did_not_do": []}) + "\n```", False),
        ("evidence without pointer", "```json\n" + json.dumps(
            {"status": "ok", "deliverable": "d", "evidence": [{"claim": "c"}],
             "unsure": [], "did_not_do": []}) + "\n```", False),
        ("json object not last block wins", f"```json\n{{\"junk\":1}}\n```\n```json\n{good}\n```", True),
    ]
    fails = 0
    for name, text, want_pass in fixtures:
        _, probs = validate_envelope(text)
        got_pass = not probs
        ok = got_pass == want_pass
        fails += 0 if ok else 1
        print(f"  {'ok ' if ok else 'BAD'}  {'expect PASS' if want_pass else 'expect FAIL'}  {name}"
              + ("" if ok else f"   <-- got {'PASS' if got_pass else 'FAIL'}: {probs}"))

    import tempfile
    with tempfile.TemporaryDirectory() as td:
        d = Path(td) / "root"; d.mkdir(); (d / "a.txt").write_text("1")
        b1 = boundary([str(d)])
        same = boundary_delta(b1, boundary([str(d)]))
        (d / "b.txt").write_text("2")
        diff = boundary_delta(b1, boundary([str(d)]))
        checks = [("boundary/manifest: unchanged root reports no delta", same == []),
                  ("boundary/manifest: added file IS detected", len(diff) == 1)]

        # The GIT path is what every real run uses, and it was previously
        # untested. The ignored-file case below was genuinely broken until
        # --ignored was added, so this is a regression test, not decoration.
        g = Path(td) / "repo"; g.mkdir()
        for cmd in (["git", "init", "-q", "."], ["git", "config", "user.email", "t@t"],
                    ["git", "config", "user.name", "t"]):
            subprocess.run(cmd, cwd=g, capture_output=True)
        (g / ".gitignore").write_text("secret.txt\n__pycache__/\n")
        (g / "tracked.txt").write_text("v1")
        subprocess.run(["git", "add", "-A"], cwd=g, capture_output=True)
        subprocess.run(["git", "commit", "-qm", "init"], cwd=g, capture_output=True)
        gb = boundary([str(g)])
        checks.append(("boundary/git: clean repo reports no delta",
                       boundary_delta(gb, boundary([str(g)])) == []))
        (g / "secret.txt").write_text("exfiltrated")
        checks.append(("boundary/git: IGNORED file IS detected (regression: --ignored)",
                       len(boundary_delta(gb, boundary([str(g)]))) == 1))
        (g / "secret.txt").unlink()
        (g / "tracked.txt").write_text("v2")
        checks.append(("boundary/git: modified tracked file IS detected",
                       len(boundary_delta(gb, boundary([str(g)]))) == 1))
        subprocess.run(["git", "commit", "-aqm", "v2"], cwd=g, capture_output=True)
        checks.append(("boundary/git: a COMMIT (HEAD move) IS detected",
                       any("HEAD" in x for x in boundary_delta(gb, boundary([str(g)])))))

        for name, cond in checks:
            fails += 0 if cond else 1
            print(f"  {'ok ' if cond else 'BAD'}  {name}")

    if fails:
        die(f"selftest: {fails} check(s) failed — the validator does not discriminate, "
            "so any PASS it reports is meaningless", 3)
    print(f"selftest: all {len(fixtures) + len(checks)} checks discriminate correctly")
    return 0


# ------------------------------------------------------------------------- main
def cmd_run(args):
    spec_path = Path(args.spec)
    spec = json.loads(spec_path.read_text())
    missing = [k for k in REQUIRED_SPEC if k not in spec]
    if missing:
        die(f"spec missing required keys: {missing}")

    profiles = (json.loads(SETTINGS.read_text()).get("pi-fork") or {}).get("effortProfiles") or {}
    prof = profiles.get(spec["effort"])
    if not prof:
        die(f"effort '{spec['effort']}' not in {SETTINGS}: pi-fork.effortProfiles ({list(profiles)})")

    stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    audit = Path(args.audit_dir) if args.audit_dir else AUDIT_ROOT / f"{stamp}-{spec['id']}"
    audit.mkdir(parents=True, exist_ok=True)
    prompt = build_prompt(spec)
    roots = spec.get("roots") or []
    budget = spec.get("budget") or {}
    wall = int(budget.get("wall_s", 600))

    argv = ["pi", "-p", "--mode", "json",
            "--provider", prof["provider"], "--model", prof["id"],
            "--thinking", prof.get("thinking", "off"),
            "--no-extensions",
            "--session-id", f"pitask-{spec['id']}-{stamp}",
            "--session-dir", str(audit / "sess"),
            prompt]

    (audit / "spec.json").write_text(json.dumps(spec, indent=2))
    (audit / "prompt.txt").write_text(prompt)
    (audit / "argv.json").write_text(json.dumps(argv[:-1] + ["<prompt in prompt.txt>"], indent=2))
    if args.dry_run:
        print(prompt); print(f"\n--- dry run, nothing launched. audit: {audit}"); return 0

    before = boundary(roots)
    (audit / "boundary-before.json").write_text(json.dumps(before, indent=2))
    t0 = time.time()
    try:
        p = subprocess.run(argv, capture_output=True, text=True, timeout=wall)
        raw, rc = p.stdout, p.returncode
    except subprocess.TimeoutExpired as e:
        raw, rc = (e.stdout or ""), 124
    elapsed = time.time() - t0
    (audit / "raw.ndjson").write_text(raw if isinstance(raw, str) else raw.decode())
    after = boundary(roots)
    (audit / "boundary-after.json").write_text(json.dumps(after, indent=2))

    # ---- parse the NDJSON: final assistant text, cost, tool calls
    events, text, cost, tools = [], "", 0.0, 0
    for line in raw.splitlines():
        try:
            events.append(json.loads(line))
        except json.JSONDecodeError:
            continue
    for ev in events:
        if ev.get("type") == "agent_end":
            for m in ev.get("messages", []):
                cost += (((m.get("usage") or {}).get("cost") or {}).get("total") or 0.0)
        if ev.get("type") == "turn_end":
            tools += len(ev.get("toolResults") or [])
        if ev.get("type") == "message_end" and (ev.get("message") or {}).get("role") == "assistant":
            text = "".join(c.get("text", "") for c in ev["message"].get("content", [])
                           if c.get("type") == "text") or text

    # ---- validate the envelope: unparseable == FAILED, no matter how fluent
    env, problems = validate_envelope(text)
    if rc != 0:
        problems.append(f"child exited rc={rc}" + (" (wall-clock budget)" if rc == 124 else ""))
    if not any(e.get("type") == "agent_end" for e in events):
        problems.append("no agent_end event — child did not complete a turn")
    delta = boundary_delta(before, after)
    allowed = writable_roots(spec)
    violations = [d for d in delta
                  if not any(d.startswith(f"{a}:") for a in allowed)]
    if violations:
        if spec.get("read_only", True):
            problems.append("BOUNDARY VIOLATION: read_only task mutated its roots")
        else:
            problems.append("BOUNDARY VIOLATION: task changed a watched root it was "
                            f"not authorised to write (authorised: {allowed or 'none'})")
    over = budget.get("usd")
    if over and cost > float(over):
        problems.append(f"cost ${cost:.4f} over budget ${float(over):.4f}")

    verdict = "PASS" if not problems else "FAIL"
    result = {"verdict": verdict, "problems": problems, "envelope": env,
              "metrics": {"cost_usd": round(cost, 6), "elapsed_s": round(elapsed, 1),
                          "tool_results": tools, "rc": rc, "model": prof["id"],
                          "effort": spec["effort"]},
              "boundary_delta": delta, "boundary_violations": violations,
              "audit_dir": str(audit),
              "raw_text_if_envelope_missing": None if env else text[-2000:]}
    (audit / "result.json").write_text(json.dumps(result, indent=2))

    # ---- report to the parent: verdict first, then what to spot-check
    print(f"pi-task {verdict}  id={spec['id']}  {prof['id']}  "
          f"${cost:.4f}  {elapsed:.0f}s  tools={tools}  rc={rc}")
    for p_ in problems:
        print(f"  ! {p_}")
    if env:
        print(f"  status={env.get('status')}")
        print("  deliverable:"); print("    " + str(env.get("deliverable", "")).replace("\n", "\n    "))
        print("  evidence (SPOT-CHECK THESE, do not trust the prose):")
        for e in env.get("evidence") or []:
            print(f"    - {e.get('claim','')}  <=  {e.get('pointer','')}")
        for k in ("unsure", "did_not_do"):
            for u in env.get(k) or []:
                print(f"  {k}: {u}")
    print(f"  audit: {audit}")
    return 0 if verdict == "PASS" else 1


def cmd_schema(_args):
    print(json.dumps({
        "id": "short-slug-used-in-session-id-and-audit-path",
        "goal": "verbatim task statement",
        "deliverable": "the exact shape of answer wanted",
        "effort": "fast | balanced | deep  (resolved via ~/.pi/agent/settings.json pi-fork.effortProfiles)",
        "read_only": True,
        "roots": ["/workspace/repo-the-task-touches"],
        "write_allowed": ["subset of roots the child may CHANGE; only meaningful when "
                          "read_only is false. Defaults to all of roots, which makes "
                          "violations undetectable — set it explicitly for write tasks."],
        "context": {"facts": ["verified fact the caller asserts"],
                    "files": ["/abs/path the child should read itself"],
                    "commands": ["exact command the child may run"]},
        "budget": {"wall_s": 600, "usd": 0.5},
    }, indent=2))
    return 0


def main():
    ap = argparse.ArgumentParser(prog="pi-task")
    sub = ap.add_subparsers(dest="cmd", required=True)
    r = sub.add_parser("run"); r.add_argument("spec")
    r.add_argument("--dry-run", action="store_true"); r.add_argument("--audit-dir")
    r.set_defaults(fn=cmd_run)
    s = sub.add_parser("schema"); s.set_defaults(fn=cmd_schema)
    t = sub.add_parser("selftest"); t.set_defaults(fn=cmd_selftest)
    a = ap.parse_args()
    if not shutil.which("pi"):
        die("pi not on PATH")
    sys.exit(a.fn(a))


if __name__ == "__main__":
    main()
