#!/usr/bin/env python3
"""MCP Discoverability Validation Tool for external servers (Issue #155)."""
import os
import sys
import json
import argparse
import subprocess

EXPECTED_JENKINS_TOOLS = {
    "jenkins_whoami",
    "jenkins_list_jobs",
    "jenkins_latest_build",
    "jenkins_build_status",
    "jenkins_get_build",
}

EXPECTED_GLITCHTIP_TOOLS = {
    "glitchtip_whoami",
    "glitchtip_list_projects",
    "glitchtip_list_unresolved",
    "glitchtip_get_issue",
    "glitchtip_recent_events",
    "glitchtip_search",
}

RELOAD_INSTRUCTIONS = """
=== MCP CLIENT RELOAD/RECONNECT RUNBOOK ===
After registering or changing external MCP servers, reload your client to discover the new tools:
- Codex: Click 'Reload Developer Tools' or restart the editor.
- Gemini / Grok / ChatGPT Desktop: Restart the client or run the reload slash command if available.
- Claude Desktop: Use 'Developer -> Reload' or restart the app.
- General MCP Clients: Restart the process or reload the server config.
===========================================
"""

def parse_gitea_mcp_config(path):
    if not path or not os.path.exists(path):
        return {}
    with open(path, "r", encoding="utf-8") as fh:
        try:
            return json.load(fh)
        except Exception:
            return {}

def read_json_rpc_response(proc, req_id):
    import time
    start_time = time.time()
    while time.time() - start_time < 5.0:
        line = proc.stdout.readline()
        if not line:
            break
        try:
            data = json.loads(line)
            if data.get("id") == req_id:
                return data
        except json.JSONDecodeError:
            continue
    return None

def query_live_tools(command, args, env):
    run_env = os.environ.copy()
    if env:
        run_env.update(env)

    proc = subprocess.Popen(
        [command] + args,
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        env=run_env,
        text=True,
        bufsize=1
    )

    tools = []
    try:
        # 1. Send initialize
        init_req = {
            "jsonrpc": "2.0",
            "id": 1,
            "method": "initialize",
            "params": {
                "protocolVersion": "2024-11-05",
                "capabilities": {},
                "clientInfo": {"name": "mcp-discoverability-check", "version": "1.0.0"}
            }
        }
        proc.stdin.write(json.dumps(init_req) + "\n")
        proc.stdin.flush()

        # Read init response
        init_resp = read_json_rpc_response(proc, 1)
        if init_resp:
            # Send initialized notification
            init_notif = {
                "jsonrpc": "2.0",
                "method": "notifications/initialized"
            }
            proc.stdin.write(json.dumps(init_notif) + "\n")
            proc.stdin.flush()

            # 2. Send tools/list
            tools_req = {
                "jsonrpc": "2.0",
                "id": 2,
                "method": "tools/list",
                "params": {}
            }
            proc.stdin.write(json.dumps(tools_req) + "\n")
            proc.stdin.flush()

            tools_resp = read_json_rpc_response(proc, 2)
            if tools_resp and "result" in tools_resp and "tools" in tools_resp["result"]:
                for t in tools_resp["result"]["tools"]:
                    tools.append(t["name"])
    except Exception as e:
        print(f"Error querying live tools: {e}", file=sys.stderr)
    finally:
        proc.terminate()
        try:
            proc.wait(timeout=2)
        except Exception:
            proc.kill()

    return set(tools)

