#!/usr/bin/env python3
"""
Kevros Governance MCP Server

Model Context Protocol server that exposes the Kevros A2A Governance Gateway
as tools for Claude Code, Cursor, Windsurf, and any MCP-compatible client.

Setup (Claude Code):
    Add to ~/.claude/settings.json:
    {
      "mcpServers": {
        "kevros-governance": {
          "command": "python3",
          "args": ["/path/to/kevros_mcp_server.py"],
          "env": {
            "KEVROS_API_KEY": "kvrs_..."
          }
        }
      }
    }

Setup (Cursor):
    Add to .cursor/mcp.json:
    {
      "mcpServers": {
        "kevros-governance": {
          "command": "python3",
          "args": ["/path/to/kevros_mcp_server.py"],
          "env": {
            "KEVROS_API_KEY": "kvrs_..."
          }
        }
      }
    }

Or install via pip and run:
    pip install kevros-governance[mcp]
    KEVROS_API_KEY=kvrs_... python -m kevros_mcp_server
"""

from __future__ import annotations

import json
import os
import sys

import httpx

GATEWAY_URL = os.environ.get("KEVROS_GATEWAY_URL", "https://governance.taskhawktech.com")
API_KEY = os.environ.get("KEVROS_API_KEY", "")
TIMEOUT = 30.0


def _headers() -> dict:
    return {"X-API-Key": API_KEY, "Content-Type": "application/json"}


def _post(path: str, body: dict) -> dict:
    with httpx.Client(base_url=GATEWAY_URL, headers=_headers(), timeout=TIMEOUT) as c:
        resp = c.post(path, json=body)
        if resp.status_code >= 400:
            try:
                detail = resp.json().get("detail", resp.text)
            except Exception:
                detail = resp.text
            return {"error": f"HTTP {resp.status_code}: {detail}"}
        return resp.json()


def _get(path: str) -> dict:
    with httpx.Client(base_url=GATEWAY_URL, headers=_headers(), timeout=TIMEOUT) as c:
        resp = c.get(path)
        if resp.status_code >= 400:
            return {"error": f"HTTP {resp.status_code}: {resp.text}"}
        return resp.json()


# ---------------------------------------------------------------------------
# MCP protocol implementation (stdio JSON-RPC)
# ---------------------------------------------------------------------------

TOOLS = [
    {
        "name": "governance_verify",
        "description": (
            "Verify an action against policy bounds before executing it. "
            "Returns ALLOW (proceed), CLAMP (proceed with modified values), or DENY (stop). "
            "Every verification is recorded in a hash-chained provenance ledger. "
            "Cost: $0.01 per call."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "action_type": {
                    "type": "string",
                    "description": "Type of action (e.g., 'trade', 'api_call', 'deploy', 'send_email')",
                },
                "action_payload": {
                    "type": "object",
                    "description": "The action details to verify (contents depend on action_type)",
                },
                "agent_id": {
                    "type": "string",
                    "description": "Your agent identifier",
                },
                "policy_context": {
                    "type": "object",
                    "description": "Optional policy constraints. Use max_values:{key:number} to set bounds, forbidden_keys:[...] to block fields.",
                },
            },
            "required": ["action_type", "action_payload", "agent_id"],
        },
    },
    {
        "name": "governance_attest",
        "description": (
            "Create a hash-chained provenance record for an action you've taken. "
            "Each attestation extends the append-only evidence chain. "
            "The hash can be independently verified by any third party. "
            "Cost: $0.02 per call."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "agent_id": {
                    "type": "string",
                    "description": "Your agent identifier",
                },
                "action_description": {
                    "type": "string",
                    "description": "What you did (human-readable)",
                },
                "action_payload": {
                    "type": "object",
                    "description": "Full action details for the provenance record",
                },
                "context": {
                    "type": "object",
                    "description": "Optional additional context (environment, state, etc.)",
                },
            },
            "required": ["agent_id", "action_description", "action_payload"],
        },
    },
    {
        "name": "governance_bind",
        "description": (
            "Declare an intent and cryptographically bind it to a command. "
            "Proves that the command was issued in service of the declared intent. "
            "Use governance_verify_outcome after execution to close the loop. "
            "Cost: $0.02 per call."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "agent_id": {
                    "type": "string",
                    "description": "Your agent identifier",
                },
                "intent_type": {
                    "type": "string",
                    "enum": [
                        "NAVIGATION", "MANIPULATION", "SENSING", "COMMUNICATION",
                        "MAINTENANCE", "EMERGENCY", "OPERATOR_COMMAND",
                        "AI_GENERATED", "AUTOMATED",
                    ],
                    "description": "Category of intent",
                },
                "intent_description": {
                    "type": "string",
                    "description": "What you intend to accomplish",
                },
                "command_payload": {
                    "type": "object",
                    "description": "The command being bound to this intent",
                },
                "goal_state": {
                    "type": "object",
                    "description": "Target state to verify against later with verify_outcome",
                },
                "intent_source": {
                    "type": "string",
                    "enum": [
                        "HUMAN_OPERATOR", "AI_PLANNER", "MISSION_SCRIPT",
                        "REMOTE_API", "SENSOR_TRIGGER", "INTERNAL",
                    ],
                    "default": "AI_PLANNER",
                    "description": "Who/what generated the intent",
                },
            },
            "required": ["agent_id", "intent_type", "intent_description", "command_payload"],
        },
    },
    {
        "name": "governance_verify_outcome",
        "description": (
            "Verify that an executed action achieved its declared intent. "
            "Closes the loop: intent -> command -> action -> outcome -> verification. "
            "Free (included with bind)."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "agent_id": {"type": "string", "description": "Your agent identifier"},
                "intent_id": {"type": "string", "description": "The intent_id from governance_bind"},
                "binding_id": {"type": "string", "description": "The binding_id from governance_bind"},
                "actual_state": {"type": "object", "description": "The state after action execution"},
                "tolerance": {
                    "type": "number",
                    "description": "Acceptable deviation from goal (0=exact, 1=any). Default 0.1",
                    "default": 0.1,
                },
            },
            "required": ["agent_id", "intent_id", "binding_id", "actual_state"],
        },
    },
    {
        "name": "governance_bundle",
        "description": (
            "Generate a certifier-grade compliance evidence bundle. "
            "Contains hash-chained provenance, intent bindings, PQC attestations, "
            "and verification instructions. Independently verifiable without Kevros access. "
            "Cost: $0.25 per call."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "agent_id": {"type": "string", "description": "Your agent identifier"},
                "time_range_start": {
                    "type": "string",
                    "description": "ISO 8601 start of time range (optional)",
                },
                "time_range_end": {
                    "type": "string",
                    "description": "ISO 8601 end of time range (optional)",
                },
            },
            "required": ["agent_id"],
        },
    },
    {
        "name": "governance_health",
        "description": "Check the governance gateway health status. Free.",
        "inputSchema": {
            "type": "object",
            "properties": {},
        },
    },
]


