"""Lab 01 — MCP server/client and the bank's tool estate.

Two halves, built together:

  * the PROTOCOL — JSON-RPC 2.0, capability negotiation, tools/list, tools/call,
    resources, prompts, change notifications;
  * the ESTATE — versioning, side-effect classes, required scopes, data classification,
    and discovery that is AUTHORIZATION-AWARE.

MCP solved the M x N integration problem. It did not solve authorization, versioning or
governance — that is the platform's job, and it is the second half of this file.

    pytest test_lab.py -v
    LAB_MODULE=solution pytest test_lab.py -v    # the reference, must be green
"""

from __future__ import annotations

import json
import re
from dataclasses import dataclass
from enum import Enum
from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple

PROTOCOL_VERSION = "2025-06-18"

# ======================================================================================
# 1. JSON-RPC 2.0
# ======================================================================================

PARSE_ERROR = -32700
INVALID_REQUEST = -32600
METHOD_NOT_FOUND = -32601
INVALID_PARAMS = -32602
INTERNAL_ERROR = -32603


class JsonRpcError(Exception):
    """A PROTOCOL failure. Distinct from a TOOL failure, which is a successful response
    carrying ``isError: true``."""

    def __init__(self, code: int, message: str, data: Any = None) -> None:
        super().__init__(message)
        self.code = code
        self.message = message
        self.data = data

    def to_payload(self) -> Dict[str, Any]:
        error: Dict[str, Any] = {"code": self.code, "message": self.message}
        if self.data is not None:
            error["data"] = self.data
        return error


def make_request(request_id: int, method: str, params: Optional[Mapping[str, Any]] = None) -> Dict[str, Any]:
    """``{"jsonrpc": "2.0", "id": ..., "method": ...}`` plus ``params`` only if given."""
    # TODO
    raise NotImplementedError


def make_notification(method: str, params: Optional[Mapping[str, Any]] = None) -> Dict[str, Any]:
    """Same, but with NO id. A notification must never be answered."""
    # TODO
    raise NotImplementedError


def validate_envelope(message: Mapping[str, Any]) -> None:
    """Raise JsonRpcError(INVALID_REQUEST) unless:
    jsonrpc == "2.0"; method is a non-empty string; id (if present) is int or str
    (NOT a float); params (if present) is an object.
    """
    # TODO
    raise NotImplementedError


# ======================================================================================
# 2. A JSON Schema subset — the contract enforcer
# ======================================================================================


@dataclass(frozen=True)
class ValidationError:
    path: str
    message: str

    def __str__(self) -> str:
        return f"{self.path}: {self.message}"


_TYPE_CHECKS: Mapping[str, Callable[[Any], bool]] = {
    "object": lambda v: isinstance(v, dict),
    "array": lambda v: isinstance(v, list),
    "string": lambda v: isinstance(v, str),
    # bool subclasses int in Python; "integer" must reject True.
    "integer": lambda v: isinstance(v, int) and not isinstance(v, bool),
    "number": lambda v: isinstance(v, (int, float)) and not isinstance(v, bool),
    "boolean": lambda v: isinstance(v, bool),
    "null": lambda v: v is None,
}


def validate_schema(value: Any, schema: Mapping[str, Any], path: str = "$") -> List[ValidationError]:
    """Validate against a subset of JSON Schema, returning ALL errors sorted by
    ``(path, message)`` so a repair loop can fix everything in one turn.

    Support: type (single or list), const, enum, required, properties,
    additionalProperties=False, items, minimum/maximum/exclusiveMinimum/exclusiveMaximum,
    minLength/maxLength/pattern, minItems/maxItems/uniqueItems.

    Paths look like ``$.a.b`` and ``$.legs[1]``.
    An unsupported ``type`` value is a PROGRAMMING error -> raise ValueError.
    """
    # TODO: call a recursive `_validate` helper, then sort
    raise NotImplementedError


def _validate(value: Any, schema: Mapping[str, Any], path: str, errors: List[ValidationError]) -> None:
    """Recursive worker.

    Important: once the TYPE check fails, RETURN — every further check on that value
    would be noise the model has to wade through.
    """
    # TODO
    raise NotImplementedError


def _type_name(value: Any) -> str:
    """JSON type name for an error message. bool before int; int is "integer"."""
    # TODO
    raise NotImplementedError


# ======================================================================================
# 3. Semantic versioning
# ======================================================================================


