"""EDA-first MCP server adapter over EDA runtime APIs."""
from __future__ import annotations

import argparse
from dataclasses import dataclass
from typing import Any, Dict, List
from uuid import uuid4
import os

from integrations.mcp.audit import write_audit_event
from integrations.mcp_eda import clients
from security.auth.core import AuthInput, AuthContext, authorize, resolve_agent_api_key

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-eda-governance")

_TOOL_SCOPE_MAP: dict[str, list[str]] = {
    "perceive": ["eda:perceive"],
    "bind_norms": ["eda:bind"],
    "decide": ["eda:decide"],
    "execute": ["eda:execute"],
    "get_traces": ["traces:read"],
}


def effective_caller_api_key(provided_api_key: str | None) -> str:
    """Resolve effective caller key with tool-arg precedence."""
    return (provided_api_key or "").strip() or os.getenv(
        "AIF_MCP_CALLER_API_KEY", ""
    ).strip()


@dataclass(frozen=True)
class _AuthorizedContext:
    auth: AuthContext
    caller_api_key: str


async def _authorize_tool_call(
    tool_name: str,
    caller_api_key: str | None,
    caller_id: str,
    request_id: str,
) -> _AuthorizedContext:
    required_scopes = _TOOL_SCOPE_MAP.get(tool_name, [])
    effective_api_key = effective_caller_api_key(caller_api_key)
    auth = await authorize(
        AuthInput(
            entry_type="mcp",
            api_key=effective_api_key,
            action=f"tool:{tool_name}",
            required_scopes=required_scopes,
            request_id=request_id,
            caller_id=caller_id,
        ),
        resolver=resolve_agent_api_key,
    )
    if auth.decision != "allow":
        if auth.reason == "missing_api_key":
            raise PermissionError("Unauthorized: missing MCP caller API key.")
        if auth.reason.startswith("missing_scope:"):
            scope = auth.reason.split(":", 1)[1]
            raise PermissionError(f"Forbidden: missing required scope '{scope}'.")
        raise PermissionError("Unauthorized: invalid MCP API key.")
    return _AuthorizedContext(auth=auth, caller_api_key=effective_api_key)


def _audit_policy_denied(request_id: str, caller_id: str, tool_name: str, error: str) -> None:
    write_audit_event(
        {
            "request_id": request_id,
            "caller_id": caller_id,
            "tool": tool_name,
            "ok": False,
            "meta": {
                "denied_reason": "auth_denied",
                "error": error,
            },
        }
    )


@mcp.tool()
async def perceive(
    input_data: Dict[str, Any],
    caller_id: str = "unknown",
    request_id: str = "",
    caller_api_key: str | None = None,
) -> Dict[str, Any]:
    """Perceive stage for the caller-bound EDA agent."""
    rid = request_id or f"mcp-eda-{uuid4()}"
    try:
        auth = await _authorize_tool_call("perceive", caller_api_key, caller_id, rid)
    except PermissionError as exc:
        _audit_policy_denied(rid, caller_id, "perceive", str(exc))
        raise
    result = await clients.perceive(
        agent_id=auth.auth.agent_id,
        input_data=input_data,
        request_id=rid,
        caller_id=caller_id,
        key_id=auth.auth.key_id,
    )
    write_audit_event(
        {
            "request_id": rid,
            "caller_id": caller_id,
            "tool": "perceive",
            "ok": True,
            "meta": {
                "agent_id": auth.auth.agent_id,
                "key_id": auth.auth.key_id,
                "belief_count": len((result.get("beliefs") or {})),
            },
        }
    )
    return result


