"""Measure common tasks using real bonsai recordings and private snapshot daemons."""

from __future__ import annotations

import argparse
import os
from pathlib import Path
import platform
import statistics
import subprocess
import sys
import time
from concurrent.futures import ThreadPoolExecutor

from compare import compare, load
from scenarios import (
    KERNEL_WARMUP_CELLS,
    NAMES,
    WARMUP_TURNS,
    harness,
    pool_processes,
    run_scenario,
)
from transport import Pacer
from safety import SecretScanner
import resources


def summary(values):
    middle = statistics.median(values)
    return {
        "median": middle,
        "mad": statistics.median(abs(value - middle) for value in values),
        "samples": values,
    }


def slope(points):
    center = statistics.mean(x for x, _ in points)
    average = statistics.mean(y for _, y in points)
    return sum((x - center) * (y - average) for x, y in points) / sum(
        (x - center) ** 2 for x, _ in points
    )


def memory_metrics(run):
    checkpoints = run["checkpoints"]
    if not checkpoints:
        return {}
    final = checkpoints[-1]
    result = {
        "retained_tree_pss": final["tree_pss"],
        "retained_tree_uss": final["tree_uss"],
    }
    result.update({"beam_" + key: value for key, value in final["beam"].items()})
    result["session_process_bytes"] = sum(
        item["bytes"] for item in final["sessions"].values()
    )
    result["kernel_uss"] = sum(
        item["uss"]
        for item in final["processes"]
        if item["kind"] in ("kernel", "forked")
    )
    return result


def aggregate(name, runs):
    memory = [memory_metrics(run) for run in runs]
    return {
        "runs": len(runs),
        "wall_ms": summary([run["wall_ms"] for run in runs]),
        "memory": {key: summary([row[key] for row in memory]) for key in memory[0]},
        "wire": runs[0]["wire"],
        "requests": runs[0]["requests"],
        "recordings": runs[0]["recordings"],
        "observations": runs,
    }


def warmup_retained(runs):
    """Memory after the warm-up turns minus before the first turn; not gated."""

    def change(run, value):
        points = {
            (point["label"], point["count"]): point for point in run["checkpoints"]
        }
        return value(points[("turn", WARMUP_TURNS)]) - value(points[("created", 0)])

    return {
        "turns": WARMUP_TURNS,
        "tree_uss": summary([change(run, lambda p: p["tree_uss"]) for run in runs]),
        "beam": summary([change(run, lambda p: p["beam"]["total"]) for run in runs]),
    }


def held_kernel(point, session):
    """The process of the kernel a session holds at a checkpoint."""
    pid = point["kernels"][session]
    return next(item for item in point["processes"] if item["pid"] == pid)


def kernel_uss(run, label, count):
    point = next(
        p for p in run["checkpoints"] if (p["label"], p["count"]) == (label, count)
    )
    session = next(iter(point["kernels"]))
    return held_kernel(point, session)["uss"]


def pool_report(runs):
    """Pool cost when warm, and one session kernel started cold vs adopted; not gated."""

    def at(run, label):
        return next(p for p in run["checkpoints"] if p["label"] == label)

    def kernel(run, label, which):
        item = held_kernel(at(run, label), run["pool_sessions"][which])
        return item

    return {
        "warm_pss": summary(
            [
                sum(
                    p["pss"]
                    for p in pool_processes(
                        at(run, "warm")["kernels"], at(run, "warm")["processes"]
                    )
                )
                for run in runs
            ]
        ),
        "first_kernel": {
            "kinds": [kernel(run, "cold", "first")["kind"] for run in runs],
            "uss": summary([kernel(run, "cold", "first")["uss"] for run in runs]),
            "pss": summary([kernel(run, "cold", "first")["pss"] for run in runs]),
        },
        "adopted_kernel": {
            "kinds": [kernel(run, "adopted", "second")["kind"] for run in runs],
            "uss": summary([kernel(run, "adopted", "second")["uss"] for run in runs]),
            "pss": summary([kernel(run, "adopted", "second")["pss"] for run in runs]),
        },
    }


