"""Measure private-daemon agent density and first-cell latency, without a real model.

Run with nix develop -c python3 test/manual/agent_density.py --agents 50
--output /tmp/density/baseline-50.json. ALBEDO_TEST_DAEMON may name a snapshot.
The default keeps kernels resident; --idle-seconds and --settle test reaping.
"""

from __future__ import annotations

import collections
import os
import argparse
from concurrent.futures import ThreadPoolExecutor
import ctypes
import json
import math
from pathlib import Path
import platform
import sqlite3
import struct
import subprocess
import sys
import threading
import time

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "test/e2e"))
import harness  # noqa: E402
from harness import Albedo, Daemon, Provider, python, text  # noqa: E402
from inspect_support import Inspection  # noqa: E402

DEFINE = """import fractions
answer = 7
rows = [1, 2, 3]
def helper(x):
    return fractions.Fraction(x, answer)
print('DENSITY_READY')
"""
RECALL = "assert helper(21) == 3 and rows == [1, 2, 3]; print('DENSITY_OK')"


def process_bytes(pid):
    """Physical footprint on Darwin, proportional resident bytes on Linux."""
    if sys.platform == "darwin":
        buffer = ctypes.create_string_buffer(96)
        libproc = ctypes.CDLL("/usr/lib/libproc.dylib")
        if libproc.proc_pid_rusage(pid, 0, buffer) != 0:
            raise OSError(f"cannot read footprint of pid {pid}")
        return struct.unpack_from("=Q", buffer.raw, 72)[0]
    for line in Path(f"/proc/{pid}/smaps_rollup").read_text().splitlines():
        if line.startswith("Pss:"):
            return int(line.split()[1]) * 1024
    raise RuntimeError(f"no PSS for pid {pid}")


def processes(app):
    """The daemon and its Python processes: kernels, spares, bridges, the
    zygote, and the children it forked (they keep the zygote's command line;
    PSS apportions the pages they share with it)."""
    rows = [
        line.strip().split(None, 2)
        for line in subprocess.check_output(
            ["ps", "-axo", "pid=,ppid=,command="], text=True
        ).splitlines()
    ]
    rows = [(int(pid), int(ppid), command) for pid, ppid, command in rows]
    daemon = app.connection["pid"]
    children: dict[int, list[int]] = {}
    for pid, ppid, _ in rows:
        children.setdefault(ppid, []).append(pid)
    descendants, frontier = set(), [daemon]
    while frontier:
        for child in children.get(frontier.pop(), []):
            if child not in descendants:
                descendants.add(child)
                frontier.append(child)
    commands = {pid: command for pid, _, command in rows}
    zygotes = {
        pid
        for pid, ppid, command in rows
        if "albedo_zygote.py" in command
        and pid in descendants
        and "albedo_zygote.py" not in commands.get(ppid, "")
    }
    result = [{"pid": daemon, "kind": "daemon"}]
    for pid, ppid, command in rows:
        ours = str(app.home) in command or pid in descendants
        if not ours:
            continue
        if pid in zygotes:
            kind = "zygote"
        elif ppid in zygotes:
            kind = "forked"
        elif "albedo_spare.py" in command:
            kind = "spare"
        elif "albedo_kernel.py" in command:
            kind = "kernel"
        elif "albedo_bridge.py" in command:
            kind = "bridge"
        else:
            continue
        result.append({"pid": pid, "kind": kind})
    return result


_TIMEBASE = None


def process_cpu(pid):
    """User plus system CPU seconds a live process has used, or None if gone."""
    global _TIMEBASE
    if sys.platform == "darwin":
        buffer = ctypes.create_string_buffer(96)
        libproc = ctypes.CDLL("/usr/lib/libproc.dylib")
        if libproc.proc_pid_rusage(pid, 0, buffer) != 0:
            return None
        user, system = struct.unpack_from("=QQ", buffer.raw, 16)
        if _TIMEBASE is None:
            # rusage times are mach absolute ticks, not nanoseconds on arm64.
            info = (ctypes.c_uint32 * 2)()
            ctypes.CDLL("/usr/lib/libSystem.dylib").mach_timebase_info(info)
            _TIMEBASE = info[0] / info[1]
        return (user + system) * _TIMEBASE / 1e9
    try:
        fields = Path(f"/proc/{pid}/stat").read_text().rsplit(")", 1)[1].split()
    except OSError:
        return None
    return (int(fields[11]) + int(fields[12])) / os.sysconf("SC_CLK_TCK")


