@tool
class_name GolemServer
extends Node

## TCP server that receives JSON-line commands from Python MCP bridges
## and routes them to the appropriate handler. Supports multiple clients.

signal client_connected(client_id: int)
signal client_disconnected(client_id: int)
signal handler_error(tool_name: String)

const PORT := 3571
const MAX_MESSAGE_SIZE := 16 * 1024 * 1024  # 16 MB
const MAX_CLIENTS := 5
const MAX_ACTIVITY := 8
var _server: TCPServer
var _peers: Dictionary = {}      # {client_id: StreamPeerTCP}
var _buffers: Dictionary = {}    # {client_id: String}
var _request_owners: Dictionary = {}  # {request_id: client_id} — routes responses
var _request_tools: Dictionary = {}   # {request_id: tool_name} — for deferred logging
var _next_client_id: int = 0
var _handlers: Dictionary = {}   # {tool_name: Callable}
var _command_queue: Array[Dictionary] = []  # {cmd: Dictionary, client_id: int}
var _deferred_responses: Array[Dictionary] = []
var _request_count: int = 0
var _activity_log: Array[Dictionary] = []  # {tool: String, time: float, ok: bool, duration: float}
var _error_count: int = 0
var _request_start_times: Dictionary = {}  # {request_id: float} — for duration tracking

# Handler references (set by the plugin for cross-handler access)
var screenshot_handler: Node
var scene_handler: Node
var inspect_handler: Node
var run_handler: Node
var script_handler: Node
var project_handler: Node
var input_handler: Node
var theme_handler: Node


func _ready() -> void:
	_server = TCPServer.new()
	var err := _server.listen(PORT, "127.0.0.1")
	if err != OK:
		push_error("GolemServer: Failed to listen on port %d (error %d)" % [PORT, err])
		return
	print("GolemServer: Listening on port %d (max %d clients)" % [PORT, MAX_CLIENTS])


func _process(delta: float) -> void:
	if _server == null:
		return

	_accept_new_connections()
	_read_all_clients()
	_poll_deferred_responses(delta)
	_process_command_queue()


# ------------------------------------------------------------------
# Connection management
# ------------------------------------------------------------------

func _accept_new_connections() -> void:
	while _server.is_connection_available():
		var incoming := _server.take_connection()
		incoming.poll()
		var status := incoming.get_status()

		# Ghost connection (port probe) — ignore
		if status == StreamPeerTCP.STATUS_NONE or status == StreamPeerTCP.STATUS_ERROR:
			incoming.disconnect_from_host()
			continue

		# Cap reached — reject
		if _peers.size() >= MAX_CLIENTS:
			push_warning("GolemServer: Rejected connection — max %d clients reached" % MAX_CLIENTS)
			incoming.disconnect_from_host()
			continue

		var id := _next_client_id
		_next_client_id += 1
		_peers[id] = incoming
		_buffers[id] = ""

		print("GolemServer: Client %d connected (%d/%d)" % [id, _peers.size(), MAX_CLIENTS])
		client_connected.emit(id)


func _read_all_clients() -> void:
	for id in _peers.keys():
		var peer: StreamPeerTCP = _peers[id]
		peer.poll()
		var status := peer.get_status()

		if status == StreamPeerTCP.STATUS_NONE or status == StreamPeerTCP.STATUS_ERROR:
			_on_peer_disconnected(id)
			continue

		if status != StreamPeerTCP.STATUS_CONNECTED:
			continue

		var available := peer.get_available_bytes()
		if available > 0:
			var data := peer.get_data(available)
			if data[0] == OK:
				_buffers[id] += data[1].get_string_from_utf8()
				_process_buffer(id)