@dataclass(frozen=True, order=True)
class Version:
    major: int
    minor: int
    patch: int

    @classmethod
    def parse(cls, text: str) -> "Version":
        """Strict ``X.Y.Z``. Anything else -> ValueError."""
        # TODO
        raise NotImplementedError

    def __str__(self) -> str:
        return f"{self.major}.{self.minor}.{self.patch}"

    def satisfies(self, constraint: str) -> bool:
        """``*`` (any) · ``X.Y.Z`` (exact) · ``^X.Y.Z`` (same major, >= minor.patch) ·
        ``~X.Y.Z`` (same major.minor, >= patch)."""
        # TODO
        raise NotImplementedError


class SchemaChange(str, Enum):
    PATCH = "patch"
    MINOR = "minor"
    MAJOR = "major"


def classify_schema_change(old: Mapping[str, Any], new: Mapping[str, Any]) -> SchemaChange:
    """Decide the release type. The rule: callers break when the contract gets STRICTER.

    MAJOR: a new required property; a removed property; a changed type; a NARROWED enum;
           a tightened bound (minimum/minLength/minItems raised, maximum/maxLength/
           maxItems lowered); additionalProperties going from open to False.
    MINOR: a new optional property; a REMOVED required property; a WIDENED enum.
    PATCH: anything else (descriptions, titles).

    Check for MAJOR first — a change can be both, and breaking wins.
    """
    # TODO
    raise NotImplementedError


# ======================================================================================
# 4. The tool estate
# ======================================================================================


class SideEffect(str, Enum):
    READ = "read"
    WRITE_IDEMPOTENT = "write_idempotent"
    WRITE_NON_IDEMPOTENT = "write_non_idempotent"
    IRREVERSIBLE = "irreversible"


# TODO: retry policy is derived from the side-effect class BY THE PLATFORM, never chosen
# by an agent author. read and write_idempotent are retryable; the other two are not
# (write_non_idempotent only with an idempotency key — Phase 10).
RETRYABLE: Mapping[SideEffect, bool] = {}


class Lifecycle(str, Enum):
    ACTIVE = "active"
    DEPRECATED = "deprecated"
    RETIRED = "retired"


@dataclass(frozen=True)
class ToolSpec:
    """One version of one tool. Immutable: a change is a new version, never an edit."""

    name: str
    version: Version
    title: str
    description: str          # the model reads this — it is PROMPT SURFACE
    input_schema: Mapping[str, Any]
    side_effect: SideEffect
    required_scopes: Tuple[str, ...] = ()
    data_classification: str = "internal"
    owner: str = "unknown"
    tenants: Tuple[str, ...] = ()      # empty means "all tenants"
    lifecycle: Lifecycle = Lifecycle.ACTIVE
    superseded_by: Optional[str] = None

    @property
    def key(self) -> str:
        return f"{self.name}@{self.version}"

    def to_mcp(self) -> Dict[str, Any]:
        """The wire shape from ``tools/list``: exactly
        ``{name, title, description, inputSchema, _meta{version, sideEffect}}``.

        Platform metadata (scopes, classification, owner, tenants) must NOT appear — it
        wastes context and leaks the shape of your control model into prompt text.
        """
        # TODO
        raise NotImplementedError


@dataclass(frozen=True)
class Principal:
    agent_id: str
    tenant: str
    scopes: Tuple[str, ...] = ()
    max_classification: str = "internal"


_CLASSIFICATION_ORDER = ("public", "internal", "confidential", "restricted")


def classification_rank(name: str) -> int:
    """Index in _CLASSIFICATION_ORDER. Unknown -> ValueError."""
    # TODO
    raise NotImplementedError


class ToolRegistry:
    def __init__(self) -> None:
        self._by_name: Dict[str, Dict[Version, ToolSpec]] = {}

    def publish(self, spec: ToolSpec) -> None:
        """Publish a new version.

        Refuse (ValueError) when: this (name, version) already exists; the version is
        not newer than the current latest; or the version bump is smaller than the
        schema change requires (a MAJOR change must bump major to previous.major + 1;
        a MINOR change must keep major and raise minor).
        """
        # TODO
        raise NotImplementedError

    def deprecate(self, name: str, version: str, *, superseded_by: Optional[str] = None) -> ToolSpec:
        # TODO
        raise NotImplementedError

    def retire(self, name: str, version: str) -> ToolSpec:
        # TODO
        raise NotImplementedError

    def _set_lifecycle(self, name, version, lifecycle, superseded_by) -> ToolSpec:
        """Replace the stored spec with a copy carrying the new lifecycle.
        Missing version -> KeyError."""
        # TODO
        raise NotImplementedError

    def versions(self, name: str) -> List[ToolSpec]:
        """All versions, ascending. Unknown name -> []."""
        # TODO
        raise NotImplementedError

    def latest(self, name: str) -> Optional[ToolSpec]:
        """Latest ACTIVE version, or None."""
        # TODO
        raise NotImplementedError

    def resolve(self, name: str, constraint: str = "*") -> Optional[ToolSpec]:
        """Highest NON-RETIRED version satisfying the constraint (deprecated counts)."""
        # TODO
        raise NotImplementedError

    def discover(self, principal: Principal, *, include_deprecated: bool = False) -> List[ToolSpec]:
        """What this principal may see RIGHT NOW — the newest visible version of each
        tool, ordered by name.

        Exclude: retired; deprecated (unless asked); tools restricted to other tenants;
        tools above the principal's classification; tools whose required_scopes are not
        a subset of the principal's scopes.

        Filter HERE, not after ranking and not in the client. A tool the principal may
        not call must not appear at all: its name and description are information, and a
        model that can see a tool will eventually call it.
        """
        # TODO
        raise NotImplementedError


