#!/usr/bin/env python3
"""
tiny-convmem-mcp ユーザ視点負荷テスト＆ベンチマーク

MCP Python SDK (mcp.client.stdio + mcp.client.session) を使い、
実際に tiny-convmem-mcp サーバプロセスを stdio 経由で起動し、
エンドツーエンドのレイテンシとスループットを測定します。

ポイント:
- 各シナリオは独立した temp data directory で実行（状態累積を排除）
- write_message は `metadata_json` 文字列を使用
- 応答の isError / status / 件数を検証
- rebuild_index は 5 回反復
"""

from __future__ import annotations

import argparse
import asyncio
import json
import os
import random
import shutil
import statistics
import subprocess
import sys
import tempfile
import time
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any, Awaitable, Callable

from mcp.client.session import ClientSession
from mcp.client.stdio import StdioServerParameters, stdio_client
from mcp.types import CallToolResult

# ---------------------------------------------------------------------------
# 定数・デフォルト
# ---------------------------------------------------------------------------
DEFAULT_ITERATIONS = 100
DEFAULT_CONTENT_MEDIUM = 1024  # bytes
DEFAULT_CONTENT_8K = 8 * 1024  # bytes
DEFAULT_CONTENT_LARGE = 32 * 1024  # bytes
DEFAULT_CONTENT_MAX = 1024 * 1024  # bytes
WARMUP_COUNT = 5
REBUILD_ITERATIONS = 5
P1_CORPUS_MESSAGES = 20

ROOT = Path(__file__).resolve().parent.parent.parent
DEFAULT_RELEASE_EXE = (
    ROOT / "build" / "windows-vs2026" / "src" / "Release" / "tiny-convmem-mcp.exe"
)


# ---------------------------------------------------------------------------
# データクラス
# ---------------------------------------------------------------------------
@dataclass
class LatencySample:
    operation: str
    duration_ms: float
    success: bool
    error: str | None = None
    bytes_in: int = 0
    bytes_out: int = 0


@dataclass
class ScenarioResult:
    name: str
    iterations: int
    successes: int
    failures: int
    samples: list[float] = field(repr=False, default_factory=list)
    throughput_ops_per_sec: float = 0.0
    latency_p50_ms: float = 0.0
    latency_p95_ms: float = 0.0
    latency_p99_ms: float = 0.0
    latency_min_ms: float = 0.0
    latency_max_ms: float = 0.0
    latency_mean_ms: float = 0.0
    total_duration_ms: float = 0.0
    total_bytes_in: int = 0
    total_bytes_out: int = 0
    notes: str = ""


@dataclass
class BenchmarkReport:
    server: str
    build_type: str
    exe_mtime: str
    iterations: int
    warmup: int
    scenarios: list[ScenarioResult] = field(default_factory=list)
    domain_index: list[dict[str, Any]] = field(default_factory=list)
    health_after: dict[str, Any] | None = None
    schema_check: dict[str, Any] | None = None
    errors: list[str] = field(default_factory=list)

    def as_dict(self) -> dict[str, Any]:
        return asdict(self)


# ---------------------------------------------------------------------------
# ユーティリティ
# ---------------------------------------------------------------------------
def percentile(sorted_values: list[float], p: float) -> float:
    if not sorted_values:
        return 0.0
    n = len(sorted_values)
    if n == 1:
        return sorted_values[0]
    k = (n - 1) * p / 100.0
    f = int(k)
    c = min(f + 1, n - 1)
    if f == c:
        return sorted_values[f]
    return sorted_values[f] + (sorted_values[c] - sorted_values[f]) * (k - f)


