"""Measure the memory footprint of mcp-calculator (see docs/memory.md).

Usage (inside the project venv):
    python scripts/measure_memory.py imports      # RSS cost of each library, cumulative
    python scripts/measure_memory.py server       # start an HTTP server, warm it up, measure all processes
    python scripts/measure_memory.py server --workers 4

Only the standard library is used, so it also runs in a container. RSS is read
via ``ps`` (macOS and Linux/procps).
"""

from __future__ import annotations

import argparse
import json
import os
import socket
import subprocess
import sys
import time
import urllib.request

# Cumulative import steps in the order the server loads them.
IMPORT_STEPS = [
    ("Python interpreter", []),
    ("+ NumPy", ["numpy"]),
    ("+ SymPy", ["sympy"]),
    ("+ SciPy (scipy.stats)", ["scipy.stats"]),
    ("+ Pint", ["pint"]),
    ("+ MCP SDK, Pydantic", ["mcp.server.mcpserver"]),
    ("+ uvicorn, Starlette", ["uvicorn", "starlette"]),
    ("+ Pebble", ["pebble"]),
]

WARMUP_CALLS = [
    ("evaluate", {"expression": "sqrt(2)", "precision": 50}),
    ("solve", {"equations": ["x^2 + y^2 = 25", "x - y = 1"]}),
    ("calculus", {"operation": "integral", "expression": "x*e^x"}),
    ("analyze_function", {"expression": "x^3 - 3x"}),
    ("matrix", {"operation": "eigenvalues", "matrix": [[2, 1], [1, 2]]}),
    ("statistics", {"data": [1, 2, 3, 4, 5]}),
    ("probability", {"distribution": "normal", "operation": "cdf", "x": 1}),
    ("convert_units", {"quantity": "100 km/h", "to_unit": "m/s"}),
]


def rss_mb(pid: int) -> float:
    return int(subprocess.check_output(["ps", "-o", "rss=", "-p", str(pid)])) / 1024


def measure_imports() -> None:
    print(f"{'step':<28}{'RSS total':>10}{'delta':>8}")
    previous = None
    for label, _ in IMPORT_STEPS:
        modules = [m for _, mods in IMPORT_STEPS[: [s[0] for s in IMPORT_STEPS].index(label) + 1] for m in mods]
        code = "import importlib, os, subprocess\n"
        code += "".join(f"importlib.import_module({m!r})\n" for m in modules)
        code += "print(int(subprocess.check_output(['ps', '-o', 'rss=', '-p', str(os.getpid())])))"
        total = int(subprocess.check_output([sys.executable, "-c", code])) / 1024
        delta = "" if previous is None else f"+{total - previous:.0f}"
        print(f"{label:<28}{total:>8.0f} MB{delta:>8}")
        previous = total


def _children(pid: int) -> list[tuple[int, str]]:
    out = subprocess.check_output(["ps", "-A", "-o", "pid=,ppid=,command="], text=True)
    result = []
    for line in out.splitlines():
        parts = line.split(None, 2)
        if len(parts) == 3 and int(parts[1]) == pid:
            result.append((int(parts[0]), parts[2]))
    return result


def _call(url: str, tool: str, arguments: dict) -> None:
    body = json.dumps({"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": tool, "arguments": arguments}}).encode()
    request = urllib.request.Request(
        url,
        data=body,
        headers={
            "Content-Type": "application/json",
            "Accept": "application/json, text/event-stream",
            "MCP-Protocol-Version": "2025-06-18",
        },
    )
    with urllib.request.urlopen(request, timeout=60) as response:
        response.read()


def measure_server(workers: int) -> None:
    with socket.socket() as s:
        s.bind(("127.0.0.1", 0))
        port = s.getsockname()[1]
    env = {**os.environ, "MCP_CALC_TRANSPORT": "http", "MCP_CALC_PORT": str(port), "MCP_CALC_WORKERS": str(workers), "MCP_CALC_LOG_LEVEL": "WARNING"}
    proc = subprocess.Popen([sys.executable, "-m", "mcp_calculator"], env=env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    base = f"http://127.0.0.1:{port}"
    try:
        started = time.perf_counter()
        while True:
            try:
                urllib.request.urlopen(f"{base}/readyz", timeout=1).read()
                break
            except OSError:
                if time.perf_counter() - started > 60:
                    raise SystemExit("server did not become ready") from None
                time.sleep(0.1)
        ready_after = time.perf_counter() - started

        first = time.perf_counter()
        _call(f"{base}/mcp", "evaluate", {"expression": "1+1"})
        first_call = time.perf_counter() - first
        for _ in range(workers):  # every worker should have computed something
            for tool, arguments in WARMUP_CALLS:
                _call(f"{base}/mcp", tool, arguments)
        time.sleep(1)

        rows = [("server", rss_mb(proc.pid))]
        for pid, command in _children(proc.pid):
            rows.append(("resource tracker" if "resource_tracker" in command else "worker", rss_mb(pid)))
        print(f"ready after {ready_after:.2f}s, first call {first_call:.2f}s, {workers} worker(s)\n")
        for name, mb in rows:
            print(f"{name:<18}{mb:>6.0f} MB")
        print(f"{'total':<18}{sum(mb for _, mb in rows):>6.0f} MB")
    finally:
        proc.terminate()
        proc.wait(timeout=30)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("mode", choices=["imports", "server"])
    parser.add_argument("--workers", type=int, default=2)
    args = parser.parse_args()
    if args.mode == "imports":
        measure_imports()
    else:
        measure_server(args.workers)


if __name__ == "__main__":
    main()