func _on_peer_disconnected(id: int) -> void:
	_peers.erase(id)
	_buffers.erase(id)
	# Cancel deferred responses owned by this client
	var cancelled := 0
	var i := _deferred_responses.size() - 1
	while i >= 0:
		if _request_owners.get(_deferred_responses[i]["id"]) == id:
			_request_owners.erase(_deferred_responses[i]["id"])
			_request_tools.erase(_deferred_responses[i]["id"])
			_deferred_responses.remove_at(i)
			cancelled += 1
		i -= 1
	# Clean up request ownership entries for this client
	for req_id in _request_owners.keys():
		if _request_owners[req_id] == id:
			_request_owners.erase(req_id)
			_request_tools.erase(req_id)

	var suffix := " (cancelled %d deferred)" % cancelled if cancelled > 0 else ""
	print("GolemServer: Client %d disconnected%s (%d/%d)" % [id, suffix, _peers.size(), MAX_CLIENTS])
	client_disconnected.emit(id)


# ------------------------------------------------------------------
# Buffer processing
# ------------------------------------------------------------------

func _process_buffer(client_id: int) -> void:
	var buf: String = _buffers[client_id]
	while true:
		var newline_pos := buf.find("\n")
		if newline_pos == -1:
			break
		var line := buf.substr(0, newline_pos).strip_edges()
		buf = buf.substr(newline_pos + 1)
		if line.is_empty():
			continue
		var parsed = JSON.parse_string(line)
		if parsed == null:
			push_warning("GolemServer: Invalid JSON from client %d" % client_id)
			continue
		_command_queue.append({"cmd": parsed, "client_id": client_id})
	_buffers[client_id] = buf


# ------------------------------------------------------------------
# Command routing
# ------------------------------------------------------------------

func register_handler(tool_name: String, callback: Callable) -> void:
	_handlers[tool_name] = callback


func defer_response(id: String, check_fn: Callable, timeout: float = 5.0, timeout_msg: String = "Operation timed out") -> void:
	_deferred_responses.append({
		"id": id,
		"check": check_fn,
		"timeout": timeout,
		"timeout_msg": timeout_msg,
		"elapsed": 0.0,
	})


func _process_command_queue() -> void:
	if _command_queue.size() == 0:
		return
	var entry: Dictionary = _command_queue.pop_front()
	var cmd: Dictionary = entry["cmd"]
	var client_id: int = entry["client_id"]
	_execute_command(cmd, client_id)


func _execute_command(cmd: Dictionary, client_id: int) -> void:
	_request_count += 1
	var id: String = cmd.get("id", "")
	var tool_name: String = cmd.get("tool", "")
	var args: Dictionary = cmd.get("args", {})

	# Track which client owns this request (for response routing)
	_request_owners[id] = client_id
	_request_start_times[id] = Time.get_unix_time_from_system()

	# Built-in: ping
	if tool_name == "ping":
		_send_response(id, {"ok": true, "result": "pong", "type": "text"})
		return

	# Built-in: graceful disconnect
	if tool_name == "disconnect":
		print("GolemServer: Client %d sent graceful disconnect" % client_id)
		if _peers.has(client_id):
			_peers[client_id].disconnect_from_host()
		_on_peer_disconnected(client_id)
		return

	if not _handlers.has(tool_name):
		_send_error(id, "Unknown tool: %s" % tool_name)
		_log_activity(tool_name, false, id)
		return

	var handler: Callable = _handlers[tool_name]
	var result = handler.call(args)

	if result == null:
		push_warning("GolemServer: Handler returned null for '%s'" % tool_name)
		_send_error(id, "Handler returned null for tool: %s (possible crash)" % tool_name)
		_log_activity(tool_name, false, id)
	elif result is Dictionary:
		if result.get("_deferred"):
			_request_tools[id] = tool_name
			defer_response(
				id,
				result["_check_fn"],
				result.get("_timeout", 5.0),
				result.get("_timeout_msg", "Operation timed out"),
			)
		else:
			_send_response(id, result)
			_log_activity(tool_name, result.get("ok", true), id)
	else:
		_send_error(id, "Handler returned invalid result for tool: %s" % tool_name)
		_log_activity(tool_name, false, id)


# ------------------------------------------------------------------
# Deferred responses
# ------------------------------------------------------------------

