"""GSAP Broker MCP Endpoint — Governance Primitives as MCP Tools.

Exposes the governance protocol as MCP tools for any MCP-compatible agent.
This is the consortia-builder interface. Capstone MCP is the MSP-operator interface.

Protocol: Streamable HTTP JSON-RPC 2.0 (MCP spec 2024-11-05)
Auth: Bearer token (Keycloak JWT or GSAP bearer token)

Tool handlers call the broker's existing internal functions — no logic duplication.
The MCP layer is protocol translation: JSON-RPC → internal function → JSON-RPC response.
"""

import json
import logging
import os
import time
from datetime import datetime, UTC

import httpx
from fastapi import APIRouter, Request
from fastapi.responses import JSONResponse

from gsap_broker.settings import settings
from gsap_broker import chronicle

logger = logging.getLogger(__name__)

router = APIRouter()

MCP_PROTOCOL_VERSION = "2024-11-05"
SERVER_NAME = "gsap-governance"
SERVER_VERSION = "0.1.0"

LLM_BROKER_URL = os.getenv("LLM_BROKER_URL", "http://localhost:8092")
DEFCON_POSTURE_SOURCE = os.getenv("DEFCON_POSTURE_SOURCE", "")

# ── Tool Definitions ─────────────────────────────────────────────

TOOLS = [
    {
        "name": "request_ac",
        "description": (
            "Request an Authorization Context for a governed session. "
            "The AC authorizes a principal to operate under a specific Accord with a defined corpus."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "principal": {"type": "string", "description": "DID of the principal requesting authorization"},
                "accord_template": {"type": "string", "description": "AccordTemplate name (e.g. shell-exec)"},
                "corpus_entry_cid": {"type": "string", "description": "Content ID of the corpus entry"},
                "playbook": {"type": "string", "description": "Playbook identifier for the operation"},
                "parameters_cid": {"type": "string", "description": "Content hash of operation parameters"},
                "session_mode": {"type": "boolean", "description": "If true, AC stays active across multiple CRs"},
            },
            "required": ["principal", "accord_template", "corpus_entry_cid"],
        },
    },
    {
        "name": "validate_ac",
        "description": "Check if an Authorization Context is currently valid and active.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "ac_id": {"type": "string", "description": "The AC poll_token or context_id to validate"},
            },
            "required": ["ac_id"],
        },
    },
    {
        "name": "post_cr",
        "description": "Post a Completion Receipt recording that a governed operation was performed.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "context_id": {"type": "string", "description": "The AC under which this operation was performed"},
                "outcome": {"type": "string", "enum": ["completed", "failed", "violated", "timed_out", "session_end"]},
                "failure_reason": {"type": "string", "description": "Reason for failure (if outcome=failed)"},
            },
            "required": ["context_id", "outcome"],
        },
    },
    {
        "name": "query_accord",
        "description": "Look up an AccordTemplate's governance parameters.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "accord_name": {"type": "string", "description": "AccordTemplate name"},
            },
            "required": ["accord_name"],
        },
    },
    {
        "name": "request_delegation",
        "description": "Delegate authority to an AI agent via the LLM Principal Broker.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "delegator_ac_id": {"type": "string", "description": "The delegator's AC ID"},
                "agent_type": {"type": "string", "description": "Agent type (e.g. claude-code)"},
                "capability_ceiling": {"type": "string", "enum": ["CAP_READ", "CAP_PROPOSE", "CAP_MUTATE"]},
                "max_ttl_minutes": {"type": "integer", "description": "Delegation TTL in minutes"},
                "ceremony_required_for": {"type": "array", "items": {"type": "string"}},
            },
            "required": ["delegator_ac_id", "agent_type"],
        },
    },
    {
        "name": "revoke_delegation",
        "description": "Revoke an active agent delegation.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "delegation_id": {"type": "string"},
                "reason": {"type": "string"},
            },
            "required": ["delegation_id"],
        },
    },
    {
        "name": "get_delegation",
        "description": "Query delegation status: active/expired/revoked, commands remaining, TTL.",
        "inputSchema": {
            "type": "object",
            "properties": {"delegation_id": {"type": "string"}},
            "required": ["delegation_id"],
        },
    },
    {
        "name": "list_agents",
        "description": "List all active agent delegations.",
        "inputSchema": {"type": "object", "properties": {}},
    },
    {
        "name": "get_posture",
        "description": "Get the current DEFCON operational posture level and restrictions.",
        "inputSchema": {
            "type": "object",
            "properties": {"scope": {"type": "string", "description": "Posture scope (default: cluster)"}},
        },
    },
    {
        "name": "check_operation",
        "description": "Dry-run: check if an operation would be allowed at current posture and delegation scope.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "operation_type": {"type": "string", "enum": ["read", "write", "delete", "execute", "delegate"]},
                "target_layer": {"type": "string", "enum": ["network", "systems", "applications"]},
                "ac_id": {"type": "string", "description": "AC to check against (optional)"},
            },
            "required": ["operation_type"],
        },
    },
    {
        "name": "session_info",
        "description": "Get current session details: principal, AC scope, delegation, DEFCON level.",
        "inputSchema": {"type": "object", "properties": {}},
    },
    {
        "name": "list_devices",
        "description": "List managed devices from Intune. Requires Intune connector enabled.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "top": {"type": "integer", "description": "Max devices to return (default: 50)"},
            },
        },
    },
    {
        "name": "get_device_compliance",
        "description": "Check compliance state of a specific device via Intune.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "device_id": {"type": "string", "description": "Intune managed device ID"},
            },
            "required": ["device_id"],
        },
    },
    {
        "name": "sync_device",
        "description": "Trigger Intune sync for a device. Requires PROPOSE capability.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "device_id": {"type": "string", "description": "Intune managed device ID"},
            },
            "required": ["device_id"],
        },
    },
    {
        "name": "remote_lock",
        "description": "Remote lock a managed device. Requires MUTATE capability. May require ceremony approval in production Accords.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "device_id": {"type": "string", "description": "Intune managed device ID"},
            },
            "required": ["device_id"],
        },
    },
]


