@tool
extends Node

## Multi-connection WebSocket client.
## Connects to multiple Node.js MCP server instances on ports 6505-6514.
## Each Claude Code session gets its own port; Godot talks to all of them.
## Ports 6505-6509: MCP servers (stdio), 6510-6514: CLI tool connections.

signal client_connected()
signal client_disconnected()
signal message_received(text: String)
signal command_executed(method: String, success: bool)
signal command_completed(method: String, success: bool, response: String, source_port: int)

var command_router: Node

const BASE_PORT := 6505
const MAX_PORT := 6514
const RECONNECT_INTERVAL := 3.0
const BUFFER_SIZE := 16 * 1024 * 1024  # 16MB
const PING_INTERVAL := 5.0  # send ping every N seconds while connected
const INACTIVITY_TIMEOUT := 30.0  # force-close if no message received for N seconds

# Per-port connection state
var _peers: Dictionary = {}  # port -> WebSocketPeer
var _connected: Dictionary = {}  # port -> bool
var _timers: Dictionary = {}  # port -> float (reconnect countdown)
var _connect_times: Dictionary = {}  # port -> float (elapsed seconds since connect)
var _last_activity: Dictionary = {}  # port -> float (seconds since last received message)
var _ping_timers: Dictionary = {}  # port -> float (seconds since last sent ping)
var _stale_ports: Dictionary = {}  # port -> bool (heartbeat timeout flag, exposed to UI)
var _running: bool = false


func start_server() -> void:
	_running = true
	for p in range(BASE_PORT, MAX_PORT + 1):
		_connected[p] = false
		_timers[p] = 0.0
		_try_connect(p)
	print("[MCP] Connecting to ports %d-%d" % [BASE_PORT, MAX_PORT])


func stop_server() -> void:
	_running = false
	for p in _peers:
		var ws: WebSocketPeer = _peers[p]
		if ws:
			ws.close(1000, "Plugin shutting down")
	_peers.clear()
	_connected.clear()
	_timers.clear()
	_last_activity.clear()
	_ping_timers.clear()
	_stale_ports.clear()
	print("[MCP] WebSocket client stopped")


func get_client_count() -> int:
	var count: int = 0
	for p in _connected:
		if _connected[p]:
			count += 1
	return count


func get_connected_ports() -> Array[int]:
	var ports: Array[int] = []
	for p: int in _connected:
		if _connected[p]:
			ports.append(p)
	return ports


func get_port_connect_time(port: int) -> float:
	return _connect_times.get(port, -1.0)


func get_port_idle_time(port: int) -> float:
	return _last_activity.get(port, -1.0)


func is_port_stale(port: int) -> bool:
	return _stale_ports.get(port, false)


func _try_connect(p: int) -> void:
	var ws := WebSocketPeer.new()
	ws.outbound_buffer_size = BUFFER_SIZE
	ws.inbound_buffer_size = BUFFER_SIZE
	var err := ws.connect_to_url("ws://127.0.0.1:%d" % p)
	if err == OK:
		_peers[p] = ws
	else:
		_peers[p] = null