def validate_mcp_client_config(client_config_path, gitea_config_path=None, live=False):
    if not client_config_path or not os.path.exists(client_config_path):
        print(f"SKIPPED: MCP client config not found at '{client_config_path}'", file=sys.stderr)
        return True

    with open(client_config_path, "r", encoding="utf-8") as fh:
        try:
            config_data = json.load(fh)
        except Exception as e:
            print(f"Error parsing client config: {e}", file=sys.stderr)
            return False

    mcp_servers = config_data.get("mcpServers", {})

    # Check for stale server names
    stale_names = {"jenkins-readonly", "glitchtip-readonly"}
    for name in mcp_servers:
        if name in stale_names:
            print(f"ERROR: Stale server name '{name}' configured. Use canonical names 'jenkins-mcp' or 'glitchtip-mcp'.", file=sys.stderr)
            return False

    gitea_data = parse_gitea_mcp_config(gitea_config_path)
    enabled_services = set()
    contexts = gitea_data.get("contexts", {})

    profile_name = os.environ.get("GITEA_MCP_PROFILE")
    if profile_name and "profiles" in gitea_data:
        profile = gitea_data["profiles"].get(profile_name)
        if profile and "context" in profile:
            ctx_name = profile["context"]
            ctx = contexts.get(ctx_name, {})
            if ctx.get("enabled"):
                services = ctx.get("services", {})
                for s_name, s_data in services.items():
                    if s_data.get("enabled"):
                        enabled_services.add(s_name)
    else:
        for ctx_name, ctx in contexts.items():
            if ctx.get("enabled"):
                services = ctx.get("services", {})
                for s_name, s_data in services.items():
                    if s_data.get("enabled"):
                        enabled_services.add(s_name)

    if not enabled_services:
        print("No external services enabled in Gitea contexts. Discoverability check complete.", file=sys.stderr)
        return True

    success = True
    for service in enabled_services:
        canonical_name = f"{service}-mcp"
        if canonical_name not in mcp_servers:
            stale_match = f"{service}-readonly"
            if stale_match in mcp_servers:
                print(f"ERROR: Server '{canonical_name}' references stale name '{stale_match}' (fail closed).", file=sys.stderr)
                success = False
                continue
            print(f"ERROR: Enabled service '{service}' is not registered under canonical name '{canonical_name}' in client config.", file=sys.stderr)
            success = False
            continue

        server_conf = mcp_servers[canonical_name]
        command = server_conf.get("command")
        args = server_conf.get("args") or []
        env = server_conf.get("env") or {}

        if not command:
            print(f"ERROR: Server '{canonical_name}' has no command configured.", file=sys.stderr)
            success = False
            continue

        expected_module = f"{service}_mcp"
        if "-m" not in args or expected_module not in args:
            print(f"ERROR: Server '{canonical_name}' args do not point to expected module '{expected_module}' (args: {args}).", file=sys.stderr)
            success = False
            continue

        profile_var = f"{service.upper()}_MCP_PROFILE"
        config_var = f"{service.upper()}_MCP_CONFIG"
        if profile_var not in env or config_var not in env:
            print(f"ERROR: Server '{canonical_name}' env is missing required variables '{profile_var}' or '{config_var}'.", file=sys.stderr)
            success = False
            continue

        if live:
            tools = query_live_tools(command, args, env)
            if not tools:
                print("SKIPPED: server enabled but no usable tools visible", file=sys.stdout)
                success = False
                continue

            expected = EXPECTED_JENKINS_TOOLS if service == "jenkins" else EXPECTED_GLITCHTIP_TOOLS
            missing = expected - tools
            if missing:
                print(f"ERROR: Server '{canonical_name}' is missing expected tools: {', '.join(missing)}", file=sys.stderr)
                success = False
            else:
                print(f"SUCCESS: Server '{canonical_name}' discoverability verified.", file=sys.stderr)
        else:
            print(f"SUCCESS: Server '{canonical_name}' static registration verified.", file=sys.stderr)

    return success

def main():
    parser = argparse.ArgumentParser(description="MCP client registration discoverability checks.")
    parser.add_argument("--client-config", help="Path to MCP client config JSON file.")
    parser.add_argument("--gitea-config", help="Path to Gitea MCP config JSON file.")
    parser.add_argument("--live", action="store_true", help="Perform live stdio checks on configured servers.")
    parser.add_argument("--runbook", action="store_true", help="Print reload/reconnect guide runbook instructions.")

    args = parser.parse_args()

    if args.runbook:
        print(RELOAD_INSTRUCTIONS)
        return 0

    if not args.client_config:
        print("ERROR: --client-config must be specified.", file=sys.stderr)
        return 2

    ok = validate_mcp_client_config(
        client_config_path=args.client_config,
        gitea_config_path=args.gitea_config,
        live=args.live
    )

    if not ok:
        print(RELOAD_INSTRUCTIONS, file=sys.stderr)
        return 1
    return 0

if __name__ == "__main__":
    sys.exit(main())
