#pragma once

#include "CoreMinimal.h"
#include "Dom/JsonObject.h"
#include "MCPUtils.h"
class FSocket;
class IMCPHandler;

/**
 * FMCPServer — plain C++ class (not a UCLASS) that implements the
 * Model Context Protocol (MCP) over a TCP socket using JSON-RPC.
 *
 * Clients connect via TCP and exchange newline-delimited JSON-RPC messages
 * (the MCP "stdio" protocol). Each connected client gets its own thread
 * for blocking I/O; tool calls are dispatched on the game thread.
 *
 * Both the standalone commandlet (UBlueprintMCPCommandlet) and the in-editor
 * subsystem (UBlueprintMCPEditorSubsystem) delegate to an instance of this
 * class. The only difference is *who ticks the engine*:
 *   - Commandlet: manual FTSTicker loop
 *   - Editor subsystem: UE editor tick via FTickableEditorObject
 */
class FMCPServer
{
public:
	/** Get the active server instance via the editor subsystem. */
	static FMCPServer* Get();

	/** Start listening for MCP clients on the given TCP port. */
	bool Start(int32 InPort, bool bEditorMode = false);

	/** Stop the server: drain pending requests, close all sockets, join threads. */
	void Stop();

	/**
	 * Process pending MCP requests on the game thread.
	 * Call this every tick from whichever host owns this server.
	 * Returns true if a request was processed.
	 */
	bool ProcessOneRequest();

	/** Whether the server is currently listening. */
	bool IsRunning() const { return bRunning; }

	/** Port the server is listening on. */
	int32 GetPort() const { return Port; }

private:
	// ----- Tool dispatch -----
	using FRequestHandler = TFunction<void(const FJsonObject* Json, FJsonObject* Result)>;
	TMap<FString, FRequestHandler> HandlerMap;         // old-style handlers
	TMap<FString, UClass*> MCPHandlerRegistry;         // new-style: tool name -> UMCPHandler subclass
	TSet<FString> MutationEndpoints;
	void RegisterHandlers();
	void BuildMCPHandlerRegistry();

	// Dispatch a tool call to the appropriate handler
	void DispatchToolCall(const FString& ToolName, const FJsonObject* Params, FJsonObject* Result);

	// ----- MCP protocol -----
	// Handle a complete JSON-RPC line and return the response line (or empty for notifications)
	FString HandleJsonRpc(const FString& Line);

	// MCP method handlers
	FString HandleInitialize(int32 Id, const FJsonObject* Params);
	FString HandleToolsList(int32 Id);
	FString HandleToolsCall(int32 Id, const FJsonObject* Params);

	// Build the tools/list response (cached after first call)
	TSharedPtr<FJsonObject> CachedToolsList;
	static TSharedRef<FJsonObject> FPropertyToPropSchema(FProperty* Prop);
	static TSharedRef<FJsonObject> HandlerClassToToolSchema(UClass* HandlerClass);
	void BuildCachedToolsList();

	// JSON-RPC helpers
	static FString MakeJsonRpcResult(int32 Id, TSharedPtr<FJsonObject> Result);
	static FString MakeJsonRpcToolResult(int32 Id, TSharedPtr<FJsonObject> ToolResult, bool bIsError);
	static FString MakeJsonRpcError(int32 Id, int32 Code, const FString& Message);

	// ----- TCP server -----
	FSocket* ListenSocket = nullptr;
	int32 Port = 9847;
	bool bRunning = false;
	bool bIsEditor = false;

	// ----- Client connections -----
	struct FClientConnection
	{
		FSocket* Socket = nullptr;
		TFuture<void> ThreadFuture;
		bool bDone = false;
	};
	TArray<TSharedPtr<FClientConnection>> Clients;
	void AcceptNewConnections();
	void CleanupFinishedClients();
	static void ClientThreadFunc(FMCPServer* Server, TSharedPtr<FClientConnection> Client);

	// ----- Thread-safe message queue -----
	struct FPendingMessage
	{
		FString Line;
		TPromise<FString> Response;
		FPendingMessage() : Response(TPromise<FString>()) {}
	};
	FCriticalSection Mutex;
	TArray<TSharedPtr<FPendingMessage>> PendingMessages;
	bool bShuttingDown = false;