def summarize(name: str, samples: list[LatencySample]) -> ScenarioResult:
    ok = [s for s in samples if s.success]
    failed = [s for s in samples if not s.success]
    durations = sorted(s.duration_ms for s in ok)
    result = ScenarioResult(
        name=name,
        iterations=len(samples),
        successes=len(ok),
        failures=len(failed),
        samples=durations,
        total_duration_ms=sum(s.duration_ms for s in samples),
        total_bytes_in=sum(s.bytes_in for s in ok),
        total_bytes_out=sum(s.bytes_out for s in ok),
        notes=failed[0].error if failed else "",
    )
    if durations:
        result.latency_min_ms = durations[0]
        result.latency_max_ms = durations[-1]
        result.latency_mean_ms = statistics.mean(durations)
        result.latency_p50_ms = percentile(durations, 50)
        result.latency_p95_ms = percentile(durations, 95)
        result.latency_p99_ms = percentile(durations, 99)
    if result.total_duration_ms > 0 and ok:
        result.throughput_ops_per_sec = len(ok) / (result.total_duration_ms / 1000.0)
    return result


def single_sample(name: str, duration_ms: float, notes: str) -> ScenarioResult:
    return ScenarioResult(
        name=name,
        iterations=1,
        successes=1,
        failures=0,
        samples=[duration_ms],
        throughput_ops_per_sec=1000.0 / duration_ms if duration_ms > 0 else 0.0,
        latency_p50_ms=duration_ms,
        latency_p95_ms=duration_ms,
        latency_p99_ms=duration_ms,
        latency_min_ms=duration_ms,
        latency_max_ms=duration_ms,
        latency_mean_ms=duration_ms,
        total_duration_ms=duration_ms,
        notes=notes,
    )


def text_content(result: CallToolResult) -> str:
    if result.isError:
        raise RuntimeError(f"tool returned isError: {result.content!r}")
    if not result.content:
        return ""
    return "".join(item.text for item in result.content if item.type == "text")


def json_result(result: CallToolResult) -> dict:
    text = text_content(result)
    payload = json.loads(text)
    return payload


def random_content(size: int) -> str:
    alphabet = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789 "
    return "".join(random.choice(alphabet) for _ in range(size))


def make_metadata_json(tags: list[str] | None = None, tier: str = "short") -> str:
    obj: dict[str, Any] = {"memory_tier": tier}
    if tags:
        obj["tags"] = tags
    return json.dumps(obj, separators=(",", ":"))


def run_domain_index_benchmarks(server: Path) -> list[dict[str, Any]]:
    if len(server.parents) < 3:
        return []
    benchmark = server.parents[2] / "tests" / server.parent.name / "tiny-convmem-search-benchmark.exe"
    if not benchmark.exists():
        return []
    results: list[dict[str, Any]] = []
    for size, iterations in ((1024, 100), (8192, 100), (32768, 50), (1048576, 5)):
        completed = subprocess.run(
            [str(benchmark), str(size), str(iterations)],
            check=True,
            capture_output=True,
            text=True,
        )
        results.append(json.loads(completed.stdout))
    return results


# ---------------------------------------------------------------------------
# クライアント・サーバランチャー
# ---------------------------------------------------------------------------
class ServerContext:
    """MCP stdio サーバを起動し、ClientSession でラップするコンテキスト。"""

    def __init__(self, exe: Path, data_dir: Path) -> None:
        self.exe = exe
        self.data_dir = data_dir
        self.read_stream = None
        self.write_stream = None
        self.session = None

    async def __aenter__(self) -> ClientSession:
        params = StdioServerParameters(
            command=str(self.exe),
            args=["--data-dir", str(self.data_dir)],
            env={**os.environ},
        )
        self._ctx = stdio_client(params)
        self.read_stream, self.write_stream = await self._ctx.__aenter__()
        self.session = ClientSession(self.read_stream, self.write_stream)
        await self.session.__aenter__()
        await self.session.initialize()
        return self.session

    async def __aexit__(self, exc_type, exc, tb) -> None:
        if self.session:
            session = self.session
            self.session = None
            await session.__aexit__(exc_type, exc, tb)
        await self._ctx.__aexit__(exc_type, exc, tb)


async def with_session(exe: Path, data_dir: Path) -> Awaitable[ClientSession]:
    return await ServerContext(exe, data_dir).__aenter__()


