#!/usr/bin/env python3
from __future__ import annotations

import argparse
import ctypes
import json
import shutil
import subprocess
import tempfile
import time
from pathlib import Path
from typing import Any


class ProcessMemoryCounters(ctypes.Structure):
    _fields_ = [
        ("cb", ctypes.c_ulong),
        ("PageFaultCount", ctypes.c_ulong),
        ("PeakWorkingSetSize", ctypes.c_size_t),
        ("WorkingSetSize", ctypes.c_size_t),
        ("QuotaPeakPagedPoolUsage", ctypes.c_size_t),
        ("QuotaPagedPoolUsage", ctypes.c_size_t),
        ("QuotaPeakNonPagedPoolUsage", ctypes.c_size_t),
        ("QuotaNonPagedPoolUsage", ctypes.c_size_t),
        ("PagefileUsage", ctypes.c_size_t),
        ("PeakPagefileUsage", ctypes.c_size_t),
    ]


class Server:
    def __init__(self, exe: Path, data_dir: Path) -> None:
        self.process = subprocess.Popen(
            [str(exe), "--data-dir", str(data_dir)],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            encoding="utf-8",
            bufsize=1,
        )
        self.request_id = 1
        self.max_response_bytes = 0
        self.request(
            "initialize",
            {
                "protocolVersion": "2024-11-05",
                "capabilities": {},
                "clientInfo": {"name": "p2-benchmark", "version": "1"},
            },
        )
        self.notify("notifications/initialized", {})

    def request(self, method: str, params: dict[str, Any]) -> dict[str, Any]:
        assert self.process.stdin and self.process.stdout
        request_id = self.request_id
        self.request_id += 1
        self.process.stdin.write(
            json.dumps(
                {"jsonrpc": "2.0", "id": request_id, "method": method, "params": params},
                separators=(",", ":"),
            )
            + "\n"
        )
        self.process.stdin.flush()
        line = self.process.stdout.readline()
        self.max_response_bytes = max(self.max_response_bytes, len(line.encode("utf-8")))
        response = json.loads(line)
        if "error" in response:
            raise RuntimeError(response)
        return response["result"]

    def call(self, name: str, arguments: dict[str, Any]) -> dict[str, Any]:
        result = self.request("tools/call", {"name": name, "arguments": arguments})
        payload = json.loads(result["content"][0]["text"])
        if result.get("isError"):
            raise RuntimeError(payload)
        return payload

    def notify(self, method: str, params: dict[str, Any]) -> None:
        assert self.process.stdin
        self.process.stdin.write(
            json.dumps({"jsonrpc": "2.0", "method": method, "params": params}) + "\n"
        )
        self.process.stdin.flush()

    def peak_working_set(self) -> int:
        counters = ProcessMemoryCounters()
        counters.cb = ctypes.sizeof(counters)
        handle = ctypes.windll.kernel32.OpenProcess(0x1000 | 0x0400, False, self.process.pid)
        try:
            if handle and ctypes.windll.psapi.GetProcessMemoryInfo(
                handle, ctypes.byref(counters), counters.cb
            ):
                return int(counters.PeakWorkingSetSize)
            return 0
        finally:
            if handle:
                ctypes.windll.kernel32.CloseHandle(handle)

    def close(self) -> None:
        if self.process.stdin:
            self.process.stdin.close()
        code = self.process.wait(timeout=30)
        if code != 0:
            raise RuntimeError(self.process.stderr.read() if self.process.stderr else code)


def read_all(server: Server, session_id: str) -> tuple[int, int]:
    cursor = None
    count = pages = 0
    while True:
        arguments: dict[str, Any] = {"session_id": session_id, "limit": 1000}
        if cursor:
            arguments["cursor"] = cursor
        payload = server.call("read_messages", arguments)
        count += len(payload["history"])
        pages += 1
        cursor = payload.get("nextCursor")
        if not cursor:
            return count, pages


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--server", type=Path, required=True)
    parser.add_argument("--messages", type=int, choices=(10000, 50000), required=True)
    parser.add_argument("--sessions", type=int, choices=(1, 100), required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    data_dir = Path(tempfile.mkdtemp(prefix="tiny-convmem-p2-"))
    try:
        server = Server(args.server.resolve(), data_dir)
        write_started = time.perf_counter()
        for index in range(args.messages):
            server.call(
                "write_message",
                {
                    "session_id": f"session-{index % args.sessions}",
                    "role": "user",
                    "content": f"history item {index}",
                },
            )
        write_ms = (time.perf_counter() - write_started) * 1000
        server.close()

        restart_started = time.perf_counter()
        server = Server(args.server.resolve(), data_dir)
        restart_ms = (time.perf_counter() - restart_started) * 1000
        read_started = time.perf_counter()
        total = pages = 0
        for session in range(args.sessions):
            session_count, session_pages = read_all(server, f"session-{session}")
            total += session_count
            pages += session_pages
        read_ms = (time.perf_counter() - read_started) * 1000
        rebuild_started = time.perf_counter()
        server.call("rebuild_index", {})
        rebuild_ms = (time.perf_counter() - rebuild_started) * 1000
        result = {
            "messages": args.messages,
            "sessions": args.sessions,
            "messages_read": total,
            "pages": pages,
            "write_ms": write_ms,
            "restart_ms": restart_ms,
            "full_read_ms": read_ms,
            "rebuild_ms": rebuild_ms,
            "peak_working_set_bytes": server.peak_working_set(),
            "max_response_bytes": server.max_response_bytes,
        }
        server.close()
        args.output.write_text(json.dumps(result, indent=2), encoding="utf-8")
        print(json.dumps(result, indent=2))
        return 0 if total == args.messages and server.max_response_bytes <= 8 * 1024 * 1024 else 1
    finally:
        shutil.rmtree(data_dir, ignore_errors=True)


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