def derived(runs, mode):
    metrics = {}

    def add(name, values):
        stats = summary(values)
        metrics[name] = {
            "value": stats["median"],
            "noise": stats["mad"],
            "samples": stats["samples"],
        }
        if len(values) == 1:
            metrics[name]["single_sample"] = True

    if runs["boot"][0]["checkpoints"]:
        add(
            "boot_tree_uss", [run["checkpoints"][0]["tree_uss"] for run in runs["boot"]]
        )
        # USS, not PSS: PSS splits shared library pages with whatever other
        # process maps the same files (another BEAM on the host moved it 8 MB).
        for label, field in (("tree_uss", "tree_uss"), ("beam", "beam")):

            def value(point):
                return point["beam"]["total"] if field == "beam" else point[field]

            add(
                "idle_session_" + label,
                [
                    slope(
                        [
                            (point["count"], value(point))
                            for point in run["checkpoints"]
                            if point["label"] == "sessions"
                        ]
                    )
                    for run in runs["fan-out"]
                ],
            )
            add(
                "retained_turn_" + label,
                [
                    slope(
                        [
                            (point["count"], value(point))
                            for point in run["checkpoints"]
                            if point["label"] == "turn"
                            and point["count"] >= WARMUP_TURNS
                        ]
                    )
                    for run in runs["long"]
                ],
            )
        add(
            "turn_peak_tree_uss",
            [
                max(
                    turn["peak_tree_uss"]
                    for name in ("chat", "coding", "long")
                    for turn in runs[name][i]["turns"]
                    if "peak_tree_uss" in turn
                )
                for i in range(len(runs["long"]))
            ],
        )
    add(
        "request_bytes_per_turn",
        [
            statistics.mean(turn["sent"] for turn in run["turns"])
            for run in runs["long"]
        ],
    )
    add(
        "request_bytes_slope",
        [
            slope([(turn["turn"], turn["sent"]) for turn in run["turns"]])
            for run in runs["long"]
        ],
    )
    add("attach_bytes", [run["wire"]["attach_received"] for run in runs["long"]])
    # Live replies are slow enough that the 3 s model-wait release frees the
    # kernel between cells, so kernel numbers come from replay only.
    if mode == "replay" and runs["kernel"][0]["checkpoints"]:
        add("kernel_rest_uss", [kernel_uss(run, "turn", 1) for run in runs["kernel"]])
        add(
            "kernel_cell_slope",
            [
                slope(
                    [
                        (point["count"], kernel_uss(run, "turn", point["count"]))
                        for point in run["checkpoints"]
                        if point["label"] == "turn"
                        and point["count"] >= 1 + KERNEL_WARMUP_CELLS
                    ]
                )
                for run in runs["kernel"]
            ],
        )
        add(
            "kernel_pool_pss",
            [
                sum(
                    item["pss"]
                    for item in pool_processes(point["kernels"], point["processes"])
                )
                for run in runs["pool"]
                for point in run["checkpoints"]
                if point["label"] == "rest"
            ],
        )
    return metrics


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--mode", choices=("replay", "live"), required=True)
    parser.add_argument("--runs", type=int, default=5)
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--baseline")
    parser.add_argument("--record", action="store_true")
    parser.add_argument(
        "--jobs",
        type=int,
        help="repetitions of a scenario run at once, each on its own daemon "
        "(default: all of them in replay, one in live mode)",
    )
    parser.add_argument(
        "--recordings-dir", type=Path, default=Path(__file__).with_name("recordings")
    )
    parser.add_argument(
        "--url",
        default=os.environ.get("BONSAI_URL") or os.environ.get("OPENAI_BASE_URL"),
    )
    parser.add_argument("--model", default=os.environ.get("OPENAI_MODEL"))
    parser.add_argument("--key-env", default="OPENAI_API_KEY")
    parser.add_argument("--provider", default="bonsai")
    parser.add_argument(
        "--rate",
        type=float,
        default=0,
        help="live provider requests per minute (0 means unpaced)",
    )
    parser.add_argument(
        "--scenario",
        choices=NAMES,
        action="append",
        help="diagnostic subset; never used by CI",
    )
    args = parser.parse_args()
    if args.runs < 1 or (args.record and (args.mode != "live" or args.runs != 1)):
        parser.error("runs must be positive; recording requires --mode live --runs 1")
    if args.mode == "live" and not args.url:
        parser.error("live mode requires --url, BONSAI_URL or OPENAI_BASE_URL")
    if args.rate < 0 or (args.record and args.provider != "bonsai"):
        parser.error("rate must be nonnegative; only bonsai recordings are committed")
    pacer = Pacer(args.rate)
    jobs = args.jobs or (args.runs if args.mode == "replay" else 1)
    scanner = SecretScanner(extra=[os.environ.get(args.key_env, "")])
    if args.mode == "replay" and sys.platform != "linux":
        parser.error("replay memory gates require Linux /proc")
    harness.snapshot_daemon()
    names = args.scenario or NAMES
    raw = {}
    result = {
        "schema": 1,
        "commit": os.environ.get("PERF_COMMIT")
        or subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip(),
        "base_commit": None,
        "ref": os.environ.get("TANGLED_REF", os.environ.get("GITHUB_REF", "local")),
        "event": os.environ.get(
            "TANGLED_PIPELINE_KIND", os.environ.get("GITHUB_EVENT_NAME", "manual")
        ),
        "mode": args.mode,
        "provider": args.provider if args.mode == "live" else "bonsai-recording",
        "model": args.model,
        "runner": {
            "cpus": resources.cpu_count(),
            "memory_mib": min(
                [
                    int((directory / "memory.max").read_text())
                    for directory in resources.cgroups()
                    if (directory / "memory.max").exists()
                    and (directory / "memory.max").read_text().strip() != "max"
                ]
                or [os.sysconf("SC_PHYS_PAGES") * os.sysconf("SC_PAGE_SIZE")]
            )
            / 1048576,
            "kernel": platform.release(),
        },
        "scenarios": {},
        "gated": {},
        "verdict": "pass",
    }
    try:
        for name in names:
            print(f"{args.mode} {name} x{args.runs} ({jobs} at a time)", flush=True)
            started = time.monotonic()

            def repetition(_index, name=name):
                return run_scenario(
                    name,
                    args.mode,
                    args.record,
                    args.url,
                    model=args.model,
                    key=os.environ.get(args.key_env),
                    pacer=pacer,
                    recordings_dir=args.recordings_dir,
                )

            with ThreadPoolExecutor(max_workers=jobs) as pool:
                raw[name] = list(pool.map(repetition, range(args.runs)))
            print(f"  {name}: {time.monotonic() - started:.1f} s", flush=True)
            result["scenarios"][name] = aggregate(name, raw[name])
            if name == "long" and raw[name][0]["checkpoints"]:
                result["scenarios"][name]["warmup_retained"] = warmup_retained(
                    raw[name]
                )
            if name == "pool" and args.mode == "replay" and raw[name][0]["checkpoints"]:
                result["scenarios"][name]["pool"] = pool_report(raw[name])
            args.out.parent.mkdir(parents=True, exist_ok=True)
            scanner.write_json(args.out, result)
        # A recording is not a measurement: slow live replies release kernels.
        if not args.scenario and not args.record:
            result["gated"] = derived(raw, args.mode)
            compare(result, load(args.baseline) if args.baseline else None)
        scanner.write_json(args.out, result)
        return int(result["verdict"] == "fail")
    finally:
        harness.shutdown()


if __name__ == "__main__":
    raise SystemExit(main())