func _process(delta: float) -> void:
	if not _running:
		return

	for p in range(BASE_PORT, MAX_PORT + 1):
		var ws: WebSocketPeer = _peers.get(p)

		# No peer - try reconnect on timer
		if ws == null:
			_timers[p] = _timers.get(p, 0.0) + delta
			if _timers[p] >= RECONNECT_INTERVAL:
				_timers[p] = 0.0
				_try_connect(p)
			continue

		ws.poll()
		var state := ws.get_ready_state()

		match state:
			WebSocketPeer.STATE_OPEN:
				if not _connected.get(p, false):
					_connected[p] = true
					_connect_times[p] = 0.0
					_last_activity[p] = 0.0
					_ping_timers[p] = 0.0
					_stale_ports[p] = false
					_timers[p] = 0.0
					print_verbose("[MCP] Connected on port %d" % p)
					client_connected.emit()
				else:
					_connect_times[p] = _connect_times.get(p, 0.0) + delta
					_last_activity[p] = _last_activity.get(p, 0.0) + delta
					_ping_timers[p] = _ping_timers.get(p, 0.0) + delta

				var received_any := false
				while ws.get_available_packet_count() > 0:
					var packet := ws.get_packet()
					var text := packet.get_string_from_utf8()
					received_any = true
					_dispatch_message(text, p)

				if received_any:
					_last_activity[p] = 0.0
					if _stale_ports.get(p, false):
						_stale_ports[p] = false
						print("[MCP] Port %d recovered from stale state" % p)

				# Force-close if no message received for INACTIVITY_TIMEOUT.
				# The MCP server pings every 10s, so 30s of silence means the
				# connection is half-open and reconnect is the only way out.
				if _last_activity.get(p, 0.0) > INACTIVITY_TIMEOUT:
					push_warning("[MCP] Port %d silent for %.1fs — forcing reconnect" % [p, _last_activity[p]])
					_stale_ports[p] = true
					ws.close(4000, "Heartbeat timeout")
					_connected[p] = false
					_peers[p] = null
					_timers[p] = 0.0
					client_disconnected.emit()
					continue

				# Send periodic ping so the server can detect our death too,
				# and so any reply resets our own inactivity timer.
				if _ping_timers.get(p, 0.0) >= PING_INTERVAL:
					_ping_timers[p] = 0.0
					ws.send_text(JSON.stringify({"jsonrpc": "2.0", "method": "ping", "params": {}}))

			WebSocketPeer.STATE_CLOSING:
				pass

			WebSocketPeer.STATE_CLOSED:
				if _connected.get(p, false):
					_connected[p] = false
					print_verbose("[MCP] Disconnected from port %d" % p)
					client_disconnected.emit()
				_peers[p] = null
				_timers[p] = 0.0
				_last_activity[p] = 0.0
				_ping_timers[p] = 0.0

			WebSocketPeer.STATE_CONNECTING:
				pass


func _send_to_port(p: int, text: String) -> void:
	var ws: WebSocketPeer = _peers.get(p)
	if ws and _connected.get(p, false):
		ws.send_text(text)


func send_message(text: String) -> void:
	# Broadcast to all connected peers
	for p in _peers:
		_send_to_port(p, text)


## Synchronous dispatch - parse JSON, handle ping/pong, queue command execution
func _dispatch_message(text: String, source_port: int) -> void:
	message_received.emit(text)

	var json := JSON.new()
	var err := json.parse(text)
	if err != OK:
		_send_response(source_port, null, null, {"code": -32700, "message": "Parse error"})
		return

	var msg: Variant = json.data
	if not msg is Dictionary:
		_send_response(source_port, null, null, {"code": -32600, "message": "Invalid request"})
		return

	var msg_dict: Dictionary = msg

	if msg_dict.get("method") == "ping":
		_send_to_port(source_port, JSON.stringify({"jsonrpc": "2.0", "method": "pong", "params": {}}))
		return

	if msg_dict.get("method") == "pong":
		return

	var id: Variant = msg_dict.get("id")
	var method: String = msg_dict.get("method", "")
	var params: Dictionary = msg_dict.get("params", {})

	if method.is_empty():
		_send_response(source_port, id, null, {"code": -32600, "message": "Missing method"})
		return

	if not command_router:
		_send_response(source_port, id, null, {"code": -32603, "message": "No command router"})
		return

	_execute_command.call_deferred(source_port, id, method, params)


func _execute_command(source_port: int, id: Variant, method: String, params: Dictionary) -> void:
	var cmd_result: Dictionary = await command_router.execute(method, params)
	if cmd_result.has("error"):
		var err_data: Variant = cmd_result["error"]
		_send_response(source_port, id, null, err_data)
		var response_text := JSON.stringify(err_data)
		command_executed.emit(method, false)
		command_completed.emit(method, false, response_text, source_port)
	else:
		var result_data: Variant = cmd_result.get("result", {})
		_send_response(source_port, id, result_data, null)
		var response_text := JSON.stringify(result_data)
		command_executed.emit(method, true)
		command_completed.emit(method, true, response_text, source_port)


func _send_response(source_port: int, id: Variant, result: Variant, err: Variant) -> void:
	var response: Dictionary = {"jsonrpc": "2.0", "id": id}
	if err != null:
		response["error"] = err
	else:
		response["result"] = result if result != null else {}
	_send_to_port(source_port, JSON.stringify(response))