"""Record every perf scenario against bonsai inside a networkless container on valefar."""

from __future__ import annotations

import argparse
import fcntl
import json
from pathlib import Path
import platform
import shutil
import subprocess
import sys
import tempfile
import time
import uuid

from safety import SecretScanner

ROOT = Path(__file__).resolve().parents[2]


def checked(*command, **kwargs):
    return subprocess.check_output(command, text=True, **kwargs).strip()


def build(directory):
    """Build inside nix develop and retain only runtime executables and their closures."""
    snapshot = directory / "daemon"
    subprocess.run([str(ROOT / "test/snapshot-daemon.sh"), str(snapshot)], check=True)
    cli = directory / "albedo"
    subprocess.run(
        ["go", "-C", str(ROOT / "cli"), "build", "-o", str(cli), "./cmd/albedo"],
        check=True,
    )
    commands = (
        "python3",
        "erl",
        "erlc",
        "bash",
        "sh",
        "env",
        "git",
        "sleep",
        "uname",
        "rm",
        "mkdir",
        "chmod",
        "dirname",
        "basename",
        "readlink",
    )
    # The harness reads process state with ps, which the dev shell does not ship;
    # build it from the flake's locked nixpkgs rather than the host's.
    procps = [
        path
        for path in checked(
            "nix",
            "build",
            "--no-link",
            "--print-out-paths",
            "--inputs-from",
            "path:" + str(ROOT),
            "nixpkgs#procps",
        ).splitlines()
        if (Path(path) / "bin/ps").exists()
    ][0]
    bins = {"ps": procps + "/bin/ps"}
    for command in commands:
        executable = shutil.which(command)
        if not executable or not str(Path(executable).resolve()).startswith(
            "/nix/store/"
        ):
            raise RuntimeError(f"{command} is not from the pinned Nix toolchain")
        # coreutils is one multicall binary that dispatches on argv[0], so keep the
        # command's own name inside the store directory it resolves to.
        named = Path(executable).resolve().parent / command
        bins[command] = str(named if named.exists() else Path(executable).resolve())
    stores = sorted({"/".join(path.split("/")[:4]) for path in bins.values()})
    closure = checked("nix-store", "--query", "--requisites", *stores).splitlines()
    (directory / "runtime.json").write_text(
        json.dumps({"bins": bins, "closure": sorted(set(closure))})
    )


def source_snapshot(destination):
    paths = checked("git", "ls-files", "-z", cwd=ROOT).split("\0")
    paths += [
        str(path.relative_to(ROOT))
        for path in (ROOT / "test/perf").rglob("*")
        if path.is_file() and path.suffix in (".py", ".toml", ".erl")
    ]
    for name in sorted(set(paths)):
        source = ROOT / name
        if not name or not source.is_file() or name.startswith("test/perf/recordings/"):
            continue
        target = destination / name
        target.parent.mkdir(parents=True, exist_ok=True)
        shutil.copyfile(source, target)
        target.chmod(source.stat().st_mode & 0o777)
    (destination / "test/perf/recordings").mkdir(parents=True, exist_ok=True)


def scan_and_copy(source, destination):
    scanner = SecretScanner()
    for path in source.rglob("*"):
        if path.is_symlink() or not (path.is_dir() or path.is_file()):
            raise ValueError(f"unexpected output file: {path.name}")
        if path.is_file():
            scanner.check(json.loads(path.read_text()), label=path.name)
    shutil.copytree(source, destination)


