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

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

from integrations.mcp.audit import write_audit_event
from integrations.mcp import clients
from integrations.mcp.security import authenticate_mcp_caller, effective_caller_api_key
from security.auth.api_keys import ResolvedAPIKey
from security.auth.tool_policy import ToolPolicy, get_tool_policy, load_tool_policies

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")
SCOPE_WILDCARD = "*"
# Fail fast on missing/invalid tool policy file.
load_tool_policies()


@dataclass(frozen=True)
class _AuthorizedContext:
    principal: ResolvedAPIKey
    caller_api_key: str


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


def _has_scope(principal: ResolvedAPIKey, required_scope: str) -> bool:
    scopes = [str(scope).strip() for scope in (principal.scopes or []) if str(scope).strip()]
    if SCOPE_WILDCARD in scopes:
        return True
    return required_scope in scopes


def _get_tool_policy(tool_name: str) -> ToolPolicy | None:
    return get_tool_policy(tool_name)


async def _authorize_tool_call(
    tool_name: str,
    caller_api_key: str | None,
    caller_id: str,
    request_id: str,
) -> _AuthorizedContext:
    policy = _get_tool_policy(tool_name)
    if policy is None:
        raise PermissionError(f"Forbidden: tool '{tool_name}' is not configured.")
    if not policy.enabled:
        raise PermissionError(f"Forbidden: tool '{tool_name}' is disabled.")
    if policy.auth_mode == "internal_only":
        raise PermissionError(f"Forbidden: tool '{tool_name}' is internal-only.")

    principal = await _require_auth(caller_api_key)
    for scope in policy.required_scopes:
        if _has_scope(principal, scope):
            continue
        raise PermissionError(f"Forbidden: missing required scope '{scope}'.")
    return _AuthorizedContext(
        principal=principal,
        caller_api_key=effective_caller_api_key(caller_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": "tool_policy_denied",
                "error": error,
            },
        }
    )


def _normalized_agent_id_or_none(raw_agent_id: str | None) -> str | None:
    text = (raw_agent_id or "").strip()
    if not text:
        return None
    try:
        return str(UUID(text))
    except ValueError as exc:
        raise PermissionError("Unauthorized: key contains non-UUID agent_id.") from exc


@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."""
    rid = request_id or f"mcp-{uuid4()}"
    try:
        auth = await _authorize_tool_call("suggest_roles", caller_api_key, caller_id, rid)
    except PermissionError as exc:
        _audit_policy_denied(rid, caller_id, "suggest_roles", str(exc))
        raise
    principal = auth.principal
    normalized_agent_id = _normalized_agent_id_or_none(principal.agent_id)
    field_ids: List[str] = []
    if normalized_agent_id:
        try:
            field_ids = await clients.get_agent_field_ids(normalized_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=normalized_agent_id or "",
        key_id=principal.key_id,
        caller_api_key=auth.caller_api_key,
    )
    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": normalized_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."""
    rid = request_id or f"mcp-{uuid4()}"
    try:
        auth = await _authorize_tool_call("retrieve_norms", caller_api_key, caller_id, rid)
    except PermissionError as exc:
        _audit_policy_denied(rid, caller_id, "retrieve_norms", str(exc))
        raise
    principal = auth.principal
    normalized_agent_id = _normalized_agent_id_or_none(principal.agent_id)
    field_ids: List[str] = []
    if normalized_agent_id:
        try:
            field_ids = await clients.get_agent_field_ids(normalized_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=normalized_agent_id or "",
        key_id=principal.key_id,
        caller_api_key=auth.caller_api_key,
    )
    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": normalized_agent_id,
                "key_id": principal.key_id,
            },
        }
    )
    return result


@mcp.tool()
async def check_action_compliance(
    usage_token: str,
    action_description: str,
    norm_ids: List[str] | None = None,
    caller_id: str = "unknown",
    request_id: str = "",
    caller_api_key: str | None = None,
) -> Dict[str, Any]:
    """Check whether an action violates selected norms."""
    rid = request_id or f"mcp-{uuid4()}"
    try:
        auth = await _authorize_tool_call("check_action_compliance", caller_api_key, caller_id, rid)
    except PermissionError as exc:
        _audit_policy_denied(rid, caller_id, "check_action_compliance", str(exc))
        raise
    principal = auth.principal
    normalized_agent_id = _normalized_agent_id_or_none(principal.agent_id)
    result = await clients.check_action_compliance(
        usage_token=usage_token,
        action_description=action_description,
        norm_ids=norm_ids or None,
        request_id=rid,
        caller_id=caller_id,
        agent_id=normalized_agent_id or "",
        key_id=principal.key_id,
        caller_api_key=auth.caller_api_key,
    )
    write_audit_event(
        {
            "request_id": rid,
            "caller_id": caller_id,
            "tool": "check_action_compliance",
            "ok": True,
            "meta": {
                "norm_ids_count": len(norm_ids or []),
                "usage_token": usage_token,
                "compliant": bool(result.get("compliant", False)),
                "violation_count": len(result.get("violations", []) or []),
                "agent_id": normalized_agent_id,
                "key_id": principal.key_id,
            },
        }
    )
    return result


@mcp.tool()
async def report_norm_outcome(
    usage_token: str,
    norm_id: str,
    outcome: str,
    details: Dict[str, Any] | None = None,
    caller_id: str = "unknown",
    request_id: str = "",
    caller_api_key: str | None = None,
) -> Dict[str, Any]:
    """Report one norm execution outcome for governance logs."""
    rid = request_id or f"mcp-{uuid4()}"
    try:
        auth = await _authorize_tool_call("report_norm_outcome", caller_api_key, caller_id, rid)
    except PermissionError as exc:
        _audit_policy_denied(rid, caller_id, "report_norm_outcome", str(exc))
        raise
    principal = auth.principal
    normalized_agent_id = _normalized_agent_id_or_none(principal.agent_id)
    result = await clients.report_norm_outcome(
        usage_token=usage_token,
        norm_id=norm_id,
        outcome=outcome,
        details=details or {},
        request_id=rid,
        caller_id=caller_id,
        agent_id=normalized_agent_id or "",
        key_id=principal.key_id,
        caller_api_key=auth.caller_api_key,
    )
    write_audit_event(
        {
            "request_id": rid,
            "caller_id": caller_id,
            "tool": "report_norm_outcome",
            "ok": True,
            "meta": {
                "usage_token": usage_token,
                "norm_id": norm_id,
                "outcome": outcome,
                "agent_id": normalized_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."""
    rid = request_id or f"mcp-{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
    principal = auth.principal
    principal_agent_id = _normalized_agent_id_or_none(principal.agent_id)
    requested_agent_id = _normalized_agent_id_or_none(agent_id)
    if principal_agent_id is None or requested_agent_id is None:
        raise PermissionError("Unauthorized: missing UUID agent_id.")
    if principal_agent_id != requested_agent_id:
        raise PermissionError("Unauthorized: key does not match target agent.")
    result = await clients.get_traces(
        agent_id=requested_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": requested_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()