# ======================================================================================
# 5. The MCP server
# ======================================================================================


@dataclass(frozen=True)
class ToolCallResult:
    content: str
    is_error: bool = False


ToolHandler = Callable[[Mapping[str, Any]], ToolCallResult]


@dataclass(frozen=True)
class Resource:
    uri: str
    name: str
    mime_type: str
    text: str
    tenants: Tuple[str, ...] = ()


@dataclass(frozen=True)
class Prompt:
    name: str
    description: str
    arguments: Tuple[str, ...]
    template: str


class MCPServer:
    def __init__(
        self,
        *,
        name: str,
        version: str,
        registry: ToolRegistry,
        handlers: Mapping[str, ToolHandler],
        principal: Principal,
        resources: Sequence[Resource] = (),
        prompts: Sequence[Prompt] = (),
        protocol_versions: Sequence[str] = (PROTOCOL_VERSION,),
    ) -> None:
        self.name = name
        self.version = version
        self.registry = registry
        self.handlers = dict(handlers)
        self.principal = principal
        self.resources = {r.uri: r for r in resources}
        self.prompts = {p.name: p for p in prompts}
        self.protocol_versions = list(protocol_versions)
        self.initialized = False
        self.negotiated_version: Optional[str] = None
        self.outbound: List[Dict[str, Any]] = []
        self.call_log: List[Tuple[str, Mapping[str, Any]]] = []

    def handle(self, message: Mapping[str, Any]) -> Optional[Dict[str, Any]]:
        """Handle one message; return a response, or None for a notification.

        Order: validate the envelope (envelope errors answer with the message's id, even
        if absent) → if the method starts with "notifications/", handle and return None →
        if there is no id, return None → if not initialized and the method is not
        "initialize", INVALID_REQUEST → dispatch → wrap in a result or an error.
        """
        # TODO
        raise NotImplementedError

    def _dispatch(self, method: str, params: Mapping[str, Any]) -> Dict[str, Any]:
        """Route to _initialize / ping / tools/list / tools/call / resources/list /
        resources/read / prompts/list / prompts/get. Unknown -> METHOD_NOT_FOUND."""
        # TODO
        raise NotImplementedError

    def _handle_notification(self, method: str, params: Mapping[str, Any]) -> None:
        """``notifications/initialized`` sets ``self.initialized``. Ignore others."""
        # TODO
        raise NotImplementedError

    def _initialize(self, params: Mapping[str, Any]) -> Dict[str, Any]:
        """Negotiate. ``protocolVersion`` is required (else INVALID_PARAMS).

        If the requested version is supported, honour it; OTHERWISE offer our newest
        rather than failing — a version skew must not become an outage.

        Return ``{protocolVersion, capabilities{tools{listChanged:True},
        resources{subscribe:False,listChanged:False}, prompts{listChanged:False}},
        serverInfo{name, version}}``.
        """
        # TODO
        raise NotImplementedError

    def _tools_list(self, params: Mapping[str, Any]) -> Dict[str, Any]:
        """``{"tools": [spec.to_mcp() for discoverable specs]}``."""
        # TODO
        raise NotImplementedError

    def _tools_call(self, params: Mapping[str, Any]) -> Dict[str, Any]:
        """Validate params, resolve against DISCOVERABLE tools, validate arguments
        against the schema, then dispatch.

        - name missing/not a string, or arguments not an object -> INVALID_PARAMS
        - not discoverable -> METHOD_NOT_FOUND, with the SAME message shape as a tool
          that does not exist (a caller must not be able to probe the estate)
        - schema errors -> INVALID_PARAMS with ``data={"errors": [...]}``; the handler is
          NEVER called
        - no handler registered -> INTERNAL_ERROR
        - success -> append to ``call_log`` and return
          ``{"content": [{"type": "text", "text": ...}], "isError": bool}``

        A tool FAILURE is a successful response with isError=True, so the model can see
        it and correct itself.
        """
        # TODO
        raise NotImplementedError

    def _resources_list(self, params: Mapping[str, Any]) -> Dict[str, Any]:
        """Tenant-visible resources, sorted by uri, as ``{uri, name, mimeType}``."""
        # TODO
        raise NotImplementedError

    def _resources_read(self, params: Mapping[str, Any]) -> Dict[str, Any]:
        """Missing/unknown/other-tenant uri -> INVALID_PARAMS (indistinguishable)."""
        # TODO
        raise NotImplementedError

    def _prompts_list(self, params: Mapping[str, Any]) -> Dict[str, Any]:
        # TODO
        raise NotImplementedError

    def _prompts_get(self, params: Mapping[str, Any]) -> Dict[str, Any]:
        """Unknown prompt or missing arguments -> INVALID_PARAMS. Otherwise interpolate
        ``{name}`` placeholders and return
        ``{"messages": [{"role": "user", "content": {"type": "text", "text": ...}}]}``."""
        # TODO
        raise NotImplementedError

    def notify_tools_changed(self) -> Dict[str, Any]:
        """Build ``notifications/tools/list_changed``, append it to ``outbound``, return
        it. Without this, a client calls a tool that was retired an hour ago."""
        # TODO
        raise NotImplementedError


