// Copyright Ralpha Team. All Rights Reserved.

#include "RalphaMCPServer.h"
#include "RalphaParameterBridge.h"
#include "RalphaScreenCapture.h"
#include "Dom/JsonObject.h"
#include "Dom/JsonValue.h"
#include "Serialization/JsonSerializer.h"
#include "Serialization/JsonReader.h"
#include "Serialization/JsonWriter.h"
#include "Misc/Base64.h"
#include "Engine/World.h"
#include "Engine/Engine.h"
#include "Kismet/GameplayStatics.h"

DEFINE_LOG_CATEGORY(LogRalphaMCP);

URalphaMCPServer* URalphaMCPServer::Instance = nullptr;

URalphaMCPServer::URalphaMCPServer()
	: ListenerSocket(nullptr)
	, ServerPort(30010)
	, bIsRunning(false)
	, ParameterBridge(nullptr)
	, ScreenCapture(nullptr)
{
}

URalphaMCPServer::~URalphaMCPServer()
{
	Stop();
}

URalphaMCPServer* URalphaMCPServer::Get()
{
	if (!Instance)
	{
		Instance = NewObject<URalphaMCPServer>();
		Instance->AddToRoot(); // Prevent garbage collection
	}
	return Instance;
}

bool URalphaMCPServer::Start(int32 Port)
{
	if (bIsRunning)
	{
		UE_LOG(LogRalphaMCP, Warning, TEXT("Server already running on port %d"), ServerPort);
		return true;
	}

	ServerPort = Port;

	// Create socket
	ISocketSubsystem* SocketSubsystem = ISocketSubsystem::Get(PLATFORM_SOCKETSUBSYSTEM);
	if (!SocketSubsystem)
	{
		UE_LOG(LogRalphaMCP, Error, TEXT("Failed to get socket subsystem"));
		return false;
	}

	ListenerSocket = SocketSubsystem->CreateSocket(NAME_Stream, TEXT("RalphaMCPServer"), false);
	if (!ListenerSocket)
	{
		UE_LOG(LogRalphaMCP, Error, TEXT("Failed to create socket"));
		return false;
	}

	// Configure socket
	ListenerSocket->SetReuseAddr(true);
	ListenerSocket->SetNonBlocking(true);

	// Bind to port
	TSharedRef<FInternetAddr> Addr = SocketSubsystem->CreateInternetAddr();
	Addr->SetAnyAddress();
	Addr->SetPort(ServerPort);

	if (!ListenerSocket->Bind(*Addr))
	{
		UE_LOG(LogRalphaMCP, Error, TEXT("Failed to bind to port %d"), ServerPort);
		SocketSubsystem->DestroySocket(ListenerSocket);
		ListenerSocket = nullptr;
		return false;
	}

	// Start listening
	if (!ListenerSocket->Listen(8))
	{
		UE_LOG(LogRalphaMCP, Error, TEXT("Failed to listen on port %d"), ServerPort);
		SocketSubsystem->DestroySocket(ListenerSocket);
		ListenerSocket = nullptr;
		return false;
	}

	// Create parameter bridge and screen capture
	ParameterBridge = NewObject<URalphaParameterBridge>(this);
	ScreenCapture = NewObject<URalphaScreenCapture>(this);

	// Register tick
	TickDelegateHandle = FTSTicker::GetCoreTicker().AddTicker(
		FTickerDelegate::CreateUObject(this, &URalphaMCPServer::Tick), 0.0f);

	bIsRunning = true;
	UE_LOG(LogRalphaMCP, Log, TEXT("Ralpha MCP Server started on port %d"), ServerPort);

	return true;
}

void URalphaMCPServer::Stop()
{
	if (!bIsRunning)
	{
		return;
	}

	// Unregister tick
	FTSTicker::GetCoreTicker().RemoveTicker(TickDelegateHandle);

	// Close all client connections
	ClientConnections.Empty();

	// Close listener socket
	if (ListenerSocket)
	{
		ListenerSocket->Close();
		ISocketSubsystem::Get(PLATFORM_SOCKETSUBSYSTEM)->DestroySocket(ListenerSocket);
		ListenerSocket = nullptr;
	}

	bIsRunning = false;
	UE_LOG(LogRalphaMCP, Log, TEXT("Ralpha MCP Server stopped"));
}

bool URalphaMCPServer::IsRunning() const
{
	return bIsRunning;
}

