#!/usr/bin/env python3
"""One admitted candidate, exact golden A/B partition, retained producer receipts."""
import datetime
import hashlib
import json
import os
from pathlib import Path
import re
import shutil
import signal
import subprocess
import sys
import time
import tempfile

SHA = "49a08a07db7f5da0ca3eb77129d5504ae107f8a9"
ROOT = Path("/home/reavus/projects/spt-core/spt-core")
TREE = ROOT / ".worktrees/consumer-linux-49a08a07"
PROOF = ROOT / ".spt/preserved/304-handoff/consumer-linux-49a08a07"
PROOF.mkdir(exist_ok=True)
assert not (PROOF / "receipt.json").exists(), "consumer already has a receipt"
TARGET = TREE / "target"
FLOOR = 32 * 1024**3
ALLOWED_SPT_KEY = "SPT_TEST_EPHEMERAL_ADVISORY_PORTS"
INHERITED_ENV = dict(os.environ)
ENV = dict(INHERITED_ENV)
for key in INHERITED_ENV:
    if key.startswith(("OWL_", "SPT_")):
        ENV.pop(key, None)
        os.environ.pop(key, None)
ENVIRONMENT = {"allowed_owl_spt_names": [ALLOWED_SPT_KEY], "snapshots": []}
OWNED_TEMP_ROOTS = []
ENV.update(PATH="/home/reavus/.cargo/bin:" + ENV["PATH"], CARGO_BUILD_JOBS="2",
           CARGO_TARGET_DIR=str(TARGET), CARGO_INCREMENTAL="0",
           RUSTFLAGS="-C link-arg=-fuse-ld=mold", NEXTEST_PROFILE="default",
           SPT_TEST_EPHEMERAL_ADVISORY_PORTS="1", CARGO_TERM_COLOR="never", CI="true")
for name in ("CARGO_ENCODED_RUSTFLAGS", "NEXTEST_TEST_THREADS", "NEXTEST_RETRIES"):
    ENV.pop(name, None)
START = datetime.datetime.now(datetime.timezone.utc).isoformat()


def utc():
    return datetime.datetime.now(datetime.timezone.utc).isoformat()


def save(name, value):
    (PROOF / name).write_text(json.dumps(value, indent=2) + "\n")


def capture_environment(label, mapping=None, require_scrubbed=True):
    mapping = dict(ENV if mapping is None else mapping)
    forbidden = sorted(key for key in mapping if key.startswith(("OWL_", "SPT_")) and key != ALLOWED_SPT_KEY)
    ENVIRONMENT["snapshots"].append({"label": label, "utc": utc(),
        "environment_names": sorted(mapping),
        "owl_spt_keys": sorted(key for key in mapping if key.startswith(("OWL_", "SPT_"))),
        "forbidden_present": forbidden})
    save("environment.json", ENVIRONMENT)
    if require_scrubbed and forbidden:
        raise RuntimeError("IDENTITY_ENV_REFUSED " + repr(forbidden))


def source_guard():
    head = subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=TREE, text=True).strip()
    status = subprocess.check_output(["git", "status", "--porcelain=v1", "--untracked-files=all"], cwd=TREE, text=True)
    if head != SHA or status:
        raise RuntimeError(f"SOURCE_REFUSED head={head} expected={SHA} status={status!r}")
    return {"sha": head, "clean": True}


def census():
    active = []
    owned = []
    for directory in Path("/proc").iterdir():
        if not directory.name.isdigit():
            continue
        try:
            name = (directory / "comm").read_text().strip()
            exe = os.readlink(directory / "exe")
            row = {"pid": int(directory.name), "name": name, "exe": exe,
                   "cmd": (directory / "cmdline").read_bytes().replace(b"\0", b" ").decode(errors="replace"),
                   "stat": (directory / "stat").read_text()}
            if any(word in name for word in ("cargo", "rustc", "nextest", "Runner.Worker")) or "/deps/" in exe:
                active.append(row)
            if exe.startswith(str(TARGET) + "/"):
                owned.append(row)
        except (OSError, PermissionError):
            pass
    return {"utc": utc(), "free_bytes": shutil.disk_usage(TREE).free,
            "load": os.getloadavg(), "active": active, "owned_survivors": owned}