def handle_tool_call(name: str, arguments: dict) -> dict:
    """Route MCP tool calls to gateway HTTP endpoints."""
    if name == "governance_verify":
        return _post("/governance/verify", arguments)
    elif name == "governance_attest":
        return _post("/governance/attest", arguments)
    elif name == "governance_bind":
        return _post("/governance/bind", arguments)
    elif name == "governance_verify_outcome":
        return _post("/governance/verify-outcome", arguments)
    elif name == "governance_bundle":
        body = {"agent_id": arguments["agent_id"]}
        if "time_range_start" in arguments:
            body["time_range_start"] = arguments["time_range_start"]
        if "time_range_end" in arguments:
            body["time_range_end"] = arguments["time_range_end"]
        body["include_intent_chains"] = True
        body["include_pqc_signatures"] = True
        body["include_verification_instructions"] = True
        return _post("/governance/bundle", body)
    elif name == "governance_health":
        return _get("/governance/health")
    else:
        return {"error": f"Unknown tool: {name}"}


def send_response(id: int | str | None, result: dict) -> None:
    """Send JSON-RPC response to stdout."""
    msg = {"jsonrpc": "2.0", "id": id, "result": result}
    out = json.dumps(msg)
    sys.stdout.write(out + "\n")
    sys.stdout.flush()


def send_error(id: int | str | None, code: int, message: str) -> None:
    """Send JSON-RPC error to stdout."""
    msg = {"jsonrpc": "2.0", "id": id, "error": {"code": code, "message": message}}
    out = json.dumps(msg)
    sys.stdout.write(out + "\n")
    sys.stdout.flush()


def main() -> None:
    if not API_KEY:
        sys.stderr.write(
            "ERROR: KEVROS_API_KEY environment variable is required.\n"
            "Get your API key at https://governance.taskhawktech.com\n"
        )
        sys.exit(1)

    sys.stderr.write(f"Kevros Governance MCP Server started (gateway: {GATEWAY_URL})\n")

    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue

        try:
            msg = json.loads(line)
        except json.JSONDecodeError:
            send_error(None, -32700, "Parse error")
            continue

        method = msg.get("method", "")
        id_ = msg.get("id")
        params = msg.get("params", {})

        if method == "initialize":
            send_response(id_, {
                "protocolVersion": "2024-11-05",
                "capabilities": {"tools": {"listChanged": False}},
                "serverInfo": {
                    "name": "kevros-governance",
                    "version": "0.1.0",
                },
            })

        elif method == "notifications/initialized":
            # Client acknowledged initialization — no response needed
            pass

        elif method == "tools/list":
            send_response(id_, {"tools": TOOLS})

        elif method == "tools/call":
            tool_name = params.get("name", "")
            arguments = params.get("arguments", {})
            result = handle_tool_call(tool_name, arguments)
            send_response(id_, {
                "content": [
                    {"type": "text", "text": json.dumps(result, indent=2)}
                ],
            })

        elif method == "ping":
            send_response(id_, {})

        else:
            send_error(id_, -32601, f"Method not found: {method}")


if __name__ == "__main__":
    main()