# ── Posture Cache ────────────────────────────────────────────────

_posture_cache = {"level": 5, "fetched_at": 0.0}


def _get_posture_level() -> int:
    now = time.time()
    if (now - _posture_cache["fetched_at"]) < 30:
        return _posture_cache["level"]

    source = DEFCON_POSTURE_SOURCE
    level = 5
    if source:
        try:
            if source.startswith(("http://", "https://")):
                resp = httpx.get(f"{source.rstrip('/')}/health", timeout=3.0)
                level = resp.json().get("defcon_level", 5)
            else:
                with open(source) as f:
                    level = json.load(f).get("defcon_level", 5)
        except Exception as e:
            logger.debug("Posture fetch failed: %s", e)

    _posture_cache["level"] = level
    _posture_cache["fetched_at"] = now
    return level


# ── Helper: extract principal from request ───────────────────────

def _extract_principal(request: Request) -> str:
    auth = request.headers.get("authorization", "")
    if auth.startswith("Bearer "):
        try:
            import base64
            payload = auth[7:].split(".")[1]
            payload += "=" * (4 - len(payload) % 4)
            claims = json.loads(base64.urlsafe_b64decode(payload))
            return claims.get("preferred_username", claims.get("sub", ""))
        except Exception:
            pass
    return ""


def _extract_delegation(request: Request) -> dict | None:
    auth = request.headers.get("authorization", "")
    delegation_id = ""
    if auth.startswith("Bearer "):
        try:
            import base64
            payload = auth[7:].split(".")[1]
            payload += "=" * (4 - len(payload) % 4)
            claims = json.loads(base64.urlsafe_b64decode(payload))
            delegation_id = claims.get("delegation_id", "")
        except Exception:
            pass

    delegation_id = delegation_id or request.headers.get("x-delegation-id", "")
    if not delegation_id:
        return None

    return {
        "delegation_id": delegation_id,
        "delegator_did": request.headers.get("x-delegator-did", ""),
        "agent_did": request.headers.get("x-agent-did", ""),
    }


# ── Tool Handlers ────────────────────────────────────────────────


async def _handle_request_ac(request: Request, args: dict) -> dict:
    """Wraps the broker's existing /governance/authorize/ endpoint."""
    from gsap_broker.routers.authorize import authorize, _extract_token_data
    from gsap_broker.models import AuthorizeRequest
    from gsap_broker.db import get_session

    body = AuthorizeRequest(
        playbook=args.get("playbook", "mcp-session"),
        corpus_entry_cid=args["corpus_entry_cid"],
        parameters_cid=args.get("parameters_cid", ""),
        accord_template=args["accord_template"],
        driver_id="keycloak",
        session_mode=args.get("session_mode", False),
        on_behalf_of=args.get("principal"),
    )

    async for session in get_session():
        try:
            result = await authorize(body, request, session)
            return result.model_dump(mode="json", exclude_none=True)
        except Exception as e:
            return {"error": str(e)}


async def _handle_validate_ac(args: dict) -> dict:
    """Wraps the broker's existing /governance/authorize/{poll_token}/ endpoint."""
    from gsap_broker.routers.authorize import authorize_poll
    from gsap_broker.db import get_session

    ac_id = args["ac_id"]
    async for session in get_session():
        try:
            result = await authorize_poll(ac_id, session)
            return result.model_dump(mode="json", exclude_none=True)
        except Exception as e:
            return {"error": str(e)}


