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

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

Setup (Claude Code — recommended):
    pip install kevros
    claude mcp add kevros -- kevros-mcp

Setup (manual config):
    Add to ~/.claude/settings.json:
    {
      "mcpServers": {
        "kevros": {
          "command": "kevros-mcp"
        }
      }
    }

Setup (Cursor):
    Add to .cursor/mcp.json:
    {
      "mcpServers": {
        "kevros": {
          "command": "kevros-mcp"
        }
      }
    }
"""

from __future__ import annotations

import json
import os
import sys

import httpx

__version__ = "0.2.2"

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


# ---------------------------------------------------------------------------
# API key resolution (env > cached > auto-signup)
# ---------------------------------------------------------------------------

def _resolve_api_key() -> str:
    """Resolve API key: env > cached > auto-signup."""
    global API_KEY
    if API_KEY and API_KEY.startswith("kvrs_"):
        return API_KEY

    # Try cached key
    from pathlib import Path
    key_file = Path.home() / ".kevros" / "api_key"
    if key_file.exists():
        key = key_file.read_text().strip()
        if key.startswith("kvrs_"):
            API_KEY = key
            return API_KEY

    # Auto-signup
    agent_id = os.environ.get("KEVROS_AGENT_ID", "mcp-agent")
    sys.stderr.write(f"No KEVROS_API_KEY found — auto-signing up as '{agent_id}'...\n")
    resp = httpx.post(
        f"{GATEWAY_URL}/signup",
        json={"agent_id": agent_id},
        timeout=15.0,
    )
    if resp.status_code == 200:
        data = resp.json()
        API_KEY = data["api_key"]
        key_dir = Path.home() / ".kevros"
        key_dir.mkdir(parents=True, exist_ok=True)
        key_file.write_text(API_KEY)
        try:
            key_file.chmod(0o600)
        except OSError:
            pass
        sys.stderr.write(
            f"Kevros free tier activated: {data.get('monthly_limit', 100)} calls/mo. "
            f"Key cached at {key_file}\n"
        )
        return API_KEY

    sys.stderr.write(f"Auto-signup failed (HTTP {resp.status_code}). Check network and try again.\n")
    return ""


# ---------------------------------------------------------------------------
# HTTP helpers
# ---------------------------------------------------------------------------

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 == 402:
            try:
                data = resp.json()
            except Exception:
                data = {}
            return {
                "error": "Free tier limit exceeded",
                "detail": data.get("detail", "Monthly call limit reached"),
                "upgrade_url": data.get("upgrade_url", "https://governance.taskhawktech.com/upgrade"),
                "tiers": data.get("tiers", {}),
            }
        if resp.status_code >= 400:
            return {"error": f"Request failed (HTTP {resp.status_code})"}
        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"Request failed (HTTP {resp.status_code})"}
        return resp.json()


# ---------------------------------------------------------------------------
# MCP Tools (9)
# ---------------------------------------------------------------------------

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": {},
        },
    },
    {
        "name": "kevros_status",
        "description": (
            "Check your Kevros identity trust status: calls used, calls remaining, "
            "tier, rate limits. Free."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {},
        },
    },
    {
        "name": "governance_check_peer",
        "description": (
            "Check another agent's trust score and provenance history. "
            "Returns chain_length, attestation_count, trust_score. "
            "Free, no API key needed."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "agent_id": {
                    "type": "string",
                    "description": "The agent ID to look up",
                },
            },
            "required": ["agent_id"],
        },
    },
    {
        "name": "governance_verify_token",
        "description": (
            "Verify a release token from another agent. Confirms the token "
            "is authentic and was issued by the Kevros gateway. "
            "Free, no API key needed."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "release_token": {
                    "type": "string",
                    "description": "The HMAC release token (hex digest)",
                },
                "token_preimage": {
                    "type": "string",
                    "description": "The token preimage (provided alongside the release token)",
                },
            },
            "required": ["release_token", "token_preimage"],
        },
    },
]


# ---------------------------------------------------------------------------
# MCP Resources (2)
# ---------------------------------------------------------------------------

RESOURCES = [
    {
        "uri": "kevros://agent-card",
        "name": "Agent Card",
        "description": "A2A agent card for the Kevros Governance Gateway. Used for agent-to-agent discovery.",
        "mimeType": "application/json",
    },
    {
        "uri": "kevros://trust-status",
        "name": "Trust Status",
        "description": "Your identity trust status: calls used, tier, limits, rate limits.",
        "mimeType": "application/json",
    },
]


# ---------------------------------------------------------------------------
# MCP Prompts (2)
# ---------------------------------------------------------------------------

PROMPTS = [
    {
        "name": "verify-before-act",
        "description": "Pre-flight governance check. Generates a verify request for an action before executing it.",
        "arguments": [
            {
                "name": "action_type",
                "description": "Type of action to verify (e.g., 'trade', 'deploy', 'api_call')",
                "required": True,
            },
            {
                "name": "action_description",
                "description": "Human-readable description of what the action does",
                "required": True,
            },
            {
                "name": "agent_id",
                "description": "Your agent identifier (defaults to 'mcp-agent')",
                "required": False,
            },
        ],
    },
    {
        "name": "governance-audit",
        "description": "Generate a compliance audit bundle for an agent's actions over a time range.",
        "arguments": [
            {
                "name": "agent_id",
                "description": "The agent ID to audit",
                "required": True,
            },
            {
                "name": "time_range_start",
                "description": "ISO 8601 start of audit window (e.g., '2026-02-01T00:00:00Z')",
                "required": False,
            },
            {
                "name": "time_range_end",
                "description": "ISO 8601 end of audit window (e.g., '2026-02-28T23:59:59Z')",
                "required": False,
            },
        ],
    },
]


# ---------------------------------------------------------------------------
# Tool call handler
# ---------------------------------------------------------------------------

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")
    elif name == "kevros_status":
        return _get("/signup/status")
    elif name == "governance_check_peer":
        peer_id = arguments.get("agent_id", "")
        # Public endpoint — no API key needed
        with httpx.Client(base_url=GATEWAY_URL, timeout=TIMEOUT) as c:
            resp = c.get(f"/governance/reputation/{peer_id}")
            if resp.status_code == 404:
                return {"error": "not_found", "detail": f"No records for agent: {peer_id}"}
            if resp.status_code >= 400:
                return {"error": f"Request failed (HTTP {resp.status_code})"}
            return resp.json()
    elif name == "governance_verify_token":
        # Public endpoint — no API key needed
        with httpx.Client(base_url=GATEWAY_URL, timeout=TIMEOUT) as c:
            resp = c.post("/governance/verify-token", json={
                "release_token": arguments.get("release_token", ""),
                "token_preimage": arguments.get("token_preimage", ""),
            })
            if resp.status_code >= 400:
                return {"error": f"Request failed (HTTP {resp.status_code})"}
            return resp.json()
    else:
        return {"error": f"Unknown tool: {name}"}


# ---------------------------------------------------------------------------
# Resource handler
# ---------------------------------------------------------------------------

def handle_resource_read(uri: str) -> dict:
    """Fetch resource content by URI."""
    if uri == "kevros://agent-card":
        data = _get("/.well-known/agent.json")
        return {
            "contents": [
                {
                    "uri": uri,
                    "mimeType": "application/json",
                    "text": json.dumps(data, indent=2),
                }
            ]
        }
    elif uri == "kevros://trust-status":
        data = _get("/signup/status")
        return {
            "contents": [
                {
                    "uri": uri,
                    "mimeType": "application/json",
                    "text": json.dumps(data, indent=2),
                }
            ]
        }
    else:
        return {"contents": [], "error": f"Unknown resource: {uri}"}


# ---------------------------------------------------------------------------
# Prompt handler
# ---------------------------------------------------------------------------

def handle_prompt_get(name: str, arguments: dict) -> dict:
    """Build prompt messages for a given prompt name."""
    if name == "verify-before-act":
        action_type = arguments.get("action_type", "unknown")
        action_description = arguments.get("action_description", "")
        agent_id = arguments.get("agent_id", "mcp-agent")
        return {
            "description": "Pre-flight governance check",
            "messages": [
                {
                    "role": "user",
                    "content": {
                        "type": "text",
                        "text": (
                            f"Before executing this action, verify it with Kevros governance.\n\n"
                            f"Action type: {action_type}\n"
                            f"Description: {action_description}\n"
                            f"Agent ID: {agent_id}\n\n"
                            f"Use the governance_verify tool with:\n"
                            f"- action_type: \"{action_type}\"\n"
                            f"- action_payload: (construct from the description above)\n"
                            f"- agent_id: \"{agent_id}\"\n\n"
                            f"If the result is DENY, do not proceed. "
                            f"If CLAMP, use the modified values. "
                            f"If ALLOW, proceed and then attest the outcome."
                        ),
                    },
                }
            ],
        }
    elif name == "governance-audit":
        agent_id = arguments.get("agent_id", "")
        time_start = arguments.get("time_range_start", "")
        time_end = arguments.get("time_range_end", "")
        time_clause = ""
        if time_start or time_end:
            time_clause = f"\nTime range: {time_start or '(open)'} to {time_end or '(open)'}"
        return {
            "description": "Generate compliance audit bundle",
            "messages": [
                {
                    "role": "user",
                    "content": {
                        "type": "text",
                        "text": (
                            f"Generate a compliance audit for agent \"{agent_id}\".{time_clause}\n\n"
                            f"Steps:\n"
                            f"1. Use governance_bundle to generate the evidence bundle\n"
                            f"2. Check chain_integrity is True\n"
                            f"3. Report: record_count, chain_integrity, bundle_hash, "
                            f"PQC signature status\n"
                            f"4. If any intent chains are present, summarize their outcomes\n"
                            f"5. Provide the bundle_id for reference"
                        ),
                    },
                }
            ],
        }
    else:
        return {"messages": [], "error": f"Unknown prompt: {name}"}


# ---------------------------------------------------------------------------
# JSON-RPC helpers
# ---------------------------------------------------------------------------

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()


# ---------------------------------------------------------------------------
# Main loop (stdio JSON-RPC)
# ---------------------------------------------------------------------------

def main() -> None:
    # Auto-signup if no key configured
    _resolve_api_key()
    if not API_KEY:
        sys.stderr.write(
            "ERROR: Could not obtain Kevros API key.\n"
            "Set KEVROS_API_KEY env var or ensure network access to "
            "https://governance.taskhawktech.com/signup\n"
        )
        sys.exit(1)

    sys.stderr.write(f"Kevros Identity Trust MCP Server v{__version__} 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},
                    "resources": {"subscribe": False, "listChanged": False},
                    "prompts": {"listChanged": False},
                },
                "serverInfo": {
                    "name": "kevros-governance",
                    "version": __version__,
                },
            })

        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 == "resources/list":
            send_response(id_, {"resources": RESOURCES})

        elif method == "resources/read":
            uri = params.get("uri", "")
            result = handle_resource_read(uri)
            send_response(id_, result)

        elif method == "prompts/list":
            send_response(id_, {"prompts": PROMPTS})

        elif method == "prompts/get":
            prompt_name = params.get("name", "")
            arguments = params.get("arguments", {})
            result = handle_prompt_get(prompt_name, arguments)
            send_response(id_, result)

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

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


if __name__ == "__main__":
    main()
