#!/usr/bin/env python3
"""Measure Pi Core process startup through its JSONL ready event."""

from __future__ import annotations

import argparse
import json
import os
import selectors
import subprocess
import time
from pathlib import Path


DEFAULT_CONFIG = {
    "providers": {
        "amazon-bedrock": {
            "models": [
                {
                    "id": "benchmark-model",
                    "name": "Benchmark Model",
                    "api": "bedrock-converse-stream",
                    "reasoning": True,
                    "input": ["text", "image"],
                    "contextWindow": 200000,
                    "maxTokens": 16384,
                    "cost": {
                        "input": 0,
                        "output": 0,
                        "cacheRead": 0,
                        "cacheWrite": 0,
                    },
                }
            ]
        }
    }
}


def process_metrics(pid: int) -> dict[str, int]:
    metrics: dict[str, int] = {}
    try:
        stat_tail = Path(f"/proc/{pid}/stat").read_text().rsplit(") ", 1)[1].split()
        metrics["minor_faults"] = int(stat_tail[7])
        metrics["major_faults"] = int(stat_tail[9])
    except (OSError, IndexError, ValueError):
        pass
    try:
        for line in Path(f"/proc/{pid}/io").read_text().splitlines():
            key, value = line.split(":", 1)
            if key in {"rchar", "read_bytes"}:
                metrics[key] = int(value.strip())
    except (OSError, ValueError):
        pass
    try:
        for line in Path(f"/proc/{pid}/smaps_rollup").read_text().splitlines():
            key, _, value = line.partition(":")
            if key in {"Rss", "Pss_File", "Shared_Clean", "Private_Clean"}:
                metrics[f"{key.lower()}_kb"] = int(value.strip().split()[0])
    except (OSError, ValueError):
        pass
    return metrics


def read_json_line(proc: subprocess.Popen[str], timeout: float) -> dict:
    assert proc.stdout is not None
    selector = selectors.DefaultSelector()
    selector.register(proc.stdout, selectors.EVENT_READ)
    deadline = time.monotonic() + timeout
    while True:
        remaining = deadline - time.monotonic()
        if remaining <= 0:
            raise TimeoutError("Pi Core did not emit JSON before the timeout")
        if not selector.select(remaining):
            continue
        line = proc.stdout.readline()
        if not line:
            stderr = proc.stderr.read() if proc.stderr else ""
            raise RuntimeError(
                f"Pi Core exited before emitting JSON (rc={proc.poll()}): {stderr}"
            )
        try:
            value = json.loads(line)
        except json.JSONDecodeError:
            continue
        if isinstance(value, dict):
            return value


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--binary",
        default="/opt/investibook/agentkit/runtime/pi-core/dist/pi-core",
    )
    parser.add_argument("--cwd", default="/workspace")
    parser.add_argument("--timeout", type=float, default=45)
    args = parser.parse_args()

    env = os.environ.copy()
    env.setdefault(
        "INVESTIBOOK_PI_CORE_CONFIG_JSON",
        json.dumps(DEFAULT_CONFIG, separators=(",", ":")),
    )
    started = time.monotonic()
    proc = subprocess.Popen(
        [args.binary, "--model", "amazon-bedrock/benchmark-model"],
        cwd=args.cwd,
        env=env,
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        bufsize=1,
    )
    try:
        ready = read_json_line(proc, args.timeout)
        ready_ms = (time.monotonic() - started) * 1000
        if ready.get("type") != "ready":
            raise RuntimeError(f"Expected Pi ready event, got: {ready}")

        command_started = time.monotonic()
        assert proc.stdin is not None
        proc.stdin.write('{"id":"probe-1","type":"get_state"}\n')
        proc.stdin.flush()
        while True:
            response = read_json_line(proc, args.timeout)
            if response.get("type") == "response" and response.get("id") == "probe-1":
                break
        command_ms = (time.monotonic() - command_started) * 1000
        timings = ready.get("data", {}).get("startupTimings", {})
        print(
            json.dumps(
                {
                    "ready_ms": round(ready_ms, 1),
                    "get_state_ms": round(command_ms, 1),
                    "total_ms": round(ready_ms + command_ms, 1),
                    "startup_timings": timings,
                    "process_metrics": process_metrics(proc.pid),
                },
                separators=(",", ":"),
            )
        )
    finally:
        proc.terminate()
        try:
            proc.wait(timeout=2)
        except subprocess.TimeoutExpired:
            proc.kill()
            proc.wait(timeout=2)


if __name__ == "__main__":
    main()