def launch(args):
    if platform.system() != "Linux":
        raise RuntimeError(
            "run on valefar: python3 test/perf/record_container.py --out <new-directory>"
        )
    if args.out.exists():
        raise ValueError("output must be a new directory")
    runtime = json.loads(checked("docker", "info", "--format", "{{json .Runtimes}}"))
    selected = "runsc" if "runsc" in runtime else "runc"
    print(f"container runtime: {selected}", flush=True)
    image = checked("docker", "image", "inspect", args.image, "--format", "{{.Id}}")
    with (
        open("/tmp/bonsai-gpu.lock", "a") as lock,
        tempfile.TemporaryDirectory(prefix="perf-record-") as temporary,
    ):
        # Only this user may reach the world-writable output directory inside.
        directory = Path(temporary)
        directory.chmod(0o700)
        source = directory / "source"
        source_snapshot(source)
        subprocess.run(
            [
                "nix",
                "develop",
                "path:" + str(source),
                "-c",
                "python3",
                str(source / "test/perf/record_container.py"),
                "--build",
                str(directory),
            ],
            check=True,
        )
        config = json.loads((directory / "runtime.json").read_text())
        bins = config["bins"]
        output = directory / "output"
        output.mkdir(mode=0o777)
        output.chmod(0o777)
        bridge_directory = directory / "bridge"
        bridge_directory.mkdir(mode=0o755)
        socket_path = bridge_directory / "bonsai.sock"
        bridge = subprocess.Popen(
            [
                sys.executable,
                str(source / "test/perf/record_bridge.py"),
                str(socket_path),
            ]
        )
        name = "albedo-perf-" + uuid.uuid4().hex
        sentinel = directory / "host-only-file"
        sentinel.write_text("host canary, must not be visible in the container")
        try:
            gateway = json.loads(checked("docker", "network", "inspect", "bridge"))[0][
                "IPAM"
            ]["Config"][0]["Gateway"]
            deadline = time.monotonic() + 10
            while not socket_path.exists():
                if bridge.poll() is not None or time.monotonic() > deadline:
                    raise RuntimeError("bonsai bridge did not start")
                time.sleep(0.02)
            command = [
                "docker",
                "run",
                "--name",
                name,
                "--runtime",
                selected,
                "--user",
                "65534:65534",
                "--cap-drop",
                "ALL",
                "--security-opt",
                "no-new-privileges",
                "--read-only",
                "--network",
                "none",
                "--pids-limit",
                str(args.pids),
                "--memory",
                args.memory,
                "--cpus",
                str(args.cpus),
                "--tmpfs",
                "/tmp:rw,nosuid,nodev,size=4g,mode=1777",
                "--tmpfs",
                "/home:rw,nosuid,nodev,size=1g,mode=1777",
                "--workdir",
                str(source),
            ]
            mounts = [
                (source, source, True),
                (directory / "daemon", directory / "daemon", True),
                (directory / "albedo", directory / "albedo", True),
                (bridge_directory, Path("/bridge"), True),
                (output, Path("/recordings"), False),
            ]
            mounts += [(Path(path), Path(path), True) for path in config["closure"]]
            for host, target, readonly in mounts:
                command += [
                    "--mount",
                    f"type=bind,src={host},dst={target}"
                    + (",readonly" if readonly else ""),
                ]
            environment = {
                "PATH": ":".join(
                    sorted({str(Path(path).parent) for path in bins.values()})
                ),
                "HOME": "/home",
                "TMPDIR": "/tmp",
                "LANG": "C.UTF-8",
                "ALBEDO_TEST_TMP": "/tmp/perf",
                "PERF_COMMIT": checked("git", "rev-parse", "HEAD", cwd=ROOT),
                "ALBEDO_NO_BROWSER": "1",
                "ALBEDO_TEST_DAEMON": str(directory / "daemon/albedo-daemon"),
                "ALBEDO_TEST_CLI": str(directory / "albedo"),
                "ALBEDO_STANDBY": "0",
                "PYTHONDONTWRITEBYTECODE": "1",
                "GIT_CONFIG_NOSYSTEM": "1",
                "GIT_CONFIG_GLOBAL": "/dev/null",
                "GIT_CONFIG_COUNT": "1",
                "GIT_CONFIG_KEY_0": "safe.directory",
                "GIT_CONFIG_VALUE_0": str(source),
            }
            clean_entry = [
                bins["env"],
                "-i",
                *[key + "=" + value for key, value in environment.items()],
            ]
            runner = [
                bins["python3"],
                str(source / "test/perf/run.py"),
                "--mode",
                "live",
                "--record",
                "--runs",
                "1",
                "--url",
                "http://127.0.0.1:8899/v1",
                "--out",
                "/recordings/results.json",
                "--recordings-dir",
                "/recordings",
                "--provider",
                "bonsai",
            ]
            for scenario in args.scenario or []:
                runner += ["--scenario", scenario]
            command += [
                "--entrypoint",
                clean_entry[0],
                image,
                *clean_entry[1:],
                bins["python3"],
                str(source / "test/perf/record_bridge.py"),
                "/bridge/bonsai.sock",
                "--",
                *runner,
            ]
            probe = [
                *command[: command.index("--entrypoint")],
                "--entrypoint",
                clean_entry[0],
                image,
                *clean_entry[1:],
                bins["python3"],
                str(source / "test/perf/record_container.py"),
                "--inside",
                str(sentinel),
                gateway,
            ]
            subprocess.run(probe, check=True, timeout=180)
            subprocess.run(
                ["docker", "rm", name], check=True, stdout=subprocess.DEVNULL
            )
            # A queued gpu.sh job stops bonsai, so only the run that talks to it
            # holds the lock; the build and the isolation probe do not.
            fcntl.flock(lock, fcntl.LOCK_EX)
            subprocess.run(command, check=True, timeout=args.timeout)
            fcntl.flock(lock, fcntl.LOCK_UN)
            subprocess.run(
                ["docker", "rm", name], check=True, stdout=subprocess.DEVNULL
            )
            # The harness writes 0600 files as the container user. With no agent
            # left running, one fixed command hands them to this user to scan.
            subprocess.run(
                [
                    *command[: command.index("--entrypoint")],
                    "--entrypoint",
                    bins["python3"],
                    image,
                    "-c",
                    "import os\n"
                    "for root, dirs, names in os.walk('/recordings'):\n"
                    "    for d in dirs: os.chmod(os.path.join(root, d), 0o777)\n"
                    "    for n in names: os.chmod(os.path.join(root, n), 0o666)\n",
                ],
                check=True,
                timeout=60,
            )
        finally:
            subprocess.run(
                ["docker", "rm", "--force", name],
                check=False,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
            fcntl.flock(lock, fcntl.LOCK_UN)
            bridge.terminate()
            bridge.wait(timeout=10)
        # No container process may remain to race the scanner or rewrite scanned files.
        if checked(
            "docker", "ps", "--all", "--quiet", "--filter", "name=^" + name + "$"
        ):
            raise RuntimeError(f"container {name} was not removed")
        scan_and_copy(output, args.out)
        print(f"scanned recordings: {args.out}", flush=True)


def inside(sentinel, gateway):
    """Require a real agent cell to fail network and host-file access before recording."""
    sys.path.insert(0, str(ROOT / "test/e2e"))
    import harness

    code = """# verify this agent cannot leave the recording sandbox
import socket, pathlib, os
checks = {}
for host, port in [('1.1.1.1', 443), ('10.0.0.30', 22), (GATEWAY, 22)]:
    try:
        with socket.create_connection((host, port), timeout=1):
            checks[host] = 'reachable'
    except OSError:
        checks[host] = 'blocked'
assert all(value == 'blocked' for value in checks.values()), checks
try:
    files.read(SENTINEL)
except (FileNotFoundError, PermissionError):
    checks['host_file'] = 'blocked'
else:
    raise AssertionError('host file was readable')
assert not pathlib.Path('/var/run/docker.sock').exists()
assert os.geteuid() != 0
print(json.dumps(checks))
""".replace("GATEWAY", repr(gateway)).replace("SENTINEL", repr(sentinel))

    def script(request):
        if request["messages"][-1].get("role") == "user":
            return harness.python(code)
        return harness.text("isolation checked")

    harness.snapshot_daemon()
    provider = harness.Provider(script)
    daemon = harness.Daemon("perf-isolation")
    try:
        with harness.on_daemon(daemon), harness.Albedo(provider) as app:
            session = app.session()
            app.prompt(session, "check isolation").close()
            app.idle(session, timeout=120)
            results = [
                json.loads(part["value"])
                for entry in app.history(session)["items"]
                if entry["kind"] == "tool_result" and entry["tool"]["name"] == "python"
                for part in entry["content"]
                if part["kind"] == "json" and part["field"] == "result"
            ]
            if not results or results[-1]["status"] != "ok":
                raise AssertionError(results)
            evidence = json.loads(results[-1]["output"])
            SecretScanner().write_json(Path("/recordings/isolation.json"), evidence)
            print("agent-cell isolation:", evidence, flush=True)
    finally:
        daemon.shutdown()
        provider.close()
        harness.shutdown()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--out", type=Path)
    parser.add_argument(
        "--image",
        default="alpine:3.22",
        help="local image, resolved to its immutable ID",
    )
    parser.add_argument("--memory", default="8g")
    parser.add_argument("--cpus", type=int, default=4)
    # fan-out runs 20 sessions; every kernel and BEAM thread counts against this.
    parser.add_argument("--pids", type=int, default=4096)
    parser.add_argument("--timeout", type=int, default=3600)
    parser.add_argument(
        "--scenario",
        action="append",
        help="record only this scenario (repeatable); omit to record every scenario",
    )
    parser.add_argument("--inside", nargs=2, help=argparse.SUPPRESS)
    parser.add_argument("--build", type=Path, help=argparse.SUPPRESS)
    args = parser.parse_args()
    if args.inside:
        inside(*args.inside)
    elif args.build:
        build(args.build)
    elif not args.out:
        parser.error("--out is required")
    else:
        launch(args)


if __name__ == "__main__":
    main()