def admit(label):
    state = {**source_guard(), **census()}
    save(label + "-admission.json", state)
    if state["free_bytes"] < FLOOR or state["active"] or state["owned_survivors"]:
        raise RuntimeError("RESOURCE_REFUSED " + json.dumps(state))
    return state


def run(label, argv):
    capture_environment("before-" + label)
    started = utc()
    clock = time.monotonic()
    path = PROOF / (label + ".log")
    print(f"BEGIN {label} sha={SHA} utc={started} argv={json.dumps(argv)}", flush=True)
    with path.open("x") as stream:
        result = subprocess.Popen(argv, cwd=TREE, env=ENV, stdout=stream, stderr=subprocess.STDOUT, start_new_session=True)
        free_samples = []
        floor_crossed = False
        while result.poll() is None:
            free = shutil.disk_usage(TREE).free
            free_samples.append({"utc": utc(), "free_bytes": free})
            if free < FLOOR:
                floor_crossed = True
                os.killpg(result.pid, signal.SIGKILL)
                result.wait()
                break
            time.sleep(5)
        save(label + "-disk-samples.json", {"floor_bytes": FLOOR, "crossed": floor_crossed, "samples": free_samples})
    capture_environment("after-" + label)
    text = path.read_text(errors="replace")
    summaries = [line.strip() for line in text.splitlines() if re.search(r"\bSummary\b", line)]
    receipt = {"sha": SHA, "command": argv, "start_utc": started, "end_utc": utc(),
               "elapsed_s": round(time.monotonic() - clock, 3), "exit": result.returncode,
               "summary_lines": summaries, "log": str(path),
               "log_sha256": hashlib.sha256(path.read_bytes()).hexdigest()}
    save(label + "-receipt.json", receipt)
    print("END " + label + " " + json.dumps(receipt), flush=True)
    if floor_crossed:
        raise RuntimeError("DISK_FLOOR_CROSSED " + label)
    return receipt


def outside_git(path):
    result = subprocess.run(["git", "-C", str(path), "rev-parse", "--show-toplevel"], env=ENV, capture_output=True, text=True)
    with (PROOF / "temp-git-probes.jsonl").open("a") as log:
        log.write(json.dumps({"path": str(path), "exit": result.returncode, "stdout": result.stdout, "stderr": result.stderr}) + "\n")
    if result.returncode == 0 or "not a git repository" not in result.stderr:
        raise RuntimeError("TEMP_PREMISE_REFUSED " + str(path) + " " + result.stdout + result.stderr)


def reap_owned(label):
    rows = census()["owned_survivors"]
    actions = []
    def born(stat):
        return stat.split(") ", 1)[1].split()[19]
    for sig in (signal.SIGTERM, signal.SIGKILL):
        for row in rows:
            fd = None
            try:
                fd = os.pidfd_open(row["pid"])
                proc = Path("/proc") / str(row["pid"])
                exe = os.readlink(proc / "exe")
                stat = (proc / "stat").read_text()
                assert exe == row["exe"] and exe.startswith(str(TARGET) + "/") and born(stat) == born(row["stat"])
                home = next((value.split(b"=", 1)[1].decode() for value in (proc / "environ").read_bytes().split(b"\0") if value.startswith(b"SPT_HOME=")), None)
                assert Path(exe).name in ("spt", "mock-session"), exe
                command = (proc / "cmdline").read_bytes().replace(b"\0", b" ").decode(errors="replace")
                assert command == row["cmd"], (command, row["cmd"])
                assert home is not None and any(Path(home).resolve().is_relative_to(root) for root in OWNED_TEMP_ROOTS), home
                signal.pidfd_send_signal(fd, sig)
                actions.append({"pid": row["pid"], "exe": exe, "birth": born(stat), "SPT_HOME": home, "signal": sig.name, "utc": utc()})
            except (FileNotFoundError, ProcessLookupError):
                pass
            finally:
                if fd is not None:
                    os.close(fd)
        if rows:
            time.sleep(3 if sig == signal.SIGTERM else 1)
    remaining = census()["owned_survivors"]
    save(label + "-owned-reap.json", {"scope": str(TARGET), "actions": actions, "remaining": remaining})
    assert not remaining, remaining
    print("REAP " + label + " " + json.dumps({"signalled": len(actions), "remaining": len(remaining)}), flush=True)


