"""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, Literal
from uuid import UUID, uuid4

from integrations.mcp.audit import write_audit_event
from integrations.mcp import clients
from integrations.mcp.security import authorize_mcp_tool_call, 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 authorize_mcp_tool_call(
        caller_api_key,
        tool_name="authenticate_mcp_caller",
        required_scopes=[],
    )


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


def _resolve_query_mode(value: str) -> Literal["structured", "semantic", "hybrid"]:
    normalized = (value or "").strip().lower()
    if normalized in {"structured", "semantic", "hybrid"}:
        return normalized  # type: ignore[return-value]
    return "hybrid"


def _resolve_agent_scope(value: str) -> Literal["prefer", "strict"]:
    normalized = (value or "").strip().lower()
    if normalized == "strict":
        return "strict"
    return "prefer"


def _extract_subject_agents(norm: Dict[str, Any]) -> set[str]:
    subject = norm.get("subject", {}) if isinstance(norm, dict) else {}
    canonical = subject.get("canonical", {}) if isinstance(subject, dict) else {}
    raw_agents = canonical.get("agents", []) if isinstance(canonical, dict) else []
    if not isinstance(raw_agents, list):
        return set()
    return {str(item).strip() for item in raw_agents if str(item).strip()}


def _annotate_agent_segments(
    result: Dict[str, Any],
    normalized_agent_id: str | None,
) -> Dict[str, Any]:
    norms = result.get("norms", []) or []
    if not isinstance(norms, list):
        return result
    agent_specific_ids: List[str] = []
    reference_ids: List[str] = []
    for item in norms:
        if not isinstance(item, dict):
            continue
        norm_id = str(item.get("id", "")).strip()
        agents = _extract_subject_agents(item)
        is_agent_specific = bool(normalized_agent_id and normalized_agent_id in agents)
        item["agent_specific"] = is_agent_specific
        item["agent_match"] = "agent_specific" if is_agent_specific else "reference"
        if not norm_id:
            continue
        if is_agent_specific:
            agent_specific_ids.append(norm_id)
        else:
            reference_ids.append(norm_id)
    result["mcp_segments"] = {
        "agent_specific_count": len(agent_specific_ids),
        "reference_count": len(reference_ids),
        "agent_specific_ids": agent_specific_ids,
        "reference_ids": reference_ids,
    }
    return result


@mcp.tool()
async def plan_retrieval(
    context: Dict[str, Any],
    top_k: int = 5,
    caller_id: str = "unknown",
    request_id: str = "",
    caller_api_key: str | None = None,
) -> Dict[str, Any]:
    """Plan retrieval filters (recommended role + suggested tag_filters)."""
    rid = request_id or f"mcp-{uuid4()}"
    try:
        auth = await _authorize_tool_call("plan_retrieval", caller_api_key, caller_id, rid)
    except PermissionError as exc:
        _audit_policy_denied(rid, caller_id, "plan_retrieval", 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.plan_retrieval(
        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": "plan_retrieval",
            "ok": True,
            "meta": {
                "top_k": top_k,
                "candidate_count": len(result.get("candidates", []) or []),
                "recommended_role": result.get("recommended_role"),
                "tag_filter_categories": sorted(
                    list((result.get("suggested_tag_filters") or {}).keys())
                ),
                "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 = "",
    tag_filters: Dict[str, List[str]] | None = None,
    query_mode: str = "hybrid",
    nl_query: str = "",
    agent_scope: str = "prefer",
    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)
    normalized_query_mode = _resolve_query_mode(query_mode)
    normalized_agent_scope = _resolve_agent_scope(agent_scope)
    normalized_nl_query = (nl_query or "").strip()
    if normalized_query_mode == "semantic" and not normalized_nl_query:
        raise ValueError("nl_query is required when query_mode='semantic'.")

    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 = []
    agent_ids_for_retrieve = (
        [normalized_agent_id]
        if normalized_agent_scope == "strict" and normalized_agent_id
        else None
    )
    result = await clients.retrieve_norms(
        context=context,
        query_mode=normalized_query_mode,
        nl_query=normalized_nl_query or None,
        agent_role=(agent_role or "").strip() or None,
        agent_ids=agent_ids_for_retrieve,
        tag_filters=tag_filters or None,
        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,
    )
    result = _annotate_agent_segments(result, normalized_agent_id)
    write_audit_event(
        {
            "request_id": rid,
            "caller_id": caller_id,
            "tool": "retrieve_norms",
            "ok": True,
            "meta": {
                "agent_role": (agent_role or "").strip() or None,
                "tag_filter_categories": sorted(list((tag_filters or {}).keys())),
                "top_k": top_k,
                "total": result.get("total", 0),
                "field_ids_count": len(field_ids),
                "query_mode": normalized_query_mode,
                "agent_scope": normalized_agent_scope,
                "agent_specific_count": (result.get("mcp_segments") or {}).get(
                    "agent_specific_count", 0
                ),
                "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()}"
    norm_id_text = str(norm_id).strip()
    try:
        normalized_norm_id = str(UUID(norm_id_text))
    except ValueError as exc:
        raise ValueError(
            "Invalid norm_id: must be UUID from retrieve_norms norms[].id. "
            "If retrieve returned total=0, do not call report_norm_outcome."
        ) from exc
    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=normalized_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": normalized_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()