@mcp.tool()
async def bind_norms(
    context: Dict[str, Any],
    caller_id: str = "unknown",
    request_id: str = "",
    caller_api_key: str | None = None,
) -> Dict[str, Any]:
    """Bind norms for the caller-bound EDA agent."""
    rid = request_id or f"mcp-eda-{uuid4()}"
    try:
        auth = await _authorize_tool_call("bind_norms", caller_api_key, caller_id, rid)
    except PermissionError as exc:
        _audit_policy_denied(rid, caller_id, "bind_norms", str(exc))
        raise
    result = await clients.bind_norms(
        agent_id=auth.auth.agent_id,
        context=context,
        request_id=rid,
        caller_id=caller_id,
        key_id=auth.auth.key_id,
    )
    write_audit_event(
        {
            "request_id": rid,
            "caller_id": caller_id,
            "tool": "bind_norms",
            "ok": True,
            "meta": {
                "agent_id": auth.auth.agent_id,
                "key_id": auth.auth.key_id,
                "norm_binding_count": len(result.get("norm_bindings", []) or []),
            },
        }
    )
    return result


@mcp.tool()
async def decide(
    candidate_actions: List[Dict[str, Any]],
    caller_id: str = "unknown",
    request_id: str = "",
    caller_api_key: str | None = None,
) -> Dict[str, Any]:
    """Decide stage for the caller-bound EDA agent."""
    rid = request_id or f"mcp-eda-{uuid4()}"
    try:
        auth = await _authorize_tool_call("decide", caller_api_key, caller_id, rid)
    except PermissionError as exc:
        _audit_policy_denied(rid, caller_id, "decide", str(exc))
        raise
    result = await clients.decide(
        agent_id=auth.auth.agent_id,
        candidate_actions=candidate_actions,
        request_id=rid,
        caller_id=caller_id,
        key_id=auth.auth.key_id,
    )
    write_audit_event(
        {
            "request_id": rid,
            "caller_id": caller_id,
            "tool": "decide",
            "ok": True,
            "meta": {
                "agent_id": auth.auth.agent_id,
                "key_id": auth.auth.key_id,
                "candidate_count": len(candidate_actions or []),
            },
        }
    )
    return result


@mcp.tool()
async def execute(
    action: Dict[str, Any],
    caller_id: str = "unknown",
    request_id: str = "",
    caller_api_key: str | None = None,
) -> Dict[str, Any]:
    """Execute stage for the caller-bound EDA agent."""
    rid = request_id or f"mcp-eda-{uuid4()}"
    try:
        auth = await _authorize_tool_call("execute", caller_api_key, caller_id, rid)
    except PermissionError as exc:
        _audit_policy_denied(rid, caller_id, "execute", str(exc))
        raise
    result = await clients.execute(
        agent_id=auth.auth.agent_id,
        action=action,
        request_id=rid,
        caller_id=caller_id,
        key_id=auth.auth.key_id,
    )
    write_audit_event(
        {
            "request_id": rid,
            "caller_id": caller_id,
            "tool": "execute",
            "ok": True,
            "meta": {
                "agent_id": auth.auth.agent_id,
                "key_id": auth.auth.key_id,
                "success": bool(result.get("success", False)),
            },
        }
    )
    return result


@mcp.tool()
async def get_traces(
    limit: int = 20,
    caller_id: str = "unknown",
    request_id: str = "",
    caller_api_key: str | None = None,
) -> List[Dict[str, Any]]:
    """Read recent traces for the caller-bound EDA agent."""
    rid = request_id or f"mcp-eda-{uuid4()}"
    try:
        auth = await _authorize_tool_call("get_traces", caller_api_key, caller_id, rid)
    except PermissionError as exc:
        _audit_policy_denied(rid, caller_id, "get_traces", str(exc))
        raise
    result = await clients.get_traces(
        agent_id=auth.auth.agent_id,
        limit=limit,
        request_id=rid,
        caller_id=caller_id,
        key_id=auth.auth.key_id,
    )
    write_audit_event(
        {
            "request_id": rid,
            "caller_id": caller_id,
            "tool": "get_traces",
            "ok": True,
            "meta": {
                "agent_id": auth.auth.agent_id,
                "key_id": auth.auth.key_id,
                "limit": limit,
                "count": len(result or []),
            },
        }
    )
    return result


def main() -> None:
    """Run AIF EDA-first MCP adapter."""
    parser = argparse.ArgumentParser(description="Run AIF EDA-first 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()
