#!/usr/bin/env python3
"""Run an allocated local lane; never infer allocation from available disk.

A JSON plan contains reserve_bytes, max_growth_bytes, rustc_version (the exact
`rustc -Vv` stdout recorded when allocating a WARM pool), and steps. A step has
name and timeout_seconds, plus either argv or nextest={packages, filter, profile}.
Nextest inventory and execution use the same constructed selection. No implicit
workspace battery, retries, pool disposal, or budget widening.

Requires the gater's --hold-released acknowledgement. Reuses ws272-w0's quiet
meter, lane lock, identity scrub and disk-floor meter. Volume movement is aggregate
capacity evidence, NOT attribution to this lane. Census has the existing reaper's
explicit scope; an unavailable/incomplete census is never clearance.
"""
import argparse
import importlib.util
import json
import math
import os
from pathlib import Path
import re
import shutil
import subprocess
import sys
import time

import psutil

_spec = importlib.util.spec_from_file_location("ws272_w0", Path(__file__).with_name("ws272-w0.py"))
w0 = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(w0)
REFUSED = 75


def save(path, value):
    path.write_text(json.dumps(value, indent=2) + "\n", encoding="utf-8")


# [impl->REQ-LOCAL-GATE-EXECUTION-EVIDENCE]
def capacity_verdict(start_free, free, reserve, growth, initial=False):
    if initial:
        return free >= reserve + growth
    return free > reserve and start_free - free < growth


# [impl->REQ-LOCAL-GATE-EXECUTION-EVIDENCE]
def selected_commands(selection):
    if not isinstance(selection, dict) or not isinstance(selection.get("packages"), list):
        raise ValueError("nextest selection requires packages")
    packages = selection["packages"]
    if not packages or any(not isinstance(p, str) or not p for p in packages):
        raise ValueError("nextest packages must be nonempty strings; no implicit workspace")
    common = [arg for package in packages for arg in ("-p", package)]
    for key, option in (("filter", "-E"), ("profile", "--profile")):
        if key in selection:
            if not isinstance(selection[key], str) or not selection[key]:
                raise ValueError(f"nextest {key} must be a nonempty string")
            common += [option, selection[key]]
    return (["cargo", "nextest", "list", *common, "--message-format", "json"],
            ["cargo", "nextest", "run", *common, "--no-fail-fast", "--success-output", "immediate"])

# [impl->REQ-LOCAL-GATE-EXECUTION-EVIDENCE]
def selected_inventory(inventory):
    """Interpret nextest-metadata's TestListSummary, retaining only matching tests."""
    suites = inventory.get("rust-suites") if isinstance(inventory, dict) else None
    if not isinstance(suites, dict):
        raise ValueError("nextest inventory has no suite map")
    selected, count = [], 0
    for binary_id, suite in suites.items():
        if not isinstance(suite, dict) or not isinstance(suite.get("testcases"), dict):
            raise ValueError(f"{binary_id}: missing test cases")
        if suite.get("status", "listed") not in ("listed", "skipped", "skipped-default-filter"):
            raise ValueError(f"{binary_id}: unknown listing status")
        matches = []
        for name, case in suite["testcases"].items():
            count += 1
            match = case.get("filter-match") if isinstance(case, dict) else None
            if not isinstance(match, dict) or match.get("status") not in ("matches", "mismatch"):
                raise ValueError(f"{binary_id}/{name}: missing filter result")
            if match["status"] == "matches":
                if suite.get("status", "listed") != "listed":
                    raise ValueError(f"{binary_id}: skipped suite contains matching tests")
                matches.append(name)
        if matches:
            path = suite.get("binary-path")
            if not isinstance(path, str) or not path:
                raise ValueError(f"{binary_id}: missing executing binary path")
            selected.append({"binary_id": binary_id, "binary_path": path, "tests": matches})
    if type(inventory.get("test-count")) is not int or inventory["test-count"] != count:
        raise ValueError("nextest inventory test count is incomplete")
    if not selected:
        raise ValueError("nextest selection matches no tests")
    return selected