bool URalphaMCPServer::Tick(float DeltaTime)
{
	if (!bIsRunning || !ListenerSocket)
	{
		return true;
	}

	// Check for new connections
	bool bHasPendingConnection = false;
	if (ListenerSocket->HasPendingConnection(bHasPendingConnection) && bHasPendingConnection)
	{
		TSharedRef<FInternetAddr> RemoteAddress = ISocketSubsystem::Get(PLATFORM_SOCKETSUBSYSTEM)->CreateInternetAddr();
		FSocket* ClientSocket = ListenerSocket->Accept(*RemoteAddress, TEXT("RalphaClient"));

		if (ClientSocket)
		{
			ClientSocket->SetNonBlocking(true);

			TSharedPtr<FRalphaClientConnection> NewConnection = MakeShared<FRalphaClientConnection>(ClientSocket, this);
			ClientConnections.Add(NewConnection);

			UE_LOG(LogRalphaMCP, Log, TEXT("New client connected from %s"), *RemoteAddress->ToString(true));
		}
	}

	// Process existing connections
	for (int32 i = ClientConnections.Num() - 1; i >= 0; --i)
	{
		if (ClientConnections[i]->IsValid())
		{
			ClientConnections[i]->Tick();
		}
		else
		{
			ClientConnections.RemoveAt(i);
		}
	}

	return true;
}

