"""AIF MCP server adapter over Normbase and EDA APIs."""
from __future__ import annotations

import argparse
from typing import Any, Dict, List
from uuid import uuid4

from integrations.mcp.audit import write_audit_event
from integrations.mcp import clients
from integrations.mcp.security import authenticate_mcp_caller
from security.auth.api_keys import ResolvedAPIKey

try:
    from mcp.server.fastmcp import FastMCP
except Exception as exc:  # pragma: no cover
    raise RuntimeError(
        "MCP SDK import failed. Install dependencies from requirements.txt first."
    ) from exc


mcp = FastMCP("aif-governance")


async def _require_auth(caller_api_key: str | None) -> ResolvedAPIKey:
    return await authenticate_mcp_caller(caller_api_key)


@mcp.tool()
async def suggest_roles(
    context: Dict[str, Any],
    top_k: int = 5,
    caller_id: str = "unknown",
    request_id: str = "",
    caller_api_key: str | None = None,
) -> Dict[str, Any]:
    """Suggest candidate agent roles for the given context."""
    principal = await _require_auth(caller_api_key)
    rid = request_id or f"mcp-{uuid4()}"
    field_ids: List[str] = []
    if principal.agent_id:
        try:
            field_ids = await clients.get_agent_field_ids(principal.agent_id)
        except Exception:
            field_ids = []
    result = await clients.suggest_roles(
        context=context,
        field_ids=field_ids or None,
        top_k=top_k,
        request_id=rid,
        caller_id=caller_id,
        agent_id=principal.agent_id,
        key_id=principal.key_id,
    )
    write_audit_event(
        {
            "request_id": rid,
            "caller_id": caller_id,
            "tool": "suggest_roles",
            "ok": True,
            "meta": {
                "top_k": top_k,
                "candidate_count": len(result.get("candidates", []) or []),
                "recommended_role": result.get("recommended_role"),
                "field_ids_count": len(field_ids),
                "agent_id": principal.agent_id,
                "key_id": principal.key_id,
            },
        }
    )
    return result


@mcp.tool()
async def retrieve_norms(
    context: Dict[str, Any],
    agent_role: str,
    top_k: int = 20,
    caller_id: str = "unknown",
    request_id: str = "",
    caller_api_key: str | None = None,
) -> Dict[str, Any]:
    """Retrieve candidate norms from Normbase."""
    principal = await _require_auth(caller_api_key)
    rid = request_id or f"mcp-{uuid4()}"
    field_ids: List[str] = []
    if principal.agent_id:
        try:
            field_ids = await clients.get_agent_field_ids(principal.agent_id)
        except Exception:
            # Field auto-discovery is best-effort; retrieval falls back to role/context.
            field_ids = []
    result = await clients.retrieve_norms(
        context=context,
        agent_role=agent_role,
        top_k=top_k,
        field_ids=field_ids or None,
        request_id=rid,
        caller_id=caller_id,
        agent_id=principal.agent_id,
        key_id=principal.key_id,
    )
    write_audit_event(
        {
            "request_id": rid,
            "caller_id": caller_id,
            "tool": "retrieve_norms",
            "ok": True,
            "meta": {
                "agent_role": agent_role,
                "top_k": top_k,
                "total": result.get("total", 0),
                "field_ids_count": len(field_ids),
                "agent_id": principal.agent_id,
                "key_id": principal.key_id,
            },
        }
    )
    return result


@mcp.tool()
async def bind_norms(
    agent_id: str,
    context: Dict[str, Any],
    caller_id: str = "unknown",
    request_id: str = "",
    caller_api_key: str | None = None,
) -> Dict[str, Any]:
    """Bind applicable norms in EDA for the given context."""
    principal = await _require_auth(caller_api_key)
    rid = request_id or f"mcp-{uuid4()}"
    if principal.agent_id != agent_id:
        raise PermissionError("Unauthorized: key does not match target agent.")
    result = await clients.bind_norms(
        agent_id=agent_id,
        context=context,
        request_id=rid,
        caller_id=caller_id,
        key_id=principal.key_id,
    )
    write_audit_event(
        {
            "request_id": rid,
            "caller_id": caller_id,
            "tool": "bind_norms",
            "ok": True,
            "meta": {"agent_id": agent_id, "key_id": principal.key_id},
        }
    )
    return result


@mcp.tool()
async def get_traces(
    agent_id: str,
    limit: int = 20,
    caller_id: str = "unknown",
    request_id: str = "",
    caller_api_key: str | None = None,
) -> List[Dict[str, Any]]:
    """Fetch recent EDA traces for audit/debug."""
    principal = await _require_auth(caller_api_key)
    rid = request_id or f"mcp-{uuid4()}"
    if principal.agent_id != agent_id:
        raise PermissionError("Unauthorized: key does not match target agent.")
    result = await clients.get_traces(
        agent_id=agent_id,
        limit=limit,
        request_id=rid,
        caller_id=caller_id,
        key_id=principal.key_id,
    )
    write_audit_event(
        {
            "request_id": rid,
            "caller_id": caller_id,
            "tool": "get_traces",
            "ok": True,
            "meta": {
                "agent_id": agent_id,
                "limit": limit,
                "count": len(result or []),
                "key_id": principal.key_id,
            },
        }
    )
    return result


def main() -> None:
    parser = argparse.ArgumentParser(description="Run AIF MCP adapter.")
    parser.add_argument(
        "--transport",
        choices=("stdio", "sse", "streamable-http"),
        default="stdio",
        help="MCP transport mode (default: stdio)",
    )
    args = parser.parse_args()
    mcp.run(args.transport)


if __name__ == "__main__":
    main()

