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

Exposes a single MCP tool "unreal" that forwards JSON commands to the
UE Wingman TCP server in the Unreal Editor.
"""

import sys
import json
import socket

HOST = "localhost"
PORT = 9851
CONNECT_TIMEOUT = 2
READ_TIMEOUT = 30

TOOL_DESCRIPTION = (
    "Send a command to the Unreal Editor's UE Wingman plugin. "
    "The 'command' field specifies which operation to perform; "
    "additional fields are command-specific parameters. "
    'Use {"command": "Documentation_Manual"} 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) + "\0"
    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()
        return {"error": "Lost connection to Unreal Editor."}


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


def parse_editor_response(result):
    """Parse and validate a raw editor response into an MCP content list.

    MCP expects `content` to be a list of objects, each with at least a
    string "type" field (e.g. {"type": "text", "text": "..."}). Anything
    else is replaced with a single error item so the client sees a clear
    message instead of a schema violation.
    """
    try:
        parsed = json.loads(result)
    except json.JSONDecodeError:
        return [{"type": "text", "text": "Malformed response from editor: invalid JSON."}]
    if not isinstance(parsed, list):
        return [{"type": "text", "text": "Malformed response from editor: expected a list."}]
    for item in parsed:
        if not isinstance(item, dict):
            return [{"type": "text", "text": "Malformed response from editor: list item is not an object."}]
        if not isinstance(item.get("type"), str):
            return [{"type": "text", "text": "Malformed response from editor: item missing string 'type' field."}]
    return parsed


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": "ue-wingman", "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:
            content = [{"type": "text", "text": result["error"]}]
        else:
            content = parse_editor_response(result)
        return make_jsonrpc(msg_id, {
            "content": content,
        })

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