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

import argparse
import json
import shutil
import subprocess
import tempfile
import threading
from collections.abc import Mapping
from pathlib import Path
from typing import Any


class Server:
    def __init__(self, executable: Path, data_dir: Path) -> None:
        self.stderr: list[str] = []
        self.process = subprocess.Popen(
            [str(executable), "--data-dir", str(data_dir)],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            encoding="utf-8",
            bufsize=1,
        )
        assert self.process.stdin and self.process.stdout and self.process.stderr
        self._stderr_thread = threading.Thread(target=self._drain_stderr, daemon=True)
        self._stderr_thread.start()

    def _drain_stderr(self) -> None:
        assert self.process.stderr
        self.stderr.extend(self.process.stderr)

    def send(self, message: Mapping[str, Any] | str) -> dict[str, Any]:
        assert self.process.stdin and self.process.stdout
        line = message if isinstance(message, str) else json.dumps(message, separators=(",", ":"))
        self.process.stdin.write(line + "\n")
        self.process.stdin.flush()
        response_line = self.process.stdout.readline()
        if not response_line:
            raise AssertionError(
                f"server closed: returncode={self.process.poll()} stderr={''.join(self.stderr)[-4000:]}"
            )
        response = json.loads(response_line)
        assert response.get("jsonrpc") == "2.0", response
        return response

    def notify(self, message: Mapping[str, Any]) -> None:
        assert self.process.stdin
        self.process.stdin.write(json.dumps(message, separators=(",", ":")) + "\n")
        self.process.stdin.flush()

    def close(self) -> None:
        if self.process.stdin:
            self.process.stdin.close()
        returncode = self.process.wait(timeout=15)
        self._stderr_thread.join(timeout=2)
        assert returncode == 0, f"returncode={returncode} stderr={''.join(self.stderr)}"
        stderr = "".join(self.stderr)
        assert "AddressSanitizer" not in stderr and "ERROR:" not in stderr, stderr


def request(request_id: int, method: str, params: Mapping[str, Any] | None = None) -> dict[str, Any]:
    message: dict[str, Any] = {"jsonrpc": "2.0", "id": request_id, "method": method}
    if params is not None:
        message["params"] = params
    return message


def tool_call(server: Server, request_id: int, name: str, arguments: Mapping[str, Any]) -> dict[str, Any]:
    response = server.send(
        request(request_id, "tools/call", {"name": name, "arguments": arguments})
    )
    assert response.get("id") == request_id, response
    return response


def tool_payload(response: Mapping[str, Any], expect_error: bool = False) -> dict[str, Any]:
    result = response["result"]
    assert bool(result.get("isError", False)) is expect_error, response
    return json.loads(result["content"][0]["text"])