FString URalphaMCPServer::ProcessCommand(const FString& JsonCommand)
{
	TSharedPtr<FJsonObject> JsonObject;
	TSharedRef<TJsonReader<>> Reader = TJsonReaderFactory<>::Create(JsonCommand);

	if (!FJsonSerializer::Deserialize(Reader, JsonObject) || !JsonObject.IsValid())
	{
		return TEXT("{\"success\": false, \"error\": \"Invalid JSON\"}");
	}

	FString CommandType;
	if (!JsonObject->TryGetStringField(TEXT("type"), CommandType))
	{
		return TEXT("{\"success\": false, \"error\": \"Missing command type\"}");
	}

	TSharedPtr<FJsonObject> ResponseObject = MakeShared<FJsonObject>();

	// Route commands
	if (CommandType == TEXT("capture_screenshot"))
	{
		int32 Width = JsonObject->HasField(TEXT("width")) ? JsonObject->GetIntegerField(TEXT("width")) : 1920;
		int32 Height = JsonObject->HasField(TEXT("height")) ? JsonObject->GetIntegerField(TEXT("height")) : 1080;
		bool bUseCamera = JsonObject->HasField(TEXT("use_camera")) ? JsonObject->GetBoolField(TEXT("use_camera")) : true;

		FString Base64Image;
		FString FilePath;
		bool bSuccess = false;

		if (bUseCamera)
		{
			// Capture from CineCamera using SceneCapture2D
			bSuccess = ScreenCapture->CaptureFromCamera(Width, Height, Base64Image, FilePath);
		}

		// Fallback to viewport capture if camera capture failed or not requested
		if (!bSuccess)
		{
			bSuccess = ScreenCapture->CaptureScreenshot(Width, Height, Base64Image, FilePath);
		}

		ResponseObject->SetBoolField(TEXT("success"), bSuccess);
		if (bSuccess)
		{
			ResponseObject->SetStringField(TEXT("base64"), Base64Image);
			ResponseObject->SetStringField(TEXT("path"), FilePath);
		}
	}
	else if (CommandType == TEXT("get_all_parameters"))
	{
		TSharedPtr<FJsonObject> ParamsObject = ParameterBridge->GetAllParameters();
		ResponseObject->SetBoolField(TEXT("success"), true);
		ResponseObject->SetObjectField(TEXT("data"), ParamsObject);
	}
	else if (CommandType == TEXT("set_post_process"))
	{
		const TSharedPtr<FJsonObject>* Params;
		if (JsonObject->TryGetObjectField(TEXT("parameters"), Params))
		{
			bool bSuccess = ParameterBridge->SetPostProcessParameters(*Params);
			ResponseObject->SetBoolField(TEXT("success"), bSuccess);
		}
		else
		{
			ResponseObject->SetBoolField(TEXT("success"), false);
			ResponseObject->SetStringField(TEXT("error"), TEXT("Missing parameters"));
		}
	}
	else if (CommandType == TEXT("set_directional_light"))
	{
		const TSharedPtr<FJsonObject>* Params;
		if (JsonObject->TryGetObjectField(TEXT("parameters"), Params))
		{
			bool bSuccess = ParameterBridge->SetDirectionalLightParameters(*Params);
			ResponseObject->SetBoolField(TEXT("success"), bSuccess);
		}
		else
		{
			ResponseObject->SetBoolField(TEXT("success"), false);
			ResponseObject->SetStringField(TEXT("error"), TEXT("Missing parameters"));
		}
	}
	else if (CommandType == TEXT("set_sky_light"))
	{
		const TSharedPtr<FJsonObject>* Params;
		if (JsonObject->TryGetObjectField(TEXT("parameters"), Params))
		{
			bool bSuccess = ParameterBridge->SetSkyLightParameters(*Params);
			ResponseObject->SetBoolField(TEXT("success"), bSuccess);
		}
		else
		{
			ResponseObject->SetBoolField(TEXT("success"), false);
			ResponseObject->SetStringField(TEXT("error"), TEXT("Missing parameters"));
		}
	}
	else if (CommandType == TEXT("set_exponential_height_fog"))
	{
		const TSharedPtr<FJsonObject>* Params;
		if (JsonObject->TryGetObjectField(TEXT("parameters"), Params))
		{
			bool bSuccess = ParameterBridge->SetFogParameters(*Params);
			ResponseObject->SetBoolField(TEXT("success"), bSuccess);
		}
		else
		{
			ResponseObject->SetBoolField(TEXT("success"), false);
			ResponseObject->SetStringField(TEXT("error"), TEXT("Missing parameters"));
		}
	}
	else if (CommandType == TEXT("set_camera"))
	{
		const TSharedPtr<FJsonObject>* Params;
		if (JsonObject->TryGetObjectField(TEXT("parameters"), Params))
		{
			bool bSuccess = ParameterBridge->SetCameraParameters(*Params);
			ResponseObject->SetBoolField(TEXT("success"), bSuccess);
		}
		else
		{
			ResponseObject->SetBoolField(TEXT("success"), false);
			ResponseObject->SetStringField(TEXT("error"), TEXT("Missing parameters"));
		}
	}
	else if (CommandType == TEXT("search_assets"))
	{
		FString Query = JsonObject->GetStringField(TEXT("query"));
		FString AssetType = JsonObject->HasField(TEXT("asset_type")) ? JsonObject->GetStringField(TEXT("asset_type")) : TEXT("All");
		int32 MaxResults = JsonObject->HasField(TEXT("max_results")) ? JsonObject->GetIntegerField(TEXT("max_results")) : 20;

		TArray<TSharedPtr<FJsonValue>> Results = ParameterBridge->SearchAssets(Query, AssetType, MaxResults);
		ResponseObject->SetBoolField(TEXT("success"), true);
		ResponseObject->SetArrayField(TEXT("results"), Results);
	}
	else if (CommandType == TEXT("spawn_actor"))
	{
		FString AssetPath = JsonObject->GetStringField(TEXT("asset_path"));

		const TSharedPtr<FJsonObject>* LocationObj;
		FVector Location = FVector::ZeroVector;
		if (JsonObject->TryGetObjectField(TEXT("location"), LocationObj))
		{
			Location.X = (*LocationObj)->GetNumberField(TEXT("x"));
			Location.Y = (*LocationObj)->GetNumberField(TEXT("y"));
			Location.Z = (*LocationObj)->GetNumberField(TEXT("z"));
		}

		FRotator Rotation = FRotator::ZeroRotator;
		const TSharedPtr<FJsonObject>* RotationObj;
		if (JsonObject->TryGetObjectField(TEXT("rotation"), RotationObj))
		{
			Rotation.Pitch = (*RotationObj)->HasField(TEXT("pitch")) ? (*RotationObj)->GetNumberField(TEXT("pitch")) : 0.0f;
			Rotation.Yaw = (*RotationObj)->HasField(TEXT("yaw")) ? (*RotationObj)->GetNumberField(TEXT("yaw")) : 0.0f;
			Rotation.Roll = (*RotationObj)->HasField(TEXT("roll")) ? (*RotationObj)->GetNumberField(TEXT("roll")) : 0.0f;
		}

		float Scale = JsonObject->HasField(TEXT("scale")) ? JsonObject->GetNumberField(TEXT("scale")) : 1.0f;
		FString ActorLabel = JsonObject->HasField(TEXT("actor_label")) ? JsonObject->GetStringField(TEXT("actor_label")) : TEXT("");

		FString ActorId;
		bool bSuccess = ParameterBridge->SpawnActor(AssetPath, Location, Rotation, Scale, ActorLabel, ActorId);
		ResponseObject->SetBoolField(TEXT("success"), bSuccess);
		if (bSuccess)
		{
			ResponseObject->SetStringField(TEXT("actor_id"), ActorId);
		}
	}
	else if (CommandType == TEXT("delete_actor"))
	{
		FString ActorId = JsonObject->GetStringField(TEXT("actor_id"));
		bool bSuccess = ParameterBridge->DeleteActor(ActorId);
		ResponseObject->SetBoolField(TEXT("success"), bSuccess);
	}
	else if (CommandType == TEXT("set_actor_material"))
	{
		FString ActorId = JsonObject->GetStringField(TEXT("actor_id"));
		FString MaterialPath = JsonObject->GetStringField(TEXT("material_path"));
		int32 MaterialIndex = JsonObject->HasField(TEXT("material_index")) ? JsonObject->GetIntegerField(TEXT("material_index")) : 0;
		bool bSuccess = ParameterBridge->SetActorMaterial(ActorId, MaterialPath, MaterialIndex);
		ResponseObject->SetBoolField(TEXT("success"), bSuccess);
	}
	else if (CommandType == TEXT("set_actor_simple_material"))
	{
		FString ActorId = JsonObject->GetStringField(TEXT("actor_id"));

		// Parse color - supports hex string or RGB object
		FLinearColor BaseColor = FLinearColor(0.1f, 0.2f, 0.4f); // Default dark blue
		if (JsonObject->HasField(TEXT("color")))
		{
			const TSharedPtr<FJsonObject>* ColorObj;
			if (JsonObject->TryGetObjectField(TEXT("color"), ColorObj))
			{
				BaseColor.R = (*ColorObj)->HasField(TEXT("r")) ? (*ColorObj)->GetNumberField(TEXT("r")) : 0.1f;
				BaseColor.G = (*ColorObj)->HasField(TEXT("g")) ? (*ColorObj)->GetNumberField(TEXT("g")) : 0.2f;
				BaseColor.B = (*ColorObj)->HasField(TEXT("b")) ? (*ColorObj)->GetNumberField(TEXT("b")) : 0.4f;
				BaseColor.A = (*ColorObj)->HasField(TEXT("a")) ? (*ColorObj)->GetNumberField(TEXT("a")) : 1.0f;
			}
		}

		float Metallic = JsonObject->HasField(TEXT("metallic")) ? JsonObject->GetNumberField(TEXT("metallic")) : 0.0f;
		float Roughness = JsonObject->HasField(TEXT("roughness")) ? JsonObject->GetNumberField(TEXT("roughness")) : 0.5f;
		float Opacity = JsonObject->HasField(TEXT("opacity")) ? JsonObject->GetNumberField(TEXT("opacity")) : 1.0f;

		bool bSuccess = ParameterBridge->SetActorSimpleMaterial(ActorId, BaseColor, Metallic, Roughness, Opacity);
		ResponseObject->SetBoolField(TEXT("success"), bSuccess);
	}
	else if (CommandType == TEXT("list_actors"))
	{
		FString ClassFilter = JsonObject->HasField(TEXT("class_filter")) ? JsonObject->GetStringField(TEXT("class_filter")) : TEXT("");
		FString NameFilter = JsonObject->HasField(TEXT("name_filter")) ? JsonObject->GetStringField(TEXT("name_filter")) : TEXT("");
		bool bIncludeTransforms = JsonObject->HasField(TEXT("include_transforms")) ? JsonObject->GetBoolField(TEXT("include_transforms")) : true;

		TArray<TSharedPtr<FJsonValue>> Actors = ParameterBridge->ListActors(ClassFilter, NameFilter, bIncludeTransforms);
		ResponseObject->SetBoolField(TEXT("success"), true);
		ResponseObject->SetArrayField(TEXT("actors"), Actors);
	}
	else if (CommandType == TEXT("set_render_quality"))
	{
		const TSharedPtr<FJsonObject>* Settings;
		if (JsonObject->TryGetObjectField(TEXT("settings"), Settings))
		{
			bool bSuccess = ParameterBridge->SetRenderQuality(*Settings);
			ResponseObject->SetBoolField(TEXT("success"), bSuccess);
		}
		else
		{
			ResponseObject->SetBoolField(TEXT("success"), false);
			ResponseObject->SetStringField(TEXT("error"), TEXT("Missing settings"));
		}
	}
	else if (CommandType == TEXT("setup_scene"))
	{
		TSharedPtr<FJsonObject> SetupResult = ParameterBridge->SetupScene();
		ResponseObject = SetupResult;
	}
	else if (CommandType == TEXT("exec_command"))
	{
		FString Command = JsonObject->GetStringField(TEXT("command"));
		if (!Command.IsEmpty())
		{
			UWorld* World = GEngine->GetWorldContexts()[0].World();
			if (World)
			{
				GEngine->Exec(World, *Command);
				ResponseObject->SetBoolField(TEXT("success"), true);
				ResponseObject->SetStringField(TEXT("executed"), Command);
			}
			else
			{
				ResponseObject->SetBoolField(TEXT("success"), false);
				ResponseObject->SetStringField(TEXT("error"), TEXT("No world available"));
			}
		}
		else
		{
			ResponseObject->SetBoolField(TEXT("success"), false);
			ResponseObject->SetStringField(TEXT("error"), TEXT("Missing command"));
		}
	}
	else if (CommandType == TEXT("new_level"))
	{
		FString Template = JsonObject->HasField(TEXT("template")) ? JsonObject->GetStringField(TEXT("template")) : TEXT("Basic");
#if WITH_EDITOR
		if (GEditor)
		{
			// Create new level from template
			GEditor->Exec(nullptr, TEXT("MAP NEW"));
			ResponseObject->SetBoolField(TEXT("success"), true);
			ResponseObject->SetStringField(TEXT("message"), TEXT("New level created"));
		}
		else
		{
			ResponseObject->SetBoolField(TEXT("success"), false);
			ResponseObject->SetStringField(TEXT("error"), TEXT("Editor not available"));
		}
#else
		ResponseObject->SetBoolField(TEXT("success"), false);
		ResponseObject->SetStringField(TEXT("error"), TEXT("Only available in editor"));
#endif
	}
	else
	{
		ResponseObject->SetBoolField(TEXT("success"), false);
		ResponseObject->SetStringField(TEXT("error"), FString::Printf(TEXT("Unknown command type: %s"), *CommandType));
	}

	// Serialize response
	FString ResponseString;
	TSharedRef<TJsonWriter<>> Writer = TJsonWriterFactory<>::Create(&ResponseString);
	FJsonSerializer::Serialize(ResponseObject.ToSharedRef(), Writer);

	return ResponseString;
}

