#!/usr/bin/env python3
"""
MCP stdio-to-TCP bridge for BlueprintMCP.

Exposes a single MCP tool "unreal" that forwards JSON commands to the
BlueprintMCP TCP server in the Unreal Editor. The tool list is static,
so it works regardless of whether the editor is running at startup.
"""

import sys
import json
import socket

HOST = "localhost"
PORT = 9847
CONNECT_TIMEOUT = 2
READ_TIMEOUT = 120

TOOL_DESCRIPTION = (
    "Send a command to the Unreal Editor's BlueprintMCP plugin. "
    "The 'command' field specifies which operation to perform; "
    "additional fields are command-specific parameters. "
    'Use {"command": "UserManual"} to get an overview. '
    "If the editor is not running, the call will return an error; "
    "just ask the user to start the editor and try again."
)

TOOL_SCHEMA = {
    "name": "unreal",
    "description": TOOL_DESCRIPTION,
    "inputSchema": {
        "type": "object",
        "properties": {
            "command": {"type": "string", "description": "The command to execute"},
        },
        "required": ["command"],
        "additionalProperties": True,
    },
}

sock = None


def connect():
    """Try to connect to the editor. Returns True on success."""
    global sock
    if sock is not None:
        return True
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.settimeout(CONNECT_TIMEOUT)
        s.connect((HOST, PORT))
        s.settimeout(READ_TIMEOUT)
        sock = s
        return True
    except (ConnectionRefusedError, socket.timeout, OSError):
        return False


def disconnect():
    global sock
    if sock is not None:
        try:
            sock.close()
        except Exception:
            pass
        sock = None


def send_and_receive(message):
    """Send a JSON message to the editor and return the null-terminated response."""
    data = json.dumps(message) + "\n"
    sock.sendall(data.encode())

    result = b""
    while True:
        chunk = sock.recv(65536)
        if not chunk:
            raise ConnectionError("Connection closed")
        result += chunk
        if b"\0" in result:
            break

    return result[:result.index(b"\0")].decode()


def forward_to_editor(arguments):
    """Forward arguments to the editor, return the result dict."""
    if not connect():
        return {"error": "Unreal Editor is not running. Start the editor and try again."}
    try:
        return send_and_receive(arguments)
    except Exception:
        disconnect()
        # Retry once in case the connection was stale
        if connect():
            try:
                return send_and_receive(arguments)
            except Exception:
                disconnect()
        return {"error": "Lost connection to Unreal Editor."}


def make_jsonrpc(msg_id, result):
    return {"jsonrpc": "2.0", "id": msg_id, "result": result}


def handle_message(msg):
    """Handle one JSON-RPC message from Claude Code."""
    msg_id = msg.get("id")
    method = msg.get("method", "")

    # Notifications don't get responses
    if msg_id is None:
        return None

    if method == "initialize":
        return make_jsonrpc(msg_id, {
            "protocolVersion": "2024-11-05",
            "capabilities": {"tools": {}},
            "serverInfo": {"name": "blueprint-mcp", "version": "1.0.0"},
        })

    if method == "tools/list":
        return make_jsonrpc(msg_id, {"tools": [TOOL_SCHEMA]})

    if method == "tools/call":
        params = msg.get("params", {})
        arguments = params.get("arguments", {})
        result = forward_to_editor(arguments)
        if isinstance(result, dict) and "error" in result:
            text = result["error"]
        else:
            text = result
        return make_jsonrpc(msg_id, {
            "content": [{"type": "text", "text": text}],
        })

    return {
        "jsonrpc": "2.0",
        "id": msg_id,
        "error": {"code": -32601, "message": f"Method not found: {method}"},
    }


def main():
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            msg = json.loads(line)
        except json.JSONDecodeError:
            continue

        response = handle_message(msg)
        if response is not None:
            sys.stdout.write(json.dumps(response) + "\n")
            sys.stdout.flush()


if __name__ == "__main__":
    main()