"""
mcp_mini.py — a minimal Model Context Protocol implementation (Knowledge Module
03 §9), so you can SEE the protocol rather than just name it.

MCP is a CLIENT–SERVER protocol (JSON-RPC) that standardizes how an LLM host
discovers and uses external capabilities. A server exposes three things:
  - TOOLS     : callable functions (with a JSON schema)   -> "tools/list", "tools/call"
  - RESOURCES : readable data/context (by URI)            -> "resources/list", "resources/read"
  - PROMPTS   : reusable prompt templates                 -> "prompts/list", "prompts/get"

The win for a bank (K03 §9): write your core-banking tools ONCE as an MCP server
with central auth/governance; ANY MCP-aware host (LangGraph, Foundry, Claude, …)
consumes them without a rewrite, and the SERVER — not the prompt — enforces
permissions. This file is in-process; the real protocol runs over stdio/HTTP.
"""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Callable


@dataclass
class Tool:
    name: str
    description: str
    input_schema: dict
    handler: Callable[[dict], dict]


class MCPServer:
    """Exposes tools/resources/prompts and answers JSON-RPC requests."""

    def __init__(self, name: str) -> None:
        self.name = name
        self._tools: dict[str, Tool] = {}
        self._resources: dict[str, tuple[str, Callable[[], str]]] = {}
        self._prompts: dict[str, tuple[str, str]] = {}   # name -> (description, template)

    # --- registration (what a tool author writes) ---
    def tool(self, name: str, description: str, input_schema: dict):
        def deco(fn: Callable[[dict], dict]):
            self._tools[name] = Tool(name, description, input_schema, fn)
            return fn
        return deco

    def resource(self, uri: str, description: str):
        def deco(fn: Callable[[], str]):
            self._resources[uri] = (description, fn)
            return fn
        return deco

    def prompt(self, name: str, description: str, template: str) -> None:
        self._prompts[name] = (description, template)

    # --- JSON-RPC dispatch (what the protocol does) ---
    def handle(self, request: dict) -> dict:
        rid, method, params = request.get("id"), request.get("method"), request.get("params", {})
        try:
            result = self._dispatch(method, params)
            return {"jsonrpc": "2.0", "id": rid, "result": result}
        except KeyError as e:
            return {"jsonrpc": "2.0", "id": rid,
                    "error": {"code": -32601, "message": f"not found: {e}"}}
        except Exception as e:  # noqa: BLE001 — surface tool errors as JSON-RPC errors
            return {"jsonrpc": "2.0", "id": rid,
                    "error": {"code": -32000, "message": str(e)}}

    def _dispatch(self, method: str, params: dict):
        if method == "tools/list":
            return {"tools": [{"name": t.name, "description": t.description,
                               "inputSchema": t.input_schema} for t in self._tools.values()]}
        if method == "tools/call":
            tool = self._tools[params["name"]]
            return {"content": tool.handler(params.get("arguments", {}))}
        if method == "resources/list":
            return {"resources": [{"uri": u, "description": d}
                                  for u, (d, _) in self._resources.items()]}
        if method == "resources/read":
            _, fn = self._resources[params["uri"]]
            return {"contents": fn()}
        if method == "prompts/list":
            return {"prompts": [{"name": n, "description": d}
                               for n, (d, _) in self._prompts.items()]}
        if method == "prompts/get":
            _, template = self._prompts[params["name"]]
            return {"messages": template.format(**params.get("arguments", {}))}
        raise KeyError(method)


class MCPClient:
    """An MCP host's view of a server. In-process here; in production it speaks
    JSON-RPC over stdio or HTTP to a separately-deployed server."""

    def __init__(self, server: MCPServer) -> None:
        self.server = server
        self._id = 0

    def _rpc(self, method: str, params: dict | None = None) -> dict:
        self._id += 1
        resp = self.server.handle({"jsonrpc": "2.0", "id": self._id,
                                   "method": method, "params": params or {}})
        if "error" in resp:
            raise RuntimeError(resp["error"]["message"])
        return resp["result"]

    def list_tools(self) -> list[dict]:
        return self._rpc("tools/list")["tools"]

    def call_tool(self, name: str, arguments: dict) -> dict:
        return self._rpc("tools/call", {"name": name, "arguments": arguments})["content"]

    def read_resource(self, uri: str) -> str:
        return self._rpc("resources/read", {"uri": uri})["contents"]

    def get_prompt(self, name: str, arguments: dict | None = None) -> str:
        return self._rpc("prompts/get", {"name": name, "arguments": arguments or {}})["messages"]


# --- an example "core-banking MCP server" -----------------------------------
def build_core_banking_server() -> MCPServer:
    """The bank writes this ONCE; any agent host consumes it. Note the tool
    enforces the user's permission itself — the host/prompt is never trusted."""
    s = MCPServer("core-banking")
    _accounts = {"AE070331234567890123456": {"owner": "user_amal", "balance": 1250.00}}

    @s.tool("get_balance", "Return the balance for an account the user owns.",
            {"type": "object",
             "properties": {"account_id": {"type": "string"}, "user": {"type": "string"}},
             "required": ["account_id", "user"]})
    def get_balance(args: dict) -> dict:
        acct = _accounts.get(args["account_id"])
        if not acct or acct["owner"] != args["user"]:
            raise PermissionError("not authorized for this account")
        return {"account_id": args["account_id"], "balance": acct["balance"], "currency": "AED"}

    @s.resource("bank://policies/dispute", "The transaction dispute policy text.")
    def dispute_policy() -> str:
        return "A customer may dispute a card transaction within 60 days of the statement date."

    s.prompt("grounded_answer", "A grounding instruction template.",
             "Answer ONLY from: {sources}. If absent, say you don't have it. Cite [source].")
    return s