# [impl->REQ-LOCAL-GATE-EXECUTION-EVIDENCE]
def observe_after(root, output, name, env, meter):
    """An unavailable disk observation must not prevent the independent reaper."""
    disk_ok = builders_ok = clear = False
    try:
        disk_ok = meter(name + "-after-producer")
    except OSError as error:
        print(f"DISK_OBSERVATION_UNPROVEN: {error}", file=sys.stderr)
    try:
        save(output / f"{name}-builders.json", builder_snapshot())
        builders_ok = True
    except (OSError, psutil.Error) as error:
        print(f"BUILDER_OBSERVATION_UNPROVEN: {error}", file=sys.stderr)
    try:
        clear = census(root, output, name + "-after-census", env, reap=True)
    except OSError as error:
        print(f"CENSUS_UNPROVEN: {error}", file=sys.stderr)
    return disk_ok and builders_ok and clear



# [impl->REQ-LOCAL-GATE-EXECUTION-EVIDENCE]
def census_clear(raw, code):
    rows = re.findall(r"^CI-CENSUS phase=\S+ .*?scoped=(\d+) .*?unreadable_path=(\d+)(?:\s|$)", raw, re.M)
    return code == 0 and bool(rows) and rows[-1] == ("0", "0")


def census(root, output, name, env, reap=False):
    if sys.platform == "win32":
        argv = ["pwsh", "-NoProfile", "-File", str(root / ".github/ci/reap-census.ps1"),
                "-Phase", "end" if reap else "start", "-Strict"]
    else:
        argv = ["bash", str(root / ".github/ci/reap-census.sh"), "end" if reap else "start"]
    try:
        result = subprocess.run(argv, cwd=root, env=env, capture_output=True, text=True, timeout=60)
        raw, code = result.stdout + result.stderr, result.returncode
    except (OSError, subprocess.TimeoutExpired) as error:
        raw, code = f"CENSUS_UNPROVEN: {error}\n", REFUSED
    (output / f"{name}.raw").write_text(raw, encoding="utf-8")
    (output / f"{name}.exit").write_text(f"{code}\n", encoding="utf-8")
    clear = census_clear(raw, code)
    print(f"LANE_CENSUS name={name} clear={str(clear).lower()} scope=existing-reaper", flush=True)
    return clear


def builder_snapshot():
    names = {"cargo", "cargo-nextest", "rustc", "rust-lld", "lld-link", "link"}
    rows, unreadable = [], []
    for proc in psutil.process_iter(["pid", "name"]):
        if (proc.info["name"] or "").lower().removesuffix(".exe") not in names:
            continue
        try:
            cpu = proc.cpu_times()
            rows.append({"pid": proc.pid, "birth": proc.create_time(), "image": proc.exe(),
                         "name": proc.info["name"], "cpu_seconds": cpu.user + cpu.system})
        except psutil.NoSuchProcess:
            continue
        except (psutil.AccessDenied, OSError) as error:
            unreadable.append({"pid": proc.pid, "error": str(error)})
    return {"rows": rows, "unreadable": unreadable, "scope": "sampled-build-images-not-load-attribution"}


def sample(root, output, name, start_free, reserve, growth, initial=False):
    try:
        free = shutil.disk_usage(root).free
        row = {"monotonic_ns": time.monotonic_ns(), "time_ns": time.time_ns(), "name": name,
               "free_bytes": free, "reserve_bytes": reserve, "max_growth_bytes": growth,
               "aggregate_consumed_bytes": start_free - free, "sample": "INSTANT"}
        admitted = capacity_verdict(start_free, free, reserve, growth, initial)
    except OSError as error:
        row, admitted = {"name": name, "error": str(error), "sample": "UNAVAILABLE"}, False
    with (output / "capacity.jsonl").open("a", encoding="utf-8") as stream:
        stream.write(json.dumps(row) + "\n")
    return admitted