# ---------------------------------------------------------------------------
# 操作ラッパー（結果検証付き）
# ---------------------------------------------------------------------------
async def write_message(
    session: ClientSession,
    session_id: str,
    content: str,
    metadata_json: str = '{"memory_tier":"short"}',
) -> dict:
    result = await session.call_tool(
        "write_message",
        {
            "session_id": session_id,
            "role": "user",
            "content": content,
            "metadata_json": metadata_json,
        },
    )
    payload = json_result(result)
    if payload.get("status") != "OK":
        raise RuntimeError(f"write failed: {payload}")
    msg = payload.get("message", {})
    if not msg.get("message_id"):
        raise RuntimeError("write response missing message_id")
    return msg


async def read_messages(session: ClientSession, session_id: str) -> list[dict]:
    result = await session.call_tool("read_messages", {"session_id": session_id})
    payload = json_result(result)
    history = payload.get("history", [])
    if not isinstance(history, list):
        raise RuntimeError("history is not a list")
    return history


async def search_messages(
    session: ClientSession, query: str, session_id: str, limit: int = 10
) -> list[dict]:
    result = await session.call_tool(
        "search_messages",
        {"query": query, "session_id": session_id, "limit": limit},
    )
    payload = json_result(result)
    results = payload.get("results", [])
    if not isinstance(results, list):
        raise RuntimeError("search results is not a list")
    return results


async def delete_message(session: ClientSession, message_id: str, session_id: str) -> str:
    result = await session.call_tool(
        "delete_message",
        {"message_id": message_id, "session_id": session_id},
    )
    payload = json_result(result)
    status = payload.get("status", "")
    if status != "OK":
        raise RuntimeError(f"delete failed: {payload}")
    return status


async def rebuild_index(session: ClientSession) -> dict:
    result = await session.call_tool("rebuild_index", {})
    return json_result(result)


async def get_system_health(session: ClientSession) -> dict:
    result = await session.call_tool("get_system_health", {})
    payload = json_result(result)
    if payload.get("status") != "OK":
        raise RuntimeError(f"health failed: {payload}")
    return payload


async def list_tools(session: ClientSession) -> list[dict]:
    result = await session.list_tools()
    return [t.model_dump() for t in result.tools]