# ======================================================================================
# 6. The MCP client
# ======================================================================================


class MCPClient:
    """One client per server connection. Caches the tool list; invalidates on
    ``notifications/tools/list_changed`` — the cache is why the notification exists."""

    def __init__(self, server: MCPServer, *, name: str = "bank-agent-kernel", version: str = "1.0.0") -> None:
        self.server = server
        self.name = name
        self.version = version
        self._next_id = 0
        self.server_capabilities: Dict[str, Any] = {}
        self.protocol_version: Optional[str] = None
        self._tool_cache: Optional[List[Dict[str, Any]]] = None
        self.list_calls = 0

    def _send(self, method: str, params: Optional[Mapping[str, Any]] = None) -> Dict[str, Any]:
        """Increment the id, send, and: raise INTERNAL_ERROR if the server returned
        nothing or a mismatched id; re-raise a returned error as JsonRpcError; else
        return ``response["result"]``."""
        # TODO
        raise NotImplementedError

    def initialize(self, protocol_version: str = PROTOCOL_VERSION) -> Dict[str, Any]:
        """Send initialize with clientInfo and capabilities, record the negotiated
        version and the server's capabilities, then send
        ``notifications/initialized``."""
        # TODO
        raise NotImplementedError

    def supports(self, capability: str) -> bool:
        # TODO
        raise NotImplementedError

    def list_tools(self, *, force: bool = False) -> List[Dict[str, Any]]:
        """Cached. Increment ``list_calls`` only on a real round-trip."""
        # TODO
        raise NotImplementedError

    def on_notification(self, message: Mapping[str, Any]) -> None:
        """Invalidate the tool cache on ``notifications/tools/list_changed``."""
        # TODO
        raise NotImplementedError

    def call_tool(self, name: str, arguments: Mapping[str, Any]) -> Dict[str, Any]:
        # TODO
        raise NotImplementedError

    def read_resource(self, uri: str) -> str:
        # TODO
        raise NotImplementedError

    def get_prompt(self, name: str, arguments: Mapping[str, Any]) -> str:
        # TODO
        raise NotImplementedError


# ======================================================================================
# 7. The repair loop
# ======================================================================================


def repair_arguments(
    arguments: Mapping[str, Any],
    schema: Mapping[str, Any],
    *,
    defaults: Optional[Mapping[str, Any]] = None,
) -> Tuple[Dict[str, Any], List[ValidationError]]:
    """Deterministic repairs before asking the model to try again.

    Fix only what is UNAMBIGUOUS:
      - a numeric string where the schema wants integer/number (and vice versa)
      - "true"/"false" where it wants a boolean
      - drop additional properties when additionalProperties is False
      - fill a missing required field ONLY from ``defaults``

    Everything else stays for the model: inventing a missing account number is
    fabrication, not repair. Return ``(repaired, remaining_errors)``. Must be idempotent.
    """
    # TODO
    raise NotImplementedError


def main() -> None:
    print("implement the TODOs, then compare with `python solution.py`")


if __name__ == "__main__":
    main()