def stop_observed(processes):
    """psutil Process objects retain birth identity; never kill a remembered PID alone."""
    errors = []
    for process in processes:
        try:
            if process.is_running():
                process.kill()
        except psutil.NoSuchProcess:
            continue
        except (psutil.AccessDenied, OSError) as error:
            errors.append({"pid": process.pid, "error": str(error)})
    try:
        _, alive = psutil.wait_procs(processes, timeout=10)
    except (psutil.Error, OSError) as error:
        # A failed cleanup observation must not hide the producer's real exit.
        errors.append({"operation": "wait_procs", "error": str(error)})
        alive = processes
    survivors = []
    for process in alive:
        try:
            if process.is_running():
                survivors.append({"pid": process.pid, "birth": process.create_time()})
        except psutil.NoSuchProcess:
            continue
        except (psutil.AccessDenied, OSError) as error:
            errors.append({"pid": process.pid, "error": str(error)})
    return errors, survivors


# [impl->REQ-LOCAL-GATE-EXECUTION-EVIDENCE]
def run_producer(root, output, name, argv, timeout, env, meter, json_stdout=False):
    save(output / f"{name}.command.json", argv)
    observed, reason, cleanup_errors, survivors = {}, None, [], []
    started = time.monotonic()
    code, child, stdout = None, None, None
    launch_failed = False
    with (output / f"{name}.raw").open("wb") as raw:
        try:
            # Inventory JSON must not be mixed with Cargo's stderr diagnostics.
            if json_stdout:
                stdout = (output / f"{name}.stdout").open("wb")
            child = subprocess.Popen(argv, cwd=root, env=env,
                                     stdout=stdout if stdout else raw,
                                     stderr=raw if stdout else subprocess.STDOUT)
            try:
                root_process = psutil.Process(child.pid)
                observed[(root_process.pid, root_process.create_time())] = root_process
            except psutil.NoSuchProcess:
                root_process = None
            while child.poll() is None:
                if root_process is not None and root_process.is_running():
                    for proc in root_process.children(recursive=True):
                        try:
                            observed[(proc.pid, proc.create_time())] = proc
                        except psutil.NoSuchProcess:
                            continue
                if not meter():
                    reason = "capacity-or-meter-refusal"
                elif time.monotonic() - started >= timeout:
                    reason = "driver-timeout"
                if reason:
                    break
                time.sleep(0.5)
        except (OSError, psutil.Error, KeyboardInterrupt) as error:
            reason = f"producer-or-observation-error: {error}"
            if child is None:
                launch_failed = True
            print(reason, file=sys.stderr)
        finally:
            if child is not None:
                if child.poll() is None:
                    try:
                        child.kill()  # retained Popen handle, not a bare PID lookup
                    except OSError as error:
                        cleanup_errors.append({"pid": child.pid, "error": str(error)})
                errors, survivors = stop_observed(list(observed.values()))
                cleanup_errors.extend(errors)
                try:
                    code = child.wait(timeout=10)
                except subprocess.TimeoutExpired:
                    reason = "producer-exit-unobserved"
            if stdout is not None:
                stdout.close()
    # A failed artifact write must not turn a real producer failure into another code.
    try:
        (output / f"{name}.exit").write_text(
            f"{code}\n" if code is not None else "UNOBSERVED\n", encoding="utf-8")
        save(output / f"{name}.observation.json", {
            "producer_exit": code, "launch_failed": launch_failed, "stop_reason": reason,
            "elapsed_seconds": time.monotonic() - started,
            "observed_identities": [{"pid": pid, "birth": birth} for pid, birth in observed],
            "cleanup_errors": cleanup_errors, "survivors": survivors,
            "scope": "sampled-descendants; post-leg reaper census is separate"})
    except OSError as error:
        reason = f"artifact-write-failed: {error}"
        print(reason, file=sys.stderr)
    status = code if code is not None else (127 if launch_failed else REFUSED)
    return status, reason is None and not cleanup_errors and not survivors