def cpu_snapshot(app):
    """CPU seconds of the daemon and its Python processes, by pid and kind."""
    snapshot = {}
    for item in processes(app):
        seconds = process_cpu(item["pid"])
        if seconds is not None:
            snapshot[item["pid"]] = (item["kind"], seconds)
    return snapshot


def cpu_spent(before, after):
    """CPU used between two snapshots, by kind; a process that started in
    between counts from zero, one that ended in between is lost."""
    spent = {}
    for pid, (kind, seconds) in after.items():
        start = before[pid][1] if pid in before else 0.0
        spent[kind] = spent.get(kind, 0.0) + max(0.0, seconds - start)
    return spent


def measure(app, label, count, baseline, window=0.0):
    items = processes(app)
    for item in items:
        item["bytes"] = process_bytes(item["pid"])
    total = sum(item["bytes"] for item in items)
    memory = Inspection(app).evaluate(
        "io:put_chars(json:encode(maps:from_list(Call(erlang, memory, []))))"
    )
    result = {
        "state": label,
        "agents": count,
        "total_bytes": total,
        "marginal_bytes_per_agent": (total - baseline) / count,
        "kinds": dict(collections.Counter(item["kind"] for item in items)),
        "processes": items,
        "beam": json.loads(memory),
    }
    if window > 0:
        before = cpu_snapshot(app)
        time.sleep(window)
        spent = cpu_spent(before, cpu_snapshot(app))
        total = sum(spent.values())
        result["idle_cpu"] = {
            "window_s": window,
            "cpu_s": total,
            "by_kind": spent,
            "per_agent_ms_per_s": total / count / window * 1000,
        }
    print(
        json.dumps(
            {
                key: value
                for key, value in result.items()
                if key not in ("processes", "beam", "idle_cpu")
            }
            | {
                "idle_cpu_ms_per_agent_s": result.get("idle_cpu", {}).get(
                    "per_agent_ms_per_s"
                )
            }
        ),
        flush=True,
    )
    return result


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--agents", type=int, default=1)
    parser.add_argument("--samples", type=int, default=10)
    parser.add_argument("--idle-seconds", type=int, default=3600)
    parser.add_argument("--settle", type=float, default=0.2)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument(
        "--idle-window",
        type=float,
        default=0.0,
        help="seconds of CPU sampling after each idle state's memory reading",
    )
    parser.add_argument(
        "--burst",
        type=int,
        default=0,
        help="recall in this many agents at once after the sequential samples",
    )
    parser.add_argument(
        "--daemon-env",
        action="append",
        default=[],
        metavar="NAME=VALUE",
        help="environment variable for the private daemon (repeatable)",
    )
    args = parser.parse_args()
    if args.agents < 1 or args.samples < 1:
        parser.error("agents and samples must be positive")
    harness.snapshot_daemon()
    daemon = Daemon("density")
    release_model = threading.Event()
    condition = threading.Condition()
    finished = {}
    held = set()
    failures = []
    recipient = ""
    gate = daemon.root / "cell-gate"

    def script(request):
        # Notes (a restored namespace, hook context) arrive as their own user
        # items; neither is the prompt nor starts a fresh turn.
        users = [
            item.get("content", "")
            for item in request["input"]
            if item.get("role") == "user"
            and not str(item.get("content", "")).startswith("<system-note>")
        ]
        prompt = (users[-1] if users else "").split("<system-note>")[0]
        last = request["input"][-1]
        if last.get("role") == "user" and not str(last.get("content", "")).startswith(
            "<system-note>"
        ):
            if prompt.startswith("define:"):
                return python(DEFINE)
            if prompt.startswith("recall:"):
                return python(RECALL)
            if prompt.startswith("mail:"):
                return python(f"await mail.request({recipient!r}, 'density waiting')")
            if prompt.startswith("model:"):
                with condition:
                    held.add(prompt)
                    condition.notify_all()
                if not release_model.wait(300):
                    failures.append("model gate timed out")
                return text("released")
            if prompt.startswith("cell:"):
                return python(
                    "while not Path(" + repr(str(gate)) + ").exists():\n"
                    "    await asyncio.sleep(0.05)\n" + RECALL,
                    timeout_ms=300000,
                )
        output = "".join(item.get("output", "") for item in request["input"])
        if prompt.startswith("recall:") and "DENSITY_OK" not in output:
            failures.append(f"namespace round-trip failed: {prompt}: {output}")
        with condition:
            finished[prompt] = time.perf_counter()
            condition.notify_all()
        return text("ok")

    def prepare(app):
        app.env.update(
            ALBEDO_INSPECT="1",
            ALBEDO_IDLE_SECONDS=str(args.idle_seconds),
            ALBEDO_CELL_BACKGROUND_SECONDS="0.2",
            **dict(item.split("=", 1) for item in args.daemon_env),
        )

    provider = Provider(script, chunk_size=100000)
    result = {
        "platform": platform.platform(),
        "python": sys.version,
        # A synced copy without .git (a remote run) names its commit itself.
        "commit": os.environ.get("ALBEDO_DENSITY_COMMIT")
        or subprocess.check_output(
            ["git", "rev-parse", "HEAD"], cwd=ROOT, text=True
        ).strip(),
        "metric": "phys_footprint" if sys.platform == "darwin" else "pss",
        "options": {**vars(args), "output": str(args.output)},
        "states": [],
    }
    try:
        with (
            harness.on_daemon(daemon),
            Albedo(provider, protocol="responses", prepare=prepare) as app,
        ):
            recipient = app.session()
            turn_prompt = "cell:recipient"
            app.prompt(recipient, turn_prompt).close()
            app.idle(recipient)
            baseline = sum(process_bytes(item["pid"]) for item in processes(app))
            result["empty_daemon_bytes"] = baseline
            if args.idle_window > 0:
                before = cpu_snapshot(app)
                time.sleep(args.idle_window)
                spent = cpu_spent(before, cpu_snapshot(app))
                result["empty_idle_cpu"] = {
                    "window_s": args.idle_window,
                    "cpu_s": sum(spent.values()),
                    "by_kind": spent,
                }
            sessions = [app.session() for _ in range(args.agents)]
            result["states"].append(
                measure(app, "created", args.agents, baseline, args.idle_window)
            )

            def turn(pair):
                session, prompt = pair
                app.prompt(session, prompt).close()
                app.idle(session, timeout=180)

            def turns(state):
                with ThreadPoolExecutor(max_workers=min(8, args.agents)) as pool:
                    list(
                        pool.map(
                            turn,
                            [
                                (session, f"{state}:{index}")
                                for index, session in enumerate(sessions)
                            ],
                        )
                    )

            phase_cpu = result["phase_cpu"] = {}

            def charge(name, operations, before):
                spent = cpu_spent(before, cpu_snapshot(app))
                total = sum(spent.values())
                phase_cpu[name] = {
                    "operations": operations,
                    "cpu_s": total,
                    "by_kind": spent,
                    "per_operation_ms": total / operations * 1000,
                }

            before = cpu_snapshot(app)
            turns("define")
            time.sleep(args.settle)
            # A cell in a new session, then whatever release the settle allows.
            charge("first_turn_and_release", args.agents, before)
            result["states"].append(
                measure(app, "resting", args.agents, baseline, args.idle_window)
            )
            before = cpu_snapshot(app)
            latencies = []
            for index in range(args.samples):
                prompt = f"recall:{index}"
                started = time.perf_counter()
                app.prompt(sessions[index % len(sessions)], prompt).close()
                with condition:
                    if not condition.wait_for(lambda: prompt in finished, timeout=60):
                        raise AssertionError("first cell result timed out")
                    latencies.append((finished[prompt] - started) * 1000)
                app.idle(sessions[index % len(sessions)])
                time.sleep(args.settle)
            # Each sample restores a namespace if it was released, runs a cell,
            # and settles, which may release it again.
            charge("recall_turn", args.samples, before)
            ordered = sorted(latencies)
            result["first_cell_ms"] = {
                "samples": latencies,
                "p50": ordered[math.ceil(len(ordered) * 0.5) - 1],
                "p95": ordered[math.ceil(len(ordered) * 0.95) - 1],
            }
            if args.burst:
                before = cpu_snapshot(app)
                prompts = [f"recall:burst{i}" for i in range(args.burst)]
                started = time.perf_counter()
                submitted = {}
                for index, prompt in enumerate(prompts):
                    submitted[prompt] = time.perf_counter()
                    app.prompt(sessions[index % len(sessions)], prompt).close()
                submit_ms = (time.perf_counter() - started) * 1000
                with condition:
                    if not condition.wait_for(
                        lambda: all(p in finished for p in prompts), timeout=180
                    ):
                        raise AssertionError("burst results timed out")
                burst = sorted((finished[p] - started) * 1000 for p in prompts)
                own = sorted((finished[p] - submitted[p]) * 1000 for p in prompts)
                result["burst_ms"] = {
                    "agents": args.burst,
                    "p50": burst[len(burst) // 2],
                    "p95": burst[math.ceil(len(burst) * 0.95) - 1],
                    "max": burst[-1],
                    "submit_ms": submit_ms,
                    "own_p50": own[len(own) // 2],
                    "own_p95": own[math.ceil(len(own) * 0.95) - 1],
                }
                for index in range(args.burst):
                    app.idle(sessions[index % len(sessions)], timeout=180)
                charge("burst", args.burst, before)
                print(json.dumps(result["burst_ms"]), flush=True)
            turns("mail")
            time.sleep(args.settle)
            result["states"].append(
                measure(app, "waiting_mail", args.agents, baseline, args.idle_window)
            )
            for index, session in enumerate(sessions):
                app.prompt(session, f"model:{index}").close()
            with condition:
                if not condition.wait_for(
                    lambda: len(held) == args.agents, timeout=180
                ):
                    raise AssertionError(f"only {len(held)} model waits admitted")
            time.sleep(args.settle)
            result["states"].append(
                measure(app, "waiting_model", args.agents, baseline, args.idle_window)
            )
            release_model.set()
            for session in sessions:
                app.idle(session, timeout=180)
            turns("cell")
            with sqlite3.connect(app.home / "albedo.sqlite") as database:
                result["running_kernel_links"] = database.execute(
                    "SELECT count(*) FROM kernel_links"
                ).fetchone()[0]
            result["states"].append(measure(app, "running_cell", args.agents, baseline))
            gate.touch()
            for session in sessions:
                app.idle(session, timeout=180)
            result["namespace_correct"] = not failures
            if failures:
                raise AssertionError(failures)
            print(json.dumps(result["first_cell_ms"]), flush=True)
    finally:
        release_model.set()
        gate.touch()
        daemon.shutdown()
        deadline = time.monotonic() + 10
        while True:
            leftovers = subprocess.check_output(
                ["ps", "-axo", "pid=,command="], text=True
            )
            result["cleanup_remaining"] = [
                line for line in leftovers.splitlines() if str(daemon.home) in line
            ]
            if not result["cleanup_remaining"] or time.monotonic() >= deadline:
                break
            time.sleep(0.05)
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(json.dumps(result, indent=2) + "\n")
        if result["cleanup_remaining"]:
            raise AssertionError(result["cleanup_remaining"])


if __name__ == "__main__":
    main()