#!/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

FIXTURE_PREFIX = "p1-cp-benchmark-"


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),
    ]


def make_ulid(timestamp_ms: int, random_bits: int) -> str:
    alphabet = "0123456789ABCDEFGHJKMNPQRSTVWXYZ"
    bytes_data = bytearray(16)
    bytes_data[0] = (timestamp_ms >> 40) & 0xFF
    bytes_data[1] = (timestamp_ms >> 32) & 0xFF
    bytes_data[2] = (timestamp_ms >> 24) & 0xFF
    bytes_data[3] = (timestamp_ms >> 16) & 0xFF
    bytes_data[4] = (timestamp_ms >> 8) & 0xFF
    bytes_data[5] = timestamp_ms & 0xFF
    random_bytes = random_bits.to_bytes(10, "big")
    bytes_data[6:16] = random_bytes
    output = []
    for index in range(26):
        bit = index * 5
        value = 0
        for part in range(5):
            source_bit = bit + part - 2
            value <<= 1
            if source_bit >= 0:
                value |= (bytes_data[source_bit // 8] >> (7 - source_bit % 8)) & 1
        output.append(alphabet[value])
    return "".join(output)


def write_fixture(path: Path, count: int) -> None:
    with open(path, "w", encoding="utf-8") as file:
        file.write(json.dumps({"type": "tiny-convmem-manifest", "format_version": 2}) + "\n")
        base_time = 1_000_000_000_000
        for index in range(count):
            timestamp = base_time + index
            message = {
                "message_id": make_ulid(timestamp, index),
                "session_id": f"session-{index % 100}",
                "timestamp": timestamp,
                "role": "user",
                "content": f"P1 benchmark message {index} with enough content to exercise import parsing.",
                "revision": 1,
                "updated_at": timestamp,
                "metadata": {
                    "memory_tier": "short",
                    "memory_kind": "conversation",
                    "importance": 50,
                },
            }
            file.write(json.dumps(message, separators=(",", ":")) + "\n")


def run_import(executable: Path, data_dir: Path, input_path: Path) -> tuple[float, int]:
    shutil.rmtree(data_dir, ignore_errors=True)
    start = time.perf_counter()
    process = subprocess.Popen(
        [str(executable), "--data-dir", str(data_dir), "--import", str(input_path)],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
    )
    stdout, stderr = process.communicate()
    elapsed = time.perf_counter() - start
    if process.returncode != 0:
        raise RuntimeError(stderr)
    counters = ProcessMemoryCounters()
    counters.cb = ctypes.sizeof(counters)
    peak = 0
    if ctypes.windll.psapi.GetProcessMemoryInfo(  # type: ignore[attr-defined]
        int(process._handle), ctypes.byref(counters), counters.cb  # type: ignore[attr-defined]
    ):
        peak = int(counters.PeakWorkingSetSize)
    print(stdout.strip())
    return elapsed, peak


def run_export(executable: Path, data_dir: Path, output_path: Path) -> float:
    start = time.perf_counter()
    result = subprocess.run(
        [str(executable), "--data-dir", str(data_dir), "--export", str(output_path)],
        capture_output=True,
        text=True,
        check=True,
    )
    elapsed = time.perf_counter() - start
    print(result.stdout.strip())
    return elapsed


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--server", type=Path, required=True)
    parser.add_argument("--counts", type=int, nargs="+", default=[10000, 50000])
    args = parser.parse_args()

    executable = args.server.resolve()
    work_dir = Path(tempfile.mkdtemp(prefix=FIXTURE_PREFIX))
    print(f"work_dir: {work_dir}")

    try:
        for count in args.counts:
            fixture = work_dir / f"fixture-{count}.jsonl"
            print(f"\n=== P1-CP benchmark: {count:,} messages ===")
            write_fixture(fixture, count)
            fixture_size = fixture.stat().st_size
            print(f"fixture size: {fixture_size:,} bytes")

            import_dir = work_dir / f"import-{count}"
            import_ms, import_peak = run_import(executable, import_dir, fixture)
            print(f"import: {import_ms:.3f}s ({count / import_ms:,.0f} msg/s)")
            print(f"import peak working set: {import_peak:,} bytes")

            export_path = work_dir / f"exported-{count}.jsonl"
            export_ms = run_export(executable, import_dir, export_path)
            print(f"export: {export_ms:.3f}s ({count / export_ms:,.0f} msg/s)")

            roundtrip_dir = work_dir / f"roundtrip-{count}"
            roundtrip_ms, roundtrip_peak = run_import(executable, roundtrip_dir, export_path)
            print(f"round-trip import: {roundtrip_ms:.3f}s ({count / roundtrip_ms:,.0f} msg/s)")

            export_size = export_path.stat().st_size
            print(f"export size: {export_size:,} bytes")
            print(
                json.dumps(
                    {
                        "count": count,
                        "fixture_bytes": fixture_size,
                        "import_ms": round(import_ms * 1000, 3),
                        "import_peak_working_set_bytes": import_peak,
                        "export_ms": round(export_ms * 1000, 3),
                        "roundtrip_ms": round(roundtrip_ms * 1000, 3),
                        "roundtrip_peak_working_set_bytes": roundtrip_peak,
                        "export_bytes": export_size,
                    }
                )
            )
    finally:
        shutil.rmtree(work_dir, ignore_errors=True)


if __name__ == "__main__":
    main()