def validate_plan(plan):
    if not isinstance(plan, dict):
        raise ValueError("plan must be an object")
    for key in ("reserve_bytes", "max_growth_bytes"):
        if type(plan.get(key)) is not int or plan[key] <= 0:
            raise ValueError(f"{key} must be positive allocated bytes")
    if plan["reserve_bytes"] < w0.FLOOR_BYTES:
        raise ValueError("allocation cannot weaken the existing disk floor")
    if not isinstance(plan.get("rustc_version"), str) or not plan["rustc_version"].strip():
        raise ValueError("allocation requires the warm pool's exact rustc -Vv stdout")
    if not isinstance(plan.get("steps"), list) or not plan["steps"]:
        raise ValueError("plan requires nonempty steps")
    names = {"pool-claim", "pool-release", "before", "start", "lane", "plan",
             "toolchain", "environment-scrub", "builders-before", "capacity"}
    for step in plan["steps"]:
        if not isinstance(step, dict):
            raise ValueError("each step must be an object")
        name = step.get("name")
        if (not isinstance(name, str) or not re.fullmatch(r"[A-Za-z0-9_-]+", name)
                or name in names or name.endswith(("-inventory", "-builders", "-census"))):
            raise ValueError("step names must be unique safe filenames, excluding driver artifacts")
        names.add(name)
        if (type(step.get("timeout_seconds")) not in (int, float)
                or not math.isfinite(step["timeout_seconds"]) or step["timeout_seconds"] <= 0):
            raise ValueError(f"{name}: explicit finite positive driver timeout required")
        if ("argv" in step) == ("nextest" in step):
            raise ValueError(f"{name}: supply argv OR nextest")
        if "nextest" in step:
            selected_commands(step["nextest"])
        elif not isinstance(step["argv"], list) or not step["argv"] or any(not isinstance(a, str) for a in step["argv"]):
            raise ValueError(f"{name}: argv must be a nonempty string array")
        elif Path(step["argv"][0]).name.lower() in ("cargo", "cargo.exe"):
            if any(arg in ("test", "nextest") for arg in step["argv"][1:]):
                raise ValueError(f"{name}: Cargo tests require a nextest selection and inventory")