func _poll_deferred_responses(delta: float) -> void:
	var i := _deferred_responses.size() - 1
	while i >= 0:
		var deferred: Dictionary = _deferred_responses[i]
		deferred["elapsed"] += delta
		var check_fn: Callable = deferred["check"]
		var result = check_fn.call()
		if result != null:
			_send_response(deferred["id"], result)
			var t: String = _request_tools.get(deferred["id"], "deferred")
			_request_tools.erase(deferred["id"])
			_log_activity(t, true, deferred["id"])
			_deferred_responses.remove_at(i)
		elif deferred["elapsed"] >= deferred["timeout"]:
			_send_error(deferred["id"], deferred.get("timeout_msg", "Deferred operation timed out"))
			var t: String = _request_tools.get(deferred["id"], "deferred")
			_request_tools.erase(deferred["id"])
			_log_activity(t, false, deferred["id"])
			_deferred_responses.remove_at(i)
		i -= 1


# ------------------------------------------------------------------
# Response routing — sends to the correct client
# ------------------------------------------------------------------

func _send_response(id: String, data: Dictionary) -> void:
	var client_id: int = _request_owners.get(id, -1)
	_request_owners.erase(id)

	if client_id == -1 or not _peers.has(client_id):
		push_warning("GolemServer: Cannot route response for '%s' — client gone" % id)
		return

	var peer: StreamPeerTCP = _peers[client_id]
	if peer.get_status() != StreamPeerTCP.STATUS_CONNECTED:
		push_warning("GolemServer: Client %d disconnected before response for '%s'" % [client_id, id])
		return

	data["id"] = id
	if not data.has("ok"):
		data["ok"] = true
	var json_str := JSON.stringify(data) + "\n"
	var buffer := json_str.to_utf8_buffer()

	# Send in chunks (64 KB) to avoid overwhelming the TCP socket buffer
	var chunk_size := 65536
	var offset := 0
	while offset < buffer.size():
		if peer.get_status() != StreamPeerTCP.STATUS_CONNECTED:
			push_warning("GolemServer: Connection lost during send to client %d" % client_id)
			return
		var end := mini(offset + chunk_size, buffer.size())
		var chunk := buffer.slice(offset, end)
		var err := peer.put_data(chunk)
		if err != OK:
			push_warning("GolemServer: put_data failed for client %d (error %d)" % [client_id, err])
			return
		offset = end


func _send_error(id: String, error_msg: String) -> void:
	_send_response(id, {"ok": false, "error": error_msg})


# ------------------------------------------------------------------
# Activity tracking
# ------------------------------------------------------------------

func _log_activity(tool_name: String, ok: bool, request_id: String = "") -> void:
	var now := Time.get_unix_time_from_system()
	var duration := 0.0
	if not request_id.is_empty() and _request_start_times.has(request_id):
		duration = now - _request_start_times[request_id]
		_request_start_times.erase(request_id)
	_activity_log.push_back({
		"tool": tool_name,
		"time": now,
		"ok": ok,
		"duration": duration,
	})
	if _activity_log.size() > MAX_ACTIVITY:
		_activity_log.pop_front()
	if not ok:
		_error_count += 1
		handler_error.emit(tool_name)


# ------------------------------------------------------------------
# Public API
# ------------------------------------------------------------------

func get_client_count() -> int:
	return _peers.size()


func get_request_count() -> int:
	return _request_count


func get_activity_log() -> Array[Dictionary]:
	return _activity_log


func get_deferred_count() -> int:
	return _deferred_responses.size()


func get_error_count() -> int:
	return _error_count


func stop() -> void:
	for id in _peers.keys():
		_peers[id].disconnect_from_host()
	_peers.clear()
	_buffers.clear()
	_request_owners.clear()
	_request_tools.clear()
	_request_start_times.clear()
	_activity_log.clear()
	if _server != null:
		_server.stop()
		_server = null
	print("GolemServer: Stopped")