# ---------------------------------------------------------------------------
# シナリオ実装
# ---------------------------------------------------------------------------
class ScenarioRunner:
    def __init__(self, exe: Path) -> None:
        self.exe = exe

    def new_data_dir(self) -> Path:
        return Path(tempfile.mkdtemp(prefix="tiny-convmem-bench-"))

    def cleanup(self, data_dir: Path) -> None:
        shutil.rmtree(data_dir, ignore_errors=True)

    async def measure(
        self,
        name: str,
        fn: Callable[[int], Awaitable[Any]],
        iterations: int,
        bytes_in: int = 0,
        bytes_out: int = 0,
    ) -> ScenarioResult:
        samples: list[LatencySample] = []
        for i in range(iterations):
            start = time.perf_counter()
            error: str | None = None
            success = False
            try:
                await fn(i)
                success = True
            except Exception as exc:  # noqa: BLE001
                error = f"{type(exc).__name__}: {exc}"
            elapsed_ms = (time.perf_counter() - start) * 1000.0
            samples.append(
                LatencySample(
                    operation=name,
                    duration_ms=elapsed_ms,
                    success=success,
                    error=error,
                    bytes_in=bytes_in,
                    bytes_out=bytes_out,
                )
            )
        return summarize(name, samples)

    async def run_write(
        self, content: str, metadata_json: str, iterations: int, name: str
    ) -> ScenarioResult:
        data_dir = self.new_data_dir()
        session_id = "bench"
        async with ServerContext(self.exe, data_dir) as session:

            async def _run(_: int) -> None:
                await write_message(session, session_id, content, metadata_json)

            try:
                return await self.measure(
                    name,
                    _run,
                    iterations,
                    bytes_in=len(content.encode("utf-8")),
                )
            finally:
                self.cleanup(data_dir)

    async def run_write_small(self, iterations: int) -> ScenarioResult:
        content = "hello durable memory"
        return await self.run_write(
            content,
            make_metadata_json(),
            iterations,
            "write_message (small ~20B)",
        )

    async def run_write_medium(self, iterations: int) -> ScenarioResult:
        content = random_content(DEFAULT_CONTENT_MEDIUM)
        return await self.run_write(
            content,
            make_metadata_json(),
            iterations,
            "write_message (medium 1KiB)",
        )

    async def run_write_large(self, iterations: int) -> ScenarioResult:
        content = random_content(DEFAULT_CONTENT_LARGE)
        return await self.run_write(
            content,
            make_metadata_json(),
            iterations,
            "write_message (large 32KiB)",
        )

    async def run_write_8k(self, iterations: int) -> ScenarioResult:
        content = random_content(DEFAULT_CONTENT_8K)
        return await self.run_write(
            content,
            make_metadata_json(),
            iterations,
            "write_message (P1 8KiB)",
        )

    async def run_write_max(self, iterations: int) -> ScenarioResult:
        content = random_content(DEFAULT_CONTENT_MAX)
        return await self.run_write(
            content,
            make_metadata_json(),
            iterations,
            "write_message (P1 1MiB)",
        )

    async def run_p1_startup_rebuild(
        self, size: int
    ) -> tuple[ScenarioResult, ScenarioResult]:
        data_dir = self.new_data_dir()
        session_id = f"p1-{size}"
        content = random_content(size)
        try:
            async with ServerContext(self.exe, data_dir) as session:
                for _ in range(P1_CORPUS_MESSAGES):
                    await write_message(session, session_id, content)

            started = time.perf_counter()
            async with ServerContext(self.exe, data_dir) as session:
                startup_ms = (time.perf_counter() - started) * 1000.0
                health = await get_system_health(session)
                if health.get("active_messages") != P1_CORPUS_MESSAGES:
                    raise RuntimeError(f"P1 corpus count mismatch: {health}")

                rebuild_started = time.perf_counter()
                payload = await rebuild_index(session)
                rebuild_ms = (time.perf_counter() - rebuild_started) * 1000.0
                if payload.get("status") != "OK":
                    raise RuntimeError(f"P1 rebuild failed: {payload}")

            label = f"{size // 1024}KiB" if size < 1024 * 1024 else "1MiB"
            return (
                single_sample(
                    f"startup index build (P1 {label} x{P1_CORPUS_MESSAGES})",
                    startup_ms,
                    "process start + initialize; corpus is loaded from the persisted log",
                ),
                single_sample(
                    f"rebuild_index (P1 {label} x{P1_CORPUS_MESSAGES})",
                    rebuild_ms,
                    "compaction and complete in-memory search index rebuild",
                ),
            )
        finally:
            self.cleanup(data_dir)

    async def run_write_with_tags(self, iterations: int) -> ScenarioResult:
        content = "tagged message"
        metadata_json = make_metadata_json(tags=["benchmark", "tagged"], tier="short")
        return await self.run_write(
            content,
            metadata_json,
            iterations,
            "write_message (with tags)",
        )

    async def run_read(self, iterations: int) -> ScenarioResult:
        data_dir = self.new_data_dir()
        session_id = "bench"
        async with ServerContext(self.exe, data_dir) as session:
            # 事前に 100 件書き込み
            for i in range(100):
                await write_message(session, session_id, f"read target {i}")

            async def _run(_: int) -> None:
                history = await read_messages(session, session_id)
                if len(history) != 100:
                    raise RuntimeError(f"expected 100 messages, got {len(history)}")

            try:
                return await self.measure(
                    "read_messages (100 msgs)",
                    _run,
                    iterations,
                )
            finally:
                self.cleanup(data_dir)

    async def run_search(self, iterations: int) -> tuple[ScenarioResult, ScenarioResult]:
        """同じデータ・同じ query で session 絞り込み有無を比較。"""
        data_dir = self.new_data_dir()
        session_id = "bench"
        other_session = "other"
        async with ServerContext(self.exe, data_dir) as session:
            # target session に 50 件
            for i in range(50):
                await write_message(
                    session,
                    session_id,
                    f"search keyword benchmark {i}",
                    make_metadata_json(tags=["target"]),
                )
            # 別 session に 50 件ノイズ
            for i in range(50):
                await write_message(
                    session,
                    other_session,
                    f"noise keyword other {i}",
                    make_metadata_json(tags=["noise"]),
                )

            async def _global(_: int) -> int:
                results = await search_messages(session, "benchmark", "", limit=10)
                if len(results) != 10:
                    raise RuntimeError(f"expected 10 global results, got {len(results)}")
                return len(results)

            async def _scoped(_: int) -> int:
                results = await search_messages(session, "benchmark", session_id, limit=10)
                if len(results) != 10:
                    raise RuntimeError(f"expected 10 scoped results, got {len(results)}")
                # 全結果が target session であることを検証
                for r in results:
                    if r.get("session_id") != session_id:
                        raise RuntimeError(f"scoped result leaked from {r.get('session_id')}")
                return len(results)

            try:
                global_result = await self.measure(
                    "search_messages (global, no session)",
                    _global,
                    iterations,
                )
                scoped_result = await self.measure(
                    "search_messages (scoped to session)",
                    _scoped,
                    iterations,
                )
                return global_result, scoped_result
            finally:
                self.cleanup(data_dir)

    async def run_delete(self, iterations: int) -> ScenarioResult:
        data_dir = self.new_data_dir()
        session_id = "bench"
        async with ServerContext(self.exe, data_dir) as session:
            ids: list[str] = []
            for _ in range(iterations):
                msg = await write_message(session, session_id, "to be deleted")
                ids.append(msg["message_id"])

            async def _run(i: int) -> None:
                await delete_message(session, ids[i], session_id)

            try:
                return await self.measure("delete_message", _run, iterations)
            finally:
                self.cleanup(data_dir)

    async def run_rebuild(self, iterations: int) -> ScenarioResult:
        data_dir = self.new_data_dir()
        session_id = "bench"
        async with ServerContext(self.exe, data_dir) as session:
            # 削除を含めた compaction 対象を 50 件作る
            for _ in range(50):
                msg = await write_message(session, session_id, "will compact")
                await delete_message(session, msg["message_id"], session_id)
            # 追加で 50 件 active
            for i in range(50):
                await write_message(session, session_id, f"active {i}")

            async def _run(_: int) -> None:
                payload = await rebuild_index(session)
                if payload.get("status") != "OK":
                    raise RuntimeError(f"rebuild failed: {payload}")

            try:
                return await self.measure("rebuild_index", _run, iterations)
            finally:
                self.cleanup(data_dir)

    async def run_health(self, iterations: int) -> ScenarioResult:
        data_dir = self.new_data_dir()
        session_id = "bench"
        async with ServerContext(self.exe, data_dir) as session:
            for _ in range(10):
                await write_message(session, session_id, "health check")

            async def _run(_: int) -> None:
                await get_system_health(session)

            try:
                return await self.measure("get_system_health", _run, iterations)
            finally:
                self.cleanup(data_dir)

    async def run_recovery(
        self, iterations: int
    ) -> tuple[ScenarioResult, ScenarioResult, ScenarioResult]:
        data_dir = self.new_data_dir()
        session_id = "bench"
        # 1. データ投入 & 終了
        async with ServerContext(self.exe, data_dir) as session:
            for i in range(100):
                await write_message(session, session_id, f"recovery marker {i}")

        # 2. 復旧確認（メッセージ件数のみ）
        start = time.perf_counter()
        async with ServerContext(self.exe, data_dir) as session:
            history = await read_messages(session, session_id)
            if len(history) != 100:
                raise RuntimeError(f"recovery count mismatch: {len(history)} != 100")
        recovery_ms = (time.perf_counter() - start) * 1000.0
        recovery_result = ScenarioResult(
            name="recovery (restart + count check)",
            iterations=1,
            successes=1,
            failures=0,
            samples=[recovery_ms],
            latency_min_ms=recovery_ms,
            latency_max_ms=recovery_ms,
            latency_mean_ms=recovery_ms,
            latency_p50_ms=recovery_ms,
            latency_p95_ms=recovery_ms,
            latency_p99_ms=recovery_ms,
            total_duration_ms=recovery_ms,
            throughput_ops_per_sec=1000.0 / recovery_ms if recovery_ms > 0 else 0.0,
            notes="includes process shutdown + startup + initialize + read count",
        )

        # 3. 再起動後の small read を iterations 回
        async def _small(_: int) -> None:
            async with ServerContext(self.exe, data_dir) as session:
                history = await read_messages(session, session_id)
                if len(history) != 100:
                    raise RuntimeError(f"small read count mismatch: {len(history)}")

        small_result = await self.measure(
            "recovery small read (1 message check)",
            _small,
            iterations,
        )

        # 4. 再起動後の full read を iterations 回
        async def _full(_: int) -> None:
            async with ServerContext(self.exe, data_dir) as session:
                history = await read_messages(session, session_id)
                if len(history) != 100:
                    raise RuntimeError(f"full read count mismatch: {len(history)}")

        full_result = await self.measure(
            "recovery full read (100 msgs)",
            _full,
            iterations,
        )

        self.cleanup(data_dir)
        return recovery_result, small_result, full_result

    async def run_schema_and_error_checks(self) -> dict[str, Any]:
        data_dir = self.new_data_dir()
        try:
            async with ServerContext(self.exe, data_dir) as session:
                tools = await list_tools(session)
                tool_names = [t.get("name") for t in tools]
                write_schema = next((t for t in tools if t.get("name") == "write_message"), {})
                write_props = write_schema.get("inputSchema", {}).get("properties", {})

                return {
                    "tool_count": len(tools),
                    "tool_names": tool_names,
                    "write_message_properties": list(write_props.keys()),
                    "write_requires": write_schema.get("inputSchema", {}).get("required", []),
                }
        finally:
            self.cleanup(data_dir)

    def check_jsonrpc_error_response(self) -> dict[str, Any]:
        """subprocess で raw stdio し、JSON-RPC エラーレスポンスを検証する。"""
        data_dir = self.new_data_dir()
        try:
            # 必須引数不足のリクエスト
            request = json.dumps(
                {
                    "jsonrpc": "2.0",
                    "id": 2,
                    "method": "tools/call",
                    "params": {"name": "write_message", "arguments": {}},
                },
                separators=(",", ":"),
            )
            proc = subprocess.Popen(
                [str(self.exe), "--data-dir", str(data_dir)],
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
            )
            stdout, stderr = proc.communicate(input=request + "\n", timeout=10)
            lines = [line.strip() for line in stdout.splitlines() if line.strip()]
            if not lines:
                return {"error_response_present": False, "stderr": stderr[:500]}
            try:
                response = json.loads(lines[-1])
            except json.JSONDecodeError as exc:
                return {
                    "error_response_present": False,
                    "raw_last_line": lines[-1][:200],
                    "parse_error": str(exc),
                }
            has_jsonrpc_field = "jsonrpc" in response
            has_error_field = "error" in response
            return {
                "error_response_present": has_error_field,
                "has_jsonrpc_field": has_jsonrpc_field,
                "has_error_field": has_error_field,
                "error_code": response.get("error", {}).get("code") if has_error_field else None,
                "raw_response": response,
            }
        except Exception as exc:  # noqa: BLE001
            return {"error_response_present": False, "exception": str(exc)}
        finally:
            try:
                proc.kill()
            except Exception:  # noqa: BLE001
                pass
            self.cleanup(data_dir)


