"""22: Capacity planning. Run N sessions concurrently and measure the peak memory of the claude subprocesses."""

import asyncio
import os
import time

from claude_agent_sdk import ClaudeAgentOptions, ResultMessage, query
from demo_util import PROJECT

N = 3


def claude_rss_mib():
    """Sum the resident memory of every 'claude' child process (Linux /proc)."""
    total, count = 0, 0
    for pid in os.listdir("/proc"):
        if not pid.isdigit():
            continue
        try:
            with open(f"/proc/{pid}/status") as f:
                status = f.read()
            if "Name:\tclaude" not in status:
                continue
            for line in status.splitlines():
                if line.startswith("VmRSS:"):
                    total += int(line.split()[1]) // 1024
                    count += 1
        except (FileNotFoundError, ProcessLookupError, PermissionError):
            pass
    return count, total


peak = {"procs": 0, "mib": 0}


async def sampler(stop: asyncio.Event):
    while not stop.is_set():
        procs, mib = claude_rss_mib()
        peak["procs"] = max(peak["procs"], procs)
        peak["mib"] = max(peak["mib"], mib)
        await asyncio.sleep(0.5)


async def one(i):
    t0 = time.time()
    async for m in query(prompt=f"Session {i}: read src/app.py and say in one sentence what find_user does.",
                         options=ClaudeAgentOptions(cwd=PROJECT, allowed_tools=["Read"], max_turns=3, effort="low")):
        if isinstance(m, ResultMessage):
            return i, m.subtype, round(time.time() - t0, 1), m.total_cost_usd or 0


async def main():
    stop = asyncio.Event()
    task = asyncio.create_task(sampler(stop))
    t0 = time.time()
    results = await asyncio.gather(*(one(i) for i in range(N)))
    wall = time.time() - t0
    stop.set(); await task

    for i, subtype, secs, cost in results:
        print(f"session {i}: {subtype} in {secs}s, ${cost:.4f}")
    print(f"\n{N} sessions in parallel: wall={wall:.1f}s, peak claude processes={peak['procs']}, "
          f"peak total RSS={peak['mib']} MiB (~{peak['mib'] // max(peak['procs'], 1)} MiB per session)")
    print(f"planning: agents per host = (host RAM - overhead) / per-session ceiling; "
          f"at ~{max(peak['mib'] // max(peak['procs'], 1), 1)} MiB a 16 GiB host fits ~{(16 * 1024 - 2048) // max(peak['mib'] // max(peak['procs'], 1), 1)} idle-sized sessions")


asyncio.run(main())