async def _handle_post_cr(request: Request, args: dict) -> dict:
    """Wraps the broker's existing /governance/complete/ endpoint."""
    from gsap_broker.routers.complete import complete
    from gsap_broker.models import CompleteRequest, ChronicleEvidence, BehavioralAttestation
    from gsap_broker.db import get_session

    body = CompleteRequest(
        context_id=args["context_id"],
        outcome=args["outcome"],
        completed_at=datetime.now(UTC),
        failure_reason=args.get("failure_reason"),
        chronicle_evidence=ChronicleEvidence(),
        behavioral_attestation=BehavioralAttestation(),
    )

    async for session in get_session():
        try:
            result = await complete(body, session)
            return result.model_dump(mode="json", exclude_none=True)
        except Exception as e:
            return {"error": str(e)}


async def _handle_query_accord(args: dict) -> dict:
    """Look up an AccordTemplate. Currently returns known templates from config."""
    name = args["accord_name"]
    # Known accord templates (from GSAP broker config)
    templates = {
        "shell-exec": {
            "name": "shell-exec",
            "capability_ceiling": "CAP_MUTATE",
            "session_ttl_minutes": 30,
            "mfa_required": False,
            "ceremony_gate": None,
        },
        "dev-operations": {
            "name": "dev-operations",
            "capability_ceiling": "CAP_MUTATE",
            "session_ttl_minutes": 60,
            "mfa_required": False,
            "ceremony_gate": None,
        },
        "network-mutate": {
            "name": "network-mutate",
            "capability_ceiling": "CAP_GOVERN",
            "session_ttl_minutes": 15,
            "mfa_required": True,
            "ceremony_gate": "network-admin-elevated",
        },
        "ai-delegation-standard": {
            "name": "ai-delegation-standard",
            "capability_ceiling": "CAP_MUTATE",
            "session_ttl_minutes": 60,
            "ceremony_required_for": ["delete", "destroy", "drop"],
            "max_commands": 500,
        },
    }
    if name in templates:
        return templates[name]
    return {"error": f"AccordTemplate '{name}' not found"}


async def _proxy_to_llm_broker(tool_name: str, args: dict, request: Request) -> dict:
    """Proxy delegation operations to the LLM Principal Broker."""
    if not LLM_BROKER_URL:
        return {"error": "LLM Principal Broker not configured"}

    endpoint_map = {
        "request_delegation": ("POST", "/delegate"),
        "revoke_delegation": ("POST", f"/delegate/{args.get('delegation_id', '')}/revoke"),
        "get_delegation": ("GET", f"/delegate/{args.get('delegation_id', '')}"),
        "list_agents": ("GET", "/agents"),
    }

    method, path = endpoint_map.get(tool_name, (None, None))
    if not method:
        return {"error": f"Unknown delegation tool: {tool_name}"}

    headers = {}
    auth_header = request.headers.get("authorization")
    if auth_header:
        headers["Authorization"] = auth_header

    principal = _extract_principal(request)
    if principal:
        headers["X-Delegator-DID"] = f"did:web:guildhouse.dev/user/{principal}"

    try:
        async with httpx.AsyncClient(timeout=10.0) as client:
            if method == "GET":
                resp = await client.get(f"{LLM_BROKER_URL}{path}", headers=headers)
            else:
                resp = await client.post(f"{LLM_BROKER_URL}{path}", json=args, headers=headers)
            return resp.json()
    except Exception as e:
        return {"error": f"LLM Principal Broker unavailable: {e}"}


async def _handle_get_posture(args: dict) -> dict:
    level = _get_posture_level()
    restrictions = []
    if level <= 2:
        restrictions.append("write operations denied")
    if level <= 1:
        restrictions.append("all operations suspended")
    if level <= 3:
        restrictions.append("ceremony required for writes")
    if level <= 4 and level > 1:
        restrictions.append("agent TTL capped at 30min")

    return {
        "defcon_level": level,
        "posture": ["peacetime", "elevated", "restricted", "critical", "lockdown"][5 - level],
        "restrictions": restrictions,
    }


async def _handle_check_operation(args: dict, request: Request) -> dict:
    level = _get_posture_level()
    op_type = args.get("operation_type", "read")

    if level == 1:
        return {"allowed": False, "reason": "DEFCON 1 — all operations suspended", "defcon_level": level}
    if level <= 2 and op_type in ("write", "delete", "execute"):
        return {"allowed": False, "reason": f"DEFCON {level} — {op_type} operations denied", "defcon_level": level}
    if level == 3 and op_type in ("write", "delete"):
        return {"allowed": False, "reason": f"DEFCON 3 — {op_type} requires ceremony", "defcon_level": level}

    return {"allowed": True, "defcon_level": level, "operation_type": op_type}