# ---------------------------------------------------------------------------
# レポート生成
# ---------------------------------------------------------------------------
def print_report(report: BenchmarkReport) -> None:
    print(f"\n=== tiny-convmem-mcp 負荷テスト結果 ===")
    print(f"server      : {report.server}")
    print(f"build_type  : {report.build_type}")
    print(f"exe_mtime   : {report.exe_mtime}")
    print(f"iterations  : {report.iterations}")
    print(f"warmup      : {report.warmup}")
    print()
    header = (
        f"{'scenario':<45} {'succ':>5} {'fail':>5} "
        f"{'p50[ms]':>10} {'p95[ms]':>10} {'p99[ms]':>10} "
        f"{'mean[ms]':>10} {'ops/s':>10}"
    )
    print(header)
    print("-" * len(header))
    for s in report.scenarios:
        print(
            f"{s.name:<45} {s.successes:>5} {s.failures:>5} "
            f"{s.latency_p50_ms:>10.3f} {s.latency_p95_ms:>10.3f} {s.latency_p99_ms:>10.3f} "
            f"{s.latency_mean_ms:>10.3f} {s.throughput_ops_per_sec:>10.2f}"
        )
    if report.errors:
        print("\nerrors:")
        for e in report.errors:
            print(f"  - {e}")
    if report.schema_check:
        print("\nschema / error checks:")
        for k, v in report.schema_check.items():
            print(f"  {k}: {v}")
    if report.health_after:
        print("\nfinal health:")
        for k, v in report.health_after.items():
            print(f"  {k}: {v}")
    if report.domain_index:
        print("\ndomain search index:")
        for item in report.domain_index:
            print(
                f"  {item['size_bytes']:>7} bytes: {item['mean_ms']:.3f} ms, "
                f"peak working set {item['peak_working_set_bytes']} bytes"
            )