	// ----- Request handlers (read-only) -----
	void HandleList(const FJsonObject* Json, FJsonObject* Result);
	void HandleGetBlueprint(const FJsonObject* Json, FJsonObject* Result);
	void HandleGetGraph(const FJsonObject* Json, FJsonObject* Result);
	void HandleSearch(const FJsonObject* Json, FJsonObject* Result);
	void HandleFindReferences(const FJsonObject* Json, FJsonObject* Result);
	void HandleSearchByType(const FJsonObject* Json, FJsonObject* Result);

	// ----- Request handlers (write) -----
	void HandleChangeVariableType(const FJsonObject* Json, FJsonObject* Result);
	void HandleChangeFunctionParamType(const FJsonObject* Json, FJsonObject* Result);
	void HandleRemoveFunctionParameter(const FJsonObject* Json, FJsonObject* Result);

	// ----- Pin introspection (read-only) -----
	void HandleGetPinInfo(const FJsonObject* Json, FJsonObject* Result);
	void HandleCheckPinCompatibility(const FJsonObject* Json, FJsonObject* Result);

	// ----- Class/function discovery (read-only) -----
	void HandleListClasses(const FJsonObject* Json, FJsonObject* Result);
	void HandleListFunctions(const FJsonObject* Json, FJsonObject* Result);
	void HandleListProperties(const FJsonObject* Json, FJsonObject* Result);

	// ----- Reparent -----
	void HandleReparentBlueprint(const FJsonObject* Json, FJsonObject* Result);

	// ----- Create -----
	void HandleCreateBlueprint(const FJsonObject* Json, FJsonObject* Result);
	void HandleCreateGraph(const FJsonObject* Json, FJsonObject* Result);

	// ----- Graph manipulation -----
	void HandleDeleteGraph(const FJsonObject* Json, FJsonObject* Result);
	void HandleRenameGraph(const FJsonObject* Json, FJsonObject* Result);

	// ----- Variables -----
	void HandleAddVariable(const FJsonObject* Json, FJsonObject* Result);
	void HandleRemoveVariable(const FJsonObject* Json, FJsonObject* Result);
	void HandleSetVariableMetadata(const FJsonObject* Json, FJsonObject* Result);

	// ----- Event Dispatchers -----
	void HandleAddEventDispatcher(const FJsonObject* Json, FJsonObject* Result);
	void HandleListEventDispatchers(const FJsonObject* Json, FJsonObject* Result);

	// ----- Function Parameters -----
	void HandleAddFunctionParameter(const FJsonObject* Json, FJsonObject* Result);

	// ----- Components -----
	void HandleAddComponent(const FJsonObject* Json, FJsonObject* Result);
	void HandleRemoveComponent(const FJsonObject* Json, FJsonObject* Result);
	void HandleListComponents(const FJsonObject* Json, FJsonObject* Result);

	// ----- Diagnostic -----
	void HandleTestSave(const FJsonObject* Json, FJsonObject* Result);

	// ----- Snapshot / Safety tools (write) -----
	void HandleSnapshotGraph(const FJsonObject* Json, FJsonObject* Result);
	void HandleDiffGraph(const FJsonObject* Json, FJsonObject* Result);
	void HandleRestoreGraph(const FJsonObject* Json, FJsonObject* Result);
	void HandleFindDisconnectedPins(const FJsonObject* Json, FJsonObject* Result);
	void HandleAnalyzeRebuildImpact(const FJsonObject* Json, FJsonObject* Result);

	// ----- Material read-only handlers (Phase 1) -----
	void HandleListMaterials(const FJsonObject* Json, FJsonObject* Result);
	void HandleGetMaterial(const FJsonObject* Json, FJsonObject* Result);
	void HandleGetMaterialGraph(const FJsonObject* Json, FJsonObject* Result);
	void HandleDescribeMaterial(const FJsonObject* Json, FJsonObject* Result);
	void HandleSearchMaterials(const FJsonObject* Json, FJsonObject* Result);
	void HandleFindMaterialReferences(const FJsonObject* Json, FJsonObject* Result);

