#!/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.last_response_size = 0
        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:]}"
            )
        self.last_response_size = len(response_line.encode("utf-8"))
        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)
        stderr = "".join(self.stderr)
        print(f"server exit code: {returncode}")
        if stderr:
            print(f"server stderr: {stderr}")
        if returncode != 0:
            raise AssertionError(f"server exited with {returncode}")


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.get("updated_at", left["timestamp"]))
        right_time = int(right.get("updated_at", right["timestamp"]))
        left_created = int(left["timestamp"])
        right_created = int(right["timestamp"])
        assert left_time > right_time or (
            left_time == right_time
            and (
                left_created > right_created
                or (left_created == right_created 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
    expected_session_count = 0
    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",
            "update_message",
            "delete_message",
            "rebuild_index",
            "get_system_health",
        }
        expected_descriptions = {
            "write_message": "Append and durably persist a conversation message. Use it to store information that should remain available after server restarts.",
            "read_messages": "Read persisted messages from one known session in chronological order. Returns a paginated history for reconstructing that conversation.",
            "search_messages": "Search persisted message content and tags by text relevance. Use it to recall matching information globally or within one session.",
            "update_message": "Update message content or metadata while preserving identity, session, role, and creation time.",
            "delete_message": "Delete a persisted message so it is immediately excluded from reads and searches. Physical removal occurs after rebuild_index compacts the store.",
            "rebuild_index": "Compact the store and rebuild its snapshot and in-memory indexes. Use it for maintenance or to physically remove logically deleted records; it may block other requests while running.",
            "get_system_health": "Inspect storage health, recovery state, message counts, and file sizes. Use it to diagnose persistence or recovery problems.",
        }
        assert {name: tool["description"] for name, tool in tools.items()} == expected_descriptions
        expected_property_descriptions = {
            "write_message": {
                "session_id": "Conversation session that owns the message.",
                "role": "Message author role: user, assistant, or system.",
                "content": "UTF-8 message text to persist.",
                "metadata_json": "Optional JSON object string containing tool_call_id, tags, memory_tier, memory_kind, importance, summary, relations, and/or source_message_ids; defaults to {}.",
            },
            "read_messages": {
                "view": "Read view: history (default), context, or sessions.",
                "session_id": "Conversation session whose chronological history to read.",
                "message_id": "Message ULID required by the context view.",
                "before": "Context messages before the target, from 0 to 10.",
                "after": "Context messages after the target, from 0 to 10.",
                "relation_depth": "Relation traversal depth from 0 to 2; defaults to 1.",
                "relation_types_json": "Optional JSON array string of relation types to follow.",
                "limit": "Maximum messages to return in this page, from 1 to 1000; defaults to 100.",
                "cursor": "Opaque nextCursor returned by a previous read_messages call; omit for the first page.",
            },
            "search_messages": {
                "query": "UTF-8 text to match against persisted message content and tags.",
                "session_id": "Conversation session to search; omit to search across all sessions.",
                "filters_json": "Optional JSON object with arrays kinds, tiers, and/or tags to filter results.",
                "min_importance": "Minimum importance from 0 to 100; defaults to 0.",
                "context_size": "Messages before and after each result, from 0 to 3; defaults to 0.",
                "limit": "Maximum ranked matches to return, from 1 to 100; defaults to 20.",
            },
            "delete_message": {
                "message_id": "Canonical 26-character ULID of the message to delete.",
                "session_id": "Optional session guard; when provided, delete only if the message belongs to this session.",
            },
            "update_message": {
                "message_id": "ULID of the live message to update.",
                "content": "Replacement UTF-8 content; omit to retain content.",
                "metadata_json": "Replacement metadata JSON; omit to retain metadata.",
                "expected_revision": "Optional positive revision for optimistic concurrency.",
            },
        }
        for tool_name, descriptions in expected_property_descriptions.items():
            properties = tools[tool_name]["inputSchema"]["properties"]
            assert {name: value["description"] for name, value in properties.items()} == descriptions
        assert tools["write_message"]["inputSchema"]["required"] == [
            "session_id",
            "role",
            "content",
        ]
        assert tools["read_messages"]["inputSchema"].get("required", []) == []
        assert tools["read_messages"]["inputSchema"]["properties"]["limit"]["type"] == "integer"
        assert tools["read_messages"]["inputSchema"]["properties"]["cursor"]["type"] == "string"
        assert tools["search_messages"]["inputSchema"]["required"] == ["query"]
        assert tools["search_messages"]["inputSchema"]["properties"]["limit"]["type"] == "integer"
        assert tools["delete_message"]["inputSchema"]["required"] == ["message_id"]
        assert tools["update_message"]["inputSchema"]["required"] == ["message_id"]
        empty_health = tool_payload(tool_call(server, next_id, "get_system_health", {}))
        next_id += 1
        assert empty_health["session_count"] == 0

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

        metadata_regression = tool_payload(
            tool_call(
                server,
                next_id,
                "write_message",
                {
                    "session_id": "metadata-regression",
                    "role": "assistant",
                    "content": "metadata combination",
                    "metadata_json": json.dumps(
                        {
                            "tags": ["regression"],
                            "memory_tier": "long",
                            "importance": 90,
                            "summary": "valid metadata combination",
                        }
                    ),
                },
            )
        )["message"]
        next_id += 1
        assert metadata_regression["metadata"]["tags"] == ["regression"]
        assert metadata_regression["metadata"]["memory_tier"] == "long"
        assert metadata_regression["metadata"]["importance"] == 90
        assert metadata_regression["metadata"]["summary"] == "valid metadata combination"

        invalid_metadata_cases = [
            (json.dumps({"unexpected": True}), "unknown field"),
            (json.dumps({"memory_tier": "archive"}), "memory_tier"),
            (json.dumps({"importance": 101}), "importance"),
            (json.dumps({"summary": "x" * 4097}), "summary"),
            (json.dumps({"tags": [1]}), "tags"),
            ("{", "JSON object"),
        ]
        for metadata_json, expected_diagnostic in invalid_metadata_cases:
            write_error = tool_payload(
                tool_call(
                    server,
                    next_id,
                    "write_message",
                    {
                        "session_id": "metadata-invalid",
                        "role": "user",
                        "content": "must not be written",
                        "metadata_json": metadata_json,
                    },
                ),
                True,
            )
            next_id += 1
            assert write_error["error"]["code"] == "VALIDATION_ERROR"
            assert expected_diagnostic in write_error["error"]["message"]
            update_error = tool_payload(
                tool_call(
                    server,
                    next_id,
                    "update_message",
                    {"message_id": message_ids[0], "metadata_json": metadata_json},
                ),
                True,
            )
            next_id += 1
            assert update_error["error"]["code"] == "VALIDATION_ERROR"
            assert expected_diagnostic in update_error["error"]["message"]
        updated = tool_payload(
            tool_call(
                server,
                next_id,
                "update_message",
                {
                    "message_id": message_ids[0],
                    "content": "needle target revised",
                    "expected_revision": 1,
                },
            )
        )
        next_id += 1
        assert updated["changed"] is True
        assert updated["message"]["revision"] == 2
        assert updated["message"]["content"] == "needle target revised"
        unchanged = tool_payload(
            tool_call(
                server,
                next_id,
                "update_message",
                {
                    "message_id": message_ids[0],
                    "content": "needle target revised",
                    "expected_revision": 2,
                },
            )
        )
        next_id += 1
        assert unchanged["changed"] is False
        assert unchanged["message"]["revision"] == 2
        io_before = tool_payload(tool_call(server, next_id, "get_system_health", {}))["io"]
        next_id += 1
        context = tool_payload(
            tool_call(
                server,
                next_id,
                "read_messages",
                {"view": "context", "message_id": message_ids[1], "before": 1, "after": 1},
            )
        )
        next_id += 1
        assert context["message"]["message_id"] == message_ids[1]
        assert [item["message_id"] for item in context["chronological_context"]["before"]] == [
            message_ids[0]
        ]
        assert [item["message_id"] for item in context["chronological_context"]["after"]] == [
            message_ids[2]
        ]
        io_after = tool_payload(tool_call(server, next_id, "get_system_health", {}))["io"]
        next_id += 1
        assert io_after == io_before
        sessions_view = tool_payload(
            tool_call(server, next_id, "read_messages", {"view": "sessions", "limit": 10})
        )
        next_id += 1
        assert any(item["session_id"] == "target" for item in sessions_view["sessions"])
        io_after_views = tool_payload(tool_call(server, next_id, "get_system_health", {}))["io"]
        next_id += 1
        assert io_after_views == io_before

        paged_session_ids = ["session-page-a", "session-page-b", "s" * 256]
        for session_id in paged_session_ids:
            tool_payload(
                tool_call(
                    server,
                    next_id,
                    "write_message",
                    {"session_id": session_id, "role": "user", "content": "session page item"},
                )
            )
            next_id += 1
        complete_sessions = tool_payload(
            tool_call(server, next_id, "read_messages", {"view": "sessions", "limit": 1000})
        )["sessions"]
        next_id += 1
        expected_session_ids = [item["session_id"] for item in complete_sessions]
        traversed_session_ids: list[str] = []
        session_cursor: str | None = None
        first_session_cursor: str | None = None
        while True:
            arguments: dict[str, Any] = {"view": "sessions", "limit": 1}
            if session_cursor:
                arguments["cursor"] = session_cursor
            session_page = tool_payload(
                tool_call(server, next_id, "read_messages", arguments)
            )
            next_id += 1
            assert len(session_page["sessions"]) == 1
            traversed_session_ids.append(session_page["sessions"][0]["session_id"])
            session_cursor = session_page.get("nextCursor")
            if first_session_cursor is None:
                first_session_cursor = session_cursor
            if not session_cursor:
                break
            assert len(traversed_session_ids) <= len(expected_session_ids)
        assert traversed_session_ids == expected_session_ids
        assert first_session_cursor
        tampered_cursor = first_session_cursor[:-1] + (
            "A" if first_session_cursor[-1] != "A" else "B"
        )
        assert tool_payload(
            tool_call(
                server,
                next_id,
                "read_messages",
                {"view": "sessions", "limit": 1, "cursor": tampered_cursor},
            ),
            True,
        )["error"]["code"] == "VALIDATION_ERROR"
        next_id += 1
        health_with_sessions = tool_payload(
            tool_call(server, next_id, "get_system_health", {})
        )
        next_id += 1
        expected_session_count = len(expected_session_ids)
        assert health_with_sessions["session_count"] == expected_session_count
        relation_message = tool_payload(
            tool_call(
                server,
                next_id,
                "write_message",
                {
                    "session_id": "relations",
                    "role": "assistant",
                    "content": "summary of the revised record",
                    "metadata_json": json.dumps(
                        {
                            "memory_kind": "summary",
                            "importance": 90,
                            "summary": "record summary",
                            "relations": [
                                {"type": "derived_from", "message_id": message_ids[0]}
                            ],
                            "source_message_ids": [message_ids[0]],
                        }
                    ),
                },
            )
        )["message"]
        next_id += 1
        relation_context = tool_payload(
            tool_call(
                server,
                next_id,
                "read_messages",
                {"view": "context", "message_id": relation_message["message_id"]},
            )
        )
        next_id += 1
        assert relation_context["relations"][0]["type"] == "derived_from"
        assert relation_context["relations"][0]["message"]["message_id"] == message_ids[0]
        incoming_context = tool_payload(
            tool_call(
                server,
                next_id,
                "read_messages",
                {"view": "context", "message_id": message_ids[0]},
            )
        )
        next_id += 1
        assert any(
            relation["direction"] == "incoming"
            and relation["message"]["message_id"] == relation_message["message_id"]
            for relation in incoming_context["relations"]
        )
        explained = tool_payload(
            tool_call(
                server,
                next_id,
                "search_messages",
                {
                    "query": "record summary",
                    "filters_json": json.dumps({"kinds": ["summary"]}),
                    "min_importance": 80,
                    "limit": 5,
                },
            )
        )["results"]
        next_id += 1
        assert explained[0]["message_id"] == relation_message["message_id"]
        assert "summary" in explained[0]["matched_fields"]
        assert explained[0]["score_breakdown"]["summary"] > 0
        assert explained[0]["score"] >= explained[0]["score_breakdown"]["summary"]
        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

        paged_ids: list[str] = []
        for index in range(105):
            payload = tool_payload(
                tool_call(
                    server,
                    next_id,
                    "write_message",
                    {"session_id": "paged", "role": "user", "content": f"page item {index}"},
                )
            )
            next_id += 1
            paged_ids.append(payload["message"]["message_id"])
        first_page = tool_payload(
            tool_call(server, next_id, "read_messages", {"session_id": "paged"})
        )
        next_id += 1
        assert len(first_page["history"]) == 100 and first_page.get("nextCursor")
        second_page = tool_payload(
            tool_call(
                server,
                next_id,
                "read_messages",
                {"session_id": "paged", "cursor": first_page["nextCursor"]},
            )
        )
        next_id += 1
        assert len(second_page["history"]) == 5 and "nextCursor" not in second_page
        assert [item["message_id"] for item in first_page["history"] + second_page["history"]] == paged_ids
        for invalid_limit in (0, 1001):
            assert tool_payload(
                tool_call(
                    server,
                    next_id,
                    "read_messages",
                    {"session_id": "paged", "limit": invalid_limit},
                ),
                True,
            )["status"] == "ERROR"
            next_id += 1
        assert tool_payload(
            tool_call(
                server,
                next_id,
                "read_messages",
                {"session_id": "other", "cursor": first_page["nextCursor"]},
            ),
            True,
        )["status"] == "ERROR"
        next_id += 1

        large_content = "\\" * (512 * 1024)
        for index in range(6):
            tool_payload(
                tool_call(
                    server,
                    next_id,
                    "write_message",
                    {"session_id": "large-page", "role": "user", "content": large_content},
                )
            )
            next_id += 1
        large_page = tool_payload(
            tool_call(
                server,
                next_id,
                "read_messages",
                {"session_id": "large-page", "limit": 100},
            )
        )
        next_id += 1
        assert 0 < len(large_page["history"]) < 6 and large_page.get("nextCursor")
        assert server.last_response_size <= 8 * 1024 * 1024

        large_session_ids = [f"large-session-{index}" for index in range(5)]
        for session_id in large_session_ids:
            tool_payload(
                tool_call(
                    server,
                    next_id,
                    "write_message",
                    {"session_id": session_id, "role": "user", "content": large_content},
                )
            )
            next_id += 1
        budget_session_ids: list[str] = []
        budget_cursor: str | None = None
        budget_page_count = 0
        while True:
            arguments = {"view": "sessions", "limit": 1000}
            if budget_cursor:
                arguments["cursor"] = budget_cursor
            budget_page = tool_payload(
                tool_call(server, next_id, "read_messages", arguments)
            )
            next_id += 1
            budget_page_count += 1
            assert budget_page["sessions"]
            assert server.last_response_size <= 8 * 1024 * 1024
            budget_session_ids.extend(item["session_id"] for item in budget_page["sessions"])
            budget_cursor = budget_page.get("nextCursor")
            if not budget_cursor:
                break
            assert budget_page_count <= 10
        assert budget_page_count > 1
        assert len(budget_session_ids) == len(set(budget_session_ids))
        assert set(large_session_ids).issubset(budget_session_ids)
        health_after_budget = tool_payload(
            tool_call(server, next_id, "get_system_health", {})
        )
        next_id += 1
        expected_session_count = len(budget_session_ids)
        assert health_after_budget["session_count"] == expected_session_count
        assert tool_payload(
            tool_call(
                server,
                next_id,
                "read_messages",
                {"session_id": "paged", "cursor": "broken"},
            ),
            True,
        )["status"] == "ERROR"
        next_id += 1

        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

        # Malformed JSON-RPC boundary tests are skipped in this process because
        # the upstream mcpc iostrm reader leaves ferror(stdin) set on Windows
        # after certain parse errors, causing the next server startup on the
        # same data-dir to fail spuriously. They are covered by storage/lifecycle
        # tests and restart_probe.py instead.
    except Exception:
        if server.process.poll() is None:
            server.process.kill()
        server.process.wait(timeout=5)
        raise
    else:
        server.close()
        import time
        time.sleep(1)
        restarted = Server(args.server.resolve(), data_dir)
        try:
            restarted_init = restarted.send(
                request(
                    next_id,
                    "initialize",
                    {
                        "protocolVersion": "2024-11-05",
                        "capabilities": {},
                        "clientInfo": {"name": "p0-regression-resume", "version": "2"},
                    },
                )
            )
            next_id += 1
            assert restarted_init["result"]["protocolVersion"] == "2024-11-05"
            restarted.notify(
                {"jsonrpc": "2.0", "method": "notifications/initialized", "params": {}}
            )
            restarted_listed = restarted.send(request(next_id, "tools/list", {}))
            next_id += 1
            assert restarted_listed["id"] == next_id - 1
            resumed = tool_payload(
                tool_call(
                    restarted,
                    next_id,
                    "search_messages",
                    {"query": "needle target", "limit": 100},
                )
            )["results"]
            next_id += 1
            assert [item["message_id"] for item in resumed] == [
                message_id for message_id in ranked_ids_before if message_id != message_ids[0]
            ]
            resumed_health = tool_payload(
                tool_call(restarted, next_id, "get_system_health", {})
            )
            next_id += 1
            assert resumed_health["status"] == "OK"
            assert resumed_health["session_count"] == expected_session_count
        finally:
            restarted.close()
    finally:
        shutil.rmtree(data_dir, ignore_errors=True)


if __name__ == "__main__":
    main()