def write_json(report: BenchmarkReport, path: Path) -> None:
    path.write_text(json.dumps(report.as_dict(), indent=2, ensure_ascii=False), encoding="utf-8")


def write_markdown(report: BenchmarkReport, json_path: Path, md_path: Path) -> None:
    lines: list[str] = []
    lines.append("# tiny-convmem-mcp 負荷テストレポート\n")
    lines.append(f"- サーバ: `{report.server}`\n")
    lines.append(f"- ビルド: {report.build_type}\n")
    lines.append(f"- exe mtime: {report.exe_mtime}\n")
    lines.append(f"- iterations: {report.iterations}\n")
    lines.append(f"- warmup: {report.warmup}\n")
    lines.append(f"- 計測日: {time.strftime('%Y-%m-%d %H:%M:%S')}\n")
    lines.append(f"- raw JSON: `{json_path.name}`\n\n")

    lines.append("## レイテンシサマリ\n\n")
    lines.append(
        "| scenario | succ | fail | p50[ms] | p95[ms] | p99[ms] | mean[ms] | ops/s |\n"
    )
    lines.append("|---|---|---|---|---|---|---|---|\n")
    for s in report.scenarios:
        lines.append(
            f"| {s.name} | {s.successes} | {s.failures} | "
            f"{s.latency_p50_ms:.3f} | {s.latency_p95_ms:.3f} | {s.latency_p99_ms:.3f} | "
            f"{s.latency_mean_ms:.3f} | {s.throughput_ops_per_sec:.2f} |\n"
        )
    lines.append("\n")

    if report.schema_check:
        lines.append("## MCP 準拠 / エラーチェック\n\n")
        lines.append("```json\n")
        lines.append(json.dumps(report.schema_check, indent=2, ensure_ascii=False))
        lines.append("\n```\n\n")

    if report.domain_index:
        lines.append("## Domain検索index構築\n\n")
        lines.append("| payload | iterations | mean[ms] | peak working set[bytes] |\n")
        lines.append("|---:|---:|---:|---:|\n")
        for item in report.domain_index:
            lines.append(
                f"| {item['size_bytes']} | {item['iterations']} | {item['mean_ms']:.3f} | "
                f"{item['peak_working_set_bytes']} |\n"
            )
        lines.append("\n")

    if report.health_after:
        lines.append("## 最終 health\n\n")
        lines.append("```json\n")
        lines.append(json.dumps(report.health_after, indent=2, ensure_ascii=False))
        lines.append("\n```\n\n")

    if report.errors:
        lines.append("## エラー\n\n")
        for e in report.errors:
            lines.append(f"- {e}\n")

    md_path.write_text("".join(lines), encoding="utf-8")