async def _handle_session_info(request: Request) -> dict:
    return {
        "principal": _extract_principal(request),
        "delegation": _extract_delegation(request),
        "defcon_level": _get_posture_level(),
        "broker_did": settings.broker_did,
        "timestamp": datetime.now(UTC).isoformat(),
    }


async def _handle_intune_tool(tool_name: str, args: dict) -> dict:
    """Route Intune MCP tools through the governed IntuneConnector."""
    from gsap_broker.routers.connectors import _registry
    from gsap_broker.connectors.base import ConnectorContext

    intune = _registry.get("intune")
    if intune is None:
        return {"error": "Intune connector not enabled. Set intune_enabled=True."}

    op_map = {
        "list_devices": "list_devices",
        "get_device_compliance": "get_compliance",
        "sync_device": "sync_device",
        "remote_lock": "remote_lock",
    }
    operation = op_map.get(tool_name)
    if not operation:
        return {"error": f"Unknown Intune tool: {tool_name}"}

    ctx = ConnectorContext(gsap_context_id=args.get("ac_id", "mcp-session"))
    result = await intune.invoke(operation, args, ctx)
    if result.success:
        return {"data": result.data, "lineage_cid": result.lineage_cid}
    return {"error": result.error}


# ── Tool Dispatch ────────────────────────────────────────────────


async def _dispatch_tool(request: Request, tool_name: str, arguments: dict) -> dict:
    handlers = {
        "request_ac": lambda: _handle_request_ac(request, arguments),
        "validate_ac": lambda: _handle_validate_ac(arguments),
        "post_cr": lambda: _handle_post_cr(request, arguments),
        "query_accord": lambda: _handle_query_accord(arguments),
        "request_delegation": lambda: _proxy_to_llm_broker(tool_name, arguments, request),
        "revoke_delegation": lambda: _proxy_to_llm_broker(tool_name, arguments, request),
        "get_delegation": lambda: _proxy_to_llm_broker(tool_name, arguments, request),
        "list_agents": lambda: _proxy_to_llm_broker(tool_name, arguments, request),
        "get_posture": lambda: _handle_get_posture(arguments),
        "check_operation": lambda: _handle_check_operation(arguments, request),
        "session_info": lambda: _handle_session_info(request),
        "list_devices": lambda: _handle_intune_tool(tool_name, arguments),
        "get_device_compliance": lambda: _handle_intune_tool(tool_name, arguments),
        "sync_device": lambda: _handle_intune_tool(tool_name, arguments),
        "remote_lock": lambda: _handle_intune_tool(tool_name, arguments),
    }

    handler = handlers.get(tool_name)
    if not handler:
        return {"error": f"Unknown tool: {tool_name}"}

    return await handler()


# ── JSON-RPC Route ───────────────────────────────────────────────


def _error(code: int, message: str, req_id):
    return {"jsonrpc": "2.0", "error": {"code": code, "message": message}, "id": req_id}


def _success(result, req_id):
    return {"jsonrpc": "2.0", "result": result, "id": req_id}


@router.post("/mcp")
async def mcp_endpoint(request: Request):
    """MCP JSON-RPC 2.0 endpoint — governance primitives as tools."""
    try:
        body = await request.json()
    except Exception:
        return JSONResponse(_error(-32700, "Parse error", None), status_code=400)

    method = body.get("method", "")
    params = body.get("params", {})
    req_id = body.get("id")

    if method == "initialize":
        return JSONResponse(_success({
            "protocolVersion": MCP_PROTOCOL_VERSION,
            "serverInfo": {"name": SERVER_NAME, "version": SERVER_VERSION},
            "capabilities": {"tools": {"listChanged": False}},
        }, req_id))

    elif method == "tools/list":
        return JSONResponse(_success({"tools": TOOLS}, req_id))

    elif method == "tools/call":
        tool_name = params.get("name", "")
        arguments = params.get("arguments", {})
        if not tool_name:
            return JSONResponse(_error(-32602, "Missing tool name", req_id))

        result = await _dispatch_tool(request, tool_name, arguments)
        is_error = "error" in result

        # Chronicle: record MCP tool call
        await chronicle.emit("MCP_TOOL_CALL", {
            "event_code": "0x3020",
            "tool": tool_name,
            "principal": _extract_principal(request),
            "outcome": "error" if is_error else "success",
        })

        return JSONResponse(_success({
            "content": [{"type": "text", "text": json.dumps(result, default=str)}],
            "isError": is_error,
        }, req_id))

    else:
        return JSONResponse(_error(-32601, f"Unknown method: {method}", req_id))