# [impl->REQ-LOCAL-GATE-EXECUTION-EVIDENCE]
def final_status(producer, observations_ok):
    return producer if producer else (0 if observations_ok else REFUSED)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--plan", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--xtask", type=Path, required=True, help="already-built meter and pool verbs; never cargo run")
    parser.add_argument("--hold-released", action="store_true")
    args = parser.parse_args()
    if not args.hold_released:
        parser.error("requires the gater's explicit --hold-released")
    root = Path(__file__).resolve().parents[2]
    plan = json.loads(args.plan.read_text(encoding="utf-8"))
    validate_plan(plan)
    target, xtask, output = root / "target", args.xtask.resolve(), args.output.resolve()
    if not xtask.is_file() or not (target / "debug/deps").is_dir() or not (target / ".rustc_info.json").is_file():
        raise ValueError("warm pool and prebuilt xtask required; cold cache invalidates allocation")
    if target.is_symlink() or getattr(target.lstat(), "st_file_attributes", 0) & 0x400:
        raise ValueError("target reparse-point ownership requires a separately approved pool plan")
    if not any(path.suffix in (".rlib", ".rmeta") and path.is_file()
               for path in (target / "debug/deps").iterdir()):
        raise ValueError("warm dependency artifacts unavailable; allocation invalid")
    output.mkdir(parents=True, exist_ok=False)
    temp = output / "temp"
    temp.mkdir()
    if len({os.stat(p).st_dev for p in (root, target, temp)}) != 1:
        raise ValueError("build/temp volumes differ; obtain a multi-volume allocation")
    save(output / "plan.json", plan)
    env, dropped = w0.battery_env()
    env.update({"GITHUB_WORKSPACE": str(root), "CARGO_TARGET_DIR": str(target),
                "RUNNER_TEMP": str(temp), "TMP": str(temp), "TEMP": str(temp), "TMPDIR": str(temp),
                "SPT_CI_REAP_STRICT": "1"})
    env.pop("SPT_HOME", None)
    save(output / "environment-scrub.json", {"removed": dropped, "fixture_temp": str(temp)})
    reserve, growth = plan["reserve_bytes"], plan["max_growth_bytes"]
    start_free = shutil.disk_usage(root).free
    meter = lambda name: sample(root, output, name, start_free, reserve, growth)
    if not sample(root, output, "start", start_free, reserve, growth, initial=True):
        return REFUSED
    actual = subprocess.run(["rustc", "-Vv"], cwd=root, env=env, capture_output=True, text=True, timeout=30)
    save(output / "toolchain.json", {"stdout": actual.stdout, "stderr": actual.stderr, "exit": actual.returncode})
    cached = json.loads((target / ".rustc_info.json").read_text(encoding="utf-8"))
    cached_outputs = cached.get("outputs", {}).values()
    if actual.returncode or actual.stdout != plan["rustc_version"] or not any(row.get("success") is True and row.get("stdout") == actual.stdout for row in cached_outputs):
        raise ValueError("toolchain differs from allocated warm cache; allocation invalid")
    lock = w0.claim_pool(root)
    if lock is None:
        return REFUSED
    claimed = False
    result = 0
    try:
        if not w0.preflight() or not census(root, output, "before", env):
            raise ValueError("quiet-window or scoped census refused")
        save(output / "builders-before.json", builder_snapshot())
        if not w0.disk_floor(root, output, "start", xtask):
            raise ValueError("disk floor refused")
        result, ok = run_producer(root, output, "pool-claim", [str(xtask), "pool-claim", "--pool", str(target), "--label", "local-lane"], 30, env, lambda: meter("claim"))
        claimed = result == 0
        result = final_status(result, ok)
        if not claimed:
            # A terminated claim writer may have persisted ownership before exit.
            # Do not release an unproven claim: it may still belong to another lane.
            print("POOL_CUSTODY_UNPROVEN: claim did not complete successfully; "
                  "inspect POOL-OWNER.json before any takeover or deletion", file=sys.stderr)
        for step in plan["steps"]:
            if result:
                break
            name = step["name"]
            if not w0.preflight() or not meter(f"{name}-before") or not census(root, output, f"{name}-before-census", env):
                result = REFUSED
                break
            if "nextest" in step:
                listing, argv = selected_commands(step["nextest"])
                result, ok = run_producer(root, output, name + "-inventory", listing, step["timeout_seconds"], env, lambda: meter(name + "-inventory"), json_stdout=True)
                after_ok = observe_after(root, output, name + "-inventory", env, meter)
                ok = ok and after_ok
                if result or not ok:
                    result = final_status(result, ok)
                    break
                inventory = json.loads((output / f"{name}-inventory.stdout").read_text(encoding="utf-8"))
                save(output / f"{name}-selection.json", selected_inventory(inventory))
            else:
                argv = step["argv"]
            result, ok = run_producer(root, output, name, argv, step["timeout_seconds"], env, lambda: meter(name))
            after_ok = observe_after(root, output, name, env, meter)
            result = final_status(result, ok and after_ok)
            if result:
                break
    except (OSError, ValueError, psutil.Error, subprocess.TimeoutExpired) as error:
        print(f"LANE_OBSERVATION_FAILED: {error}", file=sys.stderr)
        result = final_status(result, False)
    finally:
        # Release through the prebuilt executable; this driver never deletes target.
        if claimed:
            try:
                released = subprocess.run([str(xtask), "pool-release", "--pool", str(target)], cwd=root, env=env, capture_output=True, text=True, timeout=30)
                save(output / "pool-release.json", {"exit": released.returncode, "stdout": released.stdout, "stderr": released.stderr})
                result = final_status(result, released.returncode == 0)
            except (OSError, subprocess.TimeoutExpired) as error:
                print(f"POOL_RELEASE_UNPROVEN: {error}", file=sys.stderr)
                result = final_status(result, False)
        lock.close()
    try:
        (output / "lane.exit").write_text(f"{result}\n", encoding="utf-8")
    except OSError as error:
        print(f"LANE_STATUS_WRITE_FAILED: {error}; producer/driver exit={result}", file=sys.stderr)
        result = final_status(result, False)
    return result if result >= 0 else 128 - result


if __name__ == "__main__":
    try:
        sys.exit(main())
    except (OSError, ValueError, subprocess.TimeoutExpired) as error:
        print(f"LANE_REFUSED: {error}", file=sys.stderr)
        sys.exit(REFUSED)