# ---------------------------------------------------------------------------
# コマンドライン
# ---------------------------------------------------------------------------
def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="tiny-convmem-mcp ユーザ視点負荷テスト＆ベンチマーク"
    )
    parser.add_argument(
        "--server",
        type=Path,
        default=DEFAULT_RELEASE_EXE if DEFAULT_RELEASE_EXE.exists() else None,
        help="tiny-convmem-mcp.exe のパス（デフォルト: Release ビルド）",
    )
    parser.add_argument(
        "--iterations",
        type=int,
        default=DEFAULT_ITERATIONS,
        help=f"各シナリオの反復回数（デフォルト {DEFAULT_ITERATIONS}）",
    )
    parser.add_argument(
        "--warmup",
        type=int,
        default=WARMUP_COUNT,
        help=f"ウォームアップ書き込み数（デフォルト {WARMUP_COUNT}）",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("benchmark_results.json"),
        help="出力 JSON ファイル名（デフォルト: benchmark_results.json）",
    )
    return parser.parse_args(argv)


async def amain(argv: list[str]) -> int:
    args = parse_args(argv)

    exe = args.server
    if not exe or not exe.exists():
        print(f"サーバ exe が見つかりません: {exe}")
        return 1

    exe_mtime = time.strftime(
        "%Y-%m-%d %H:%M:%S", time.localtime(exe.stat().st_mtime)
    )

    print(f"server : {exe}")
    print(f"mtime  : {exe_mtime}")
    print(f"iters  : {args.iterations}")

    runner = ScenarioRunner(exe)
    report = BenchmarkReport(
        server=str(exe),
        build_type="release" if "Release" in str(exe) else "debug",
        exe_mtime=exe_mtime,
        iterations=args.iterations,
        warmup=args.warmup,
    )
    report.domain_index = run_domain_index_benchmarks(exe.resolve())

    # シナリオ実行
    report.scenarios.append(await runner.run_write_small(args.iterations))
    report.scenarios.append(await runner.run_write_medium(args.iterations))
    report.scenarios.append(await runner.run_write_8k(min(args.iterations, 100)))
    report.scenarios.append(await runner.run_write_large(args.iterations))
    report.scenarios.append(await runner.run_write_max(min(args.iterations, 5)))
    report.scenarios.append(await runner.run_write_with_tags(args.iterations))

    for payload_size in (
        DEFAULT_CONTENT_MEDIUM,
        DEFAULT_CONTENT_8K,
        DEFAULT_CONTENT_LARGE,
        DEFAULT_CONTENT_MAX,
    ):
        startup_result, p1_rebuild_result = await runner.run_p1_startup_rebuild(payload_size)
        report.scenarios.extend((startup_result, p1_rebuild_result))

    report.scenarios.append(await runner.run_read(args.iterations))

    search_global, search_scoped = await runner.run_search(args.iterations)
    report.scenarios.append(search_global)
    report.scenarios.append(search_scoped)

    report.scenarios.append(await runner.run_delete(args.iterations))

    # rebuild は重いので回数を少なく
    rebuild_iters = min(args.iterations, REBUILD_ITERATIONS)
    report.scenarios.append(await runner.run_rebuild(rebuild_iters))

    report.scenarios.append(await runner.run_health(args.iterations))

    # recovery はプロセス起動/終了を伴うため回数を少なく
    recovery_iters = min(args.iterations, REBUILD_ITERATIONS)
    recovery_results = await runner.run_recovery(recovery_iters)
    report.scenarios.extend(recovery_results)

    # schema / error checks
    try:
        schema_info = await runner.run_schema_and_error_checks()
        jsonrpc_err = runner.check_jsonrpc_error_response()
        report.schema_check = {
            **schema_info,
            "jsonrpc_error_response": jsonrpc_err,
        }
    except Exception as exc:  # noqa: BLE001
        report.errors.append(f"schema check failed: {exc}")

    # 最終 health（独立したディレクトリ）
    try:
        data_dir = runner.new_data_dir()
        async with ServerContext(exe, data_dir) as session:
            for _ in range(10):
                await write_message(session, "final", "final check")
            report.health_after = await get_system_health(session)
        runner.cleanup(data_dir)
    except Exception as exc:  # noqa: BLE001
        report.errors.append(f"final health failed: {exc}")

    print_report(report)

    json_path = args.output.resolve()
    md_path = json_path.with_suffix(".md")
    write_json(report, json_path)
    write_markdown(report, json_path, md_path)
    print(f"\nJSON report: {json_path}")
    print(f"MD report  : {md_path}")

    return 0


def main() -> int:
    return asyncio.run(amain(sys.argv[1:]))


if __name__ == "__main__":
    sys.exit(main())