// FRalphaClientConnection implementation

FRalphaClientConnection::FRalphaClientConnection(FSocket* InSocket, URalphaMCPServer* InServer)
	: Socket(InSocket)
	, Server(InServer)
{
}

FRalphaClientConnection::~FRalphaClientConnection()
{
	if (Socket)
	{
		Socket->Close();
		ISocketSubsystem::Get(PLATFORM_SOCKETSUBSYSTEM)->DestroySocket(Socket);
		Socket = nullptr;
	}
}

bool FRalphaClientConnection::IsValid() const
{
	return Socket != nullptr && Socket->GetConnectionState() == SCS_Connected;
}

void FRalphaClientConnection::Tick()
{
	if (!IsValid())
	{
		return;
	}

	// Receive data
	uint32 PendingDataSize = 0;
	while (Socket->HasPendingData(PendingDataSize) && PendingDataSize > 0)
	{
		TArray<uint8> ReceivedData;
		ReceivedData.SetNumUninitialized(FMath::Min(PendingDataSize, 65536u));

		int32 BytesRead = 0;
		if (Socket->Recv(ReceivedData.GetData(), ReceivedData.Num(), BytesRead))
		{
			ReceiveBuffer += FString(UTF8_TO_TCHAR(reinterpret_cast<const char*>(ReceivedData.GetData())));
		}
	}

	// Process complete messages (newline-delimited JSON)
	int32 NewlineIndex;
	while (ReceiveBuffer.FindChar(TEXT('\n'), NewlineIndex))
	{
		FString Message = ReceiveBuffer.Left(NewlineIndex);
		ReceiveBuffer = ReceiveBuffer.Mid(NewlineIndex + 1);

		if (!Message.IsEmpty())
		{
			FString Response = Server->ProcessCommand(Message);
			SendResponse(Response);
		}
	}
}

void FRalphaClientConnection::SendResponse(const FString& Response)
{
	if (!IsValid())
	{
		return;
	}

	FString ResponseWithNewline = Response + TEXT("\n");
	FTCHARToUTF8 Converter(*ResponseWithNewline);
	int32 BytesSent = 0;
	Socket->Send(reinterpret_cast<const uint8*>(Converter.Get()), Converter.Length(), BytesSent);
}