	// ----- Material mutation handlers (Phase 2) -----
	void HandleCreateMaterial(const FJsonObject* Json, FJsonObject* Result);
	void HandleSetMaterialProperty(const FJsonObject* Json, FJsonObject* Result);
	void HandleAddMaterialExpression(const FJsonObject* Json, FJsonObject* Result);
	void HandleDeleteMaterialExpression(const FJsonObject* Json, FJsonObject* Result);
	void HandleConnectMaterialPins(const FJsonObject* Json, FJsonObject* Result);
	void HandleDisconnectMaterialPin(const FJsonObject* Json, FJsonObject* Result);
	void HandleSetExpressionValue(const FJsonObject* Json, FJsonObject* Result);
	void HandleMoveMaterialExpression(const FJsonObject* Json, FJsonObject* Result);

	// ----- Material instance handlers (Phase 3) -----
	void HandleCreateMaterialInstance(const FJsonObject* Json, FJsonObject* Result);
	void HandleSetMaterialInstanceParameter(const FJsonObject* Json, FJsonObject* Result);
	void HandleGetMaterialInstanceParameters(const FJsonObject* Json, FJsonObject* Result);
	void HandleReparentMaterialInstance(const FJsonObject* Json, FJsonObject* Result);

	// ----- Material function handlers (Phase 4) -----
	void HandleListMaterialFunctions(const FJsonObject* Json, FJsonObject* Result);
	void HandleGetMaterialFunction(const FJsonObject* Json, FJsonObject* Result);
	void HandleCreateMaterialFunction(const FJsonObject* Json, FJsonObject* Result);

	// ----- Material validation -----
	void HandleValidateMaterial(const FJsonObject* Json, FJsonObject* Result);

	// ----- Material snapshot/diff/restore (Phase 5) -----
	void HandleSnapshotMaterialGraph(const FJsonObject* Json, FJsonObject* Result);
	void HandleDiffMaterialGraph(const FJsonObject* Json, FJsonObject* Result);
	void HandleRestoreMaterialGraph(const FJsonObject* Json, FJsonObject* Result);

	// ----- Animation Blueprint handlers -----
	void HandleCreateAnimBlueprint(const FJsonObject* Json, FJsonObject* Result);
	void HandleAddAnimState(const FJsonObject* Json, FJsonObject* Result);
	void HandleRemoveAnimState(const FJsonObject* Json, FJsonObject* Result);
	void HandleAddAnimTransition(const FJsonObject* Json, FJsonObject* Result);
	void HandleSetTransitionRule(const FJsonObject* Json, FJsonObject* Result);
	void HandleAddAnimNode(const FJsonObject* Json, FJsonObject* Result);
	void HandleAddStateMachine(const FJsonObject* Json, FJsonObject* Result);
	void HandleSetStateAnimation(const FJsonObject* Json, FJsonObject* Result);
	void HandleListAnimSlots(const FJsonObject* Json, FJsonObject* Result);
	void HandleListSyncGroups(const FJsonObject* Json, FJsonObject* Result);
	void HandleCreateBlendSpace(const FJsonObject* Json, FJsonObject* Result);
	void HandleSetBlendSpaceSamples(const FJsonObject* Json, FJsonObject* Result);
	void HandleSetStateBlendSpace(const FJsonObject* Json, FJsonObject* Result);

public:
	// ----- Snapshot storage -----
	TMap<FString, FGraphSnapshot> Snapshots;
	TMap<FString, FGraphSnapshot> MaterialSnapshots;
	static const int32 MaxSnapshots = 50;

	// Snapshot helpers
	FString GenerateSnapshotId(const FString& BlueprintName);
	FGraphSnapshotData CaptureGraphSnapshot(UEdGraph* Graph);
	void PruneOldSnapshots();
	bool SaveSnapshotToDisk(const FString& SnapshotId, const FGraphSnapshot& Snapshot);
	bool LoadSnapshotFromDisk(const FString& SnapshotId, FGraphSnapshot& OutSnapshot);
};

// Transitional alias — old-style handlers use this to access the server instance.
using FBlueprintMCPServer = FMCPServer;
using MCPHelper = FMCPServer;