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

from __future__ import annotations

import json
import os
import sys

import httpx

__version__ = "0.2.3"

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



def _resolve_api_key() -> str:
    """Resolve API key."""
    global API_KEY
    if API_KEY and API_KEY.startswith("kvrs_"):
        return API_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

    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 ""



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



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 release token to verify",
                },
                "token_preimage": {
                    "type": "string",
                    "description": "The token preimage (provided alongside the release token)",
                },
            },
            "required": ["release_token", "token_preimage"],
        },
    },
]



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",
    },
]



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,
            },
        ],
    },
]



def handle_tool_call(name: str, arguments: dict) -> dict:
    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", "")
        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"}
            if resp.status_code >= 400:
                return {"error": f"Request failed (HTTP {resp.status_code})"}
            return resp.json()
    elif name == "governance_verify_token":
        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}"}



def handle_resource_read(uri: str) -> dict:
    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}"}



def handle_prompt_get(name: str, arguments: dict) -> dict:
    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}"}



def send_response(id: int | str | None, result: dict) -> None:
    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:
    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:
    # 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("Kevros MCP Server started\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":
            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()