def assert_stable_recency_order(messages: list[dict[str, Any]]) -> None:
    for left, right in zip(messages, messages[1:]):
        left_time = int(left["timestamp"])
        right_time = int(right["timestamp"])
        assert left_time > right_time or (
            left_time == right_time and left["message_id"] <= right["message_id"]
        ), (left, right)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--server", type=Path, required=True)
    parser.add_argument("--iterations", type=int, default=1000)
    args = parser.parse_args()
    data_dir = Path(tempfile.mkdtemp(prefix="tiny-convmem-p0-"))
    server = Server(args.server.resolve(), data_dir)
    next_id = 1
    try:
        response = server.send(
            request(
                next_id,
                "initialize",
                {
                    "protocolVersion": "2024-11-05",
                    "capabilities": {},
                    "clientInfo": {"name": "p0-regression", "version": "1"},
                },
            )
        )
        assert response["result"]["protocolVersion"] == "2024-11-05"
        next_id += 1
        server.notify({"jsonrpc": "2.0", "method": "notifications/initialized", "params": {}})

        listed = server.send(request(next_id, "tools/list", {}))
        assert listed["id"] == next_id  # also proves the notification emitted no response
        next_id += 1
        tools = {item["name"]: item for item in listed["result"]["tools"]}
        assert set(tools) == {
            "write_message",
            "read_messages",
            "search_messages",
            "delete_message",
            "rebuild_index",
            "get_system_health",
        }
        assert tools["write_message"]["inputSchema"]["required"] == [
            "session_id",
            "role",
            "content",
        ]
        assert tools["read_messages"]["inputSchema"]["required"] == ["session_id"]
        assert tools["search_messages"]["inputSchema"]["required"] == ["query"]
        assert tools["search_messages"]["inputSchema"]["properties"]["limit"]["type"] == "integer"
        assert tools["delete_message"]["inputSchema"]["required"] == ["message_id"]

        message_ids: list[str] = []
        for index in range(20):
            payload = tool_payload(
                tool_call(
                    server,
                    next_id,
                    "write_message",
                    {"session_id": "target", "role": "user", "content": f"needle target {index}"},
                )
            )
            next_id += 1
            message_ids.append(payload["message"]["message_id"])
        for index in range(10):
            tool_payload(
                tool_call(
                    server,
                    next_id,
                    "write_message",
                    {"session_id": "other", "role": "user", "content": f"needle other {index}"},
                )
            )
            next_id += 1

        history = tool_payload(
            tool_call(server, next_id, "read_messages", {"session_id": "target"})
        )["history"]
        next_id += 1
        assert len(history) == 20

        ranked = tool_payload(
            tool_call(
                server,
                next_id,
                "search_messages",
                {"query": "needle target", "limit": 100},
            )
        )["results"]
        next_id += 1
        assert len(ranked) == 30
        assert all(item["session_id"] == "target" for item in ranked[:20])
        assert all(item["session_id"] == "other" for item in ranked[20:])
        assert_stable_recency_order(ranked[:20])
        assert_stable_recency_order(ranked[20:])
        ranked_ids_before = [item["message_id"] for item in ranked]

        for index in range(args.iterations):
            arguments: dict[str, Any] = {"query": "needle"}
            if index % 2:
                arguments["session_id"] = "target"
            results = tool_payload(
                tool_call(server, next_id, "search_messages", arguments)
            )["results"]
            next_id += 1
            assert len(results) == 20
            if index % 2:
                assert all(item["session_id"] == "target" for item in results)
            if (index + 1) % 100 == 0:
                print(f"completed {index + 1}/{args.iterations} alternating searches", flush=True)

        assert tool_payload(
            tool_call(server, next_id, "search_messages", {"query": "needle", "limit": 0}), True
        )["status"] == "ERROR"
        next_id += 1
        assert tool_payload(
            tool_call(server, next_id, "search_messages", {"query": "needle", "limit": 101}), True
        )["status"] == "ERROR"
        next_id += 1
        fractional = tool_call(
            server, next_id, "search_messages", {"query": "needle", "limit": 1.5}
        )
        assert fractional["error"]["code"] == -32602
        next_id += 1
        assert tool_payload(tool_call(server, next_id, "search_messages", {}), True)["status"] == "ERROR"
        next_id += 1

        assert tool_payload(
            tool_call(server, next_id, "delete_message", {"message_id": message_ids[0]})
        )["status"] == "OK"
        next_id += 1
        tool_payload(tool_call(server, next_id, "rebuild_index", {}))
        next_id += 1
        ranked_after_rebuild = tool_payload(
            tool_call(
                server,
                next_id,
                "search_messages",
                {"query": "needle target", "limit": 100},
            )
        )["results"]
        next_id += 1
        assert [item["message_id"] for item in ranked_after_rebuild] == [
            message_id for message_id in ranked_ids_before if message_id != message_ids[0]
        ]
        assert tool_payload(tool_call(server, next_id, "get_system_health", {}))["status"] == "OK"
        next_id += 1

        assert server.send("{bad}")["error"]["code"] == -32700
        invalid = server.send({"id": next_id, "method": "tools/list"})
        assert invalid["error"]["code"] == -32600
        next_id += 1
        missing_method = server.send({"jsonrpc": "2.0", "id": next_id})
        assert missing_method["error"]["code"] == -32600
        next_id += 1
        null_id = server.send(
            {"jsonrpc": "2.0", "id": None, "method": "tools/list", "params": {}}
        )
        assert null_id["error"]["code"] == -32600 and null_id["id"] is None
        invalid_params = server.send(
            {"jsonrpc": "2.0", "id": next_id, "method": "tools/list", "params": []}
        )
        assert invalid_params["error"]["code"] == -32602
        next_id += 1
        unknown_method = server.send(request(next_id, "unknown/method", {}))
        assert unknown_method["error"]["code"] == -32601
        next_id += 1
        unknown_tool = tool_call(server, next_id, "unknown_tool", {})
        assert unknown_tool["error"]["code"] == -32602
        next_id += 1
        missing_arguments = server.send(
            request(next_id, "tools/call", {"name": "search_messages"})
        )
        assert missing_arguments["error"]["code"] == -32602
    except Exception:
        if server.process.poll() is None:
            server.process.kill()
        server.process.wait(timeout=5)
        raise
    else:
        server.close()
    finally:
        shutil.rmtree(data_dir, ignore_errors=True)


if __name__ == "__main__":
    main()