phases = []
claimed = False
try:
    capture_environment("inherited-before-scrub", INHERITED_ENV, require_scrubbed=False)
    capture_environment("scrubbed-driver-process", os.environ)
    source_guard()
    preflight = census()
    save("pre-reap-census.json", preflight)
    if preflight["active"]:
        raise RuntimeError("ACTIVE_PRODUCER_REFUSED " + json.dumps(preflight))
    reap_owned("preflight")
    admit("start")
    workflow = TREE / ".github/workflows/golden.yml"
    lines = [line for line in workflow.read_text().splitlines() if re.match(r"^\s+HEAVY:", line)]
    assert len(lines) == 1, "HEAVY must have exactly one hoisted definition"
    match = re.fullmatch(r"      HEAVY: '(.*)'", lines[0])
    assert match is not None and "''" not in match[1], "unexpected HEAVY scalar shape"
    heavy = match[1]
    ENV["HEAVY"] = heavy
    (PROOF / "heavy-source-line.txt").write_text(lines[0] + "\n")
    (PROOF / "heavy.txt").write_text(heavy + "\n")
    runner_temp = Path(tempfile.mkdtemp(prefix="spt-consumer-49a08a07-", dir="/tmp"))
    assert runner_temp.is_dir()
    outside_git(runner_temp)
    github_env = PROOF / "temp.env"
    assert not github_env.exists()
    ENV.update(RUNNER_TEMP=str(runner_temp), GITHUB_ENV=str(github_env), GITHUB_RUN_ID="49a08a07", GITHUB_RUN_ATTEMPT="1")
    setup = run("temp-setup", ["bash", ".github/ci/test-temp-sandbox.sh", "setup"])
    assert setup["exit"] == 0
    for line in github_env.read_text().splitlines():
        key, value = line.split("=", 1)
        ENV[key] = value
    for key in ("TEMP", "TMP", "TMPDIR", "SPT_CI_TEST_TMP"):
        outside_git(Path(ENV[key]))
    OWNED_TEMP_ROOTS.append(Path(ENV["SPT_CI_TEST_TMP"]).resolve())
    save("temp-boundary.json", {key: ENV[key] for key in ("TEMP", "TMP", "TMPDIR", "SPT_CI_TEST_TMP")})
    ENV.pop("SPT_CI_TEST_TMP")
    claim = run("pool-claim", [str(ROOT / ".worktrees/consumer-linux-304-b8482445/target/debug/xtask"), "pool-claim", "--pool", str(TARGET), "--label", "todlando-linux-49a08a07"])
    assert claim["exit"] == 0
    claimed = True
    for name, expression in (("phase-a", "not ( " + heavy + " )"), ("phase-b", heavy)):
        admit(name)
        receipt = run(name, ["cargo", "nextest", "run", "--workspace", "--no-fail-fast", "-E", expression])
        receipt["source_after"] = source_guard()
        phases.append(receipt)
        save("phases.json", phases)
        reap_owned(name)
    final = {"sha": SHA, "start_utc": START, "end_utc": utc(), "phases": phases, "final": census()}
    save("receipt.json", final)
    print("CONSUMER_DONE " + json.dumps({"sha": SHA, "exits": [r["exit"] for r in phases]}), flush=True)
except BaseException as error:
    save("driver-error.json", {"sha": SHA, "start_utc": START, "end_utc": utc(), "error": repr(error), "phases": phases, "census": census()})
    raise
finally:
    capture_environment("final")
    if claimed:
        run("pool-release", [str(ROOT / ".worktrees/consumer-linux-304-b8482445/target/debug/xtask"), "pool-release", "--pool", str(TARGET)])
    save("final-census.json", census())
sys.exit(0 if len(phases) == 2 and all(row["exit"] == 0 for row in phases) else 1)