#include "HyperTwistSolverLibrary.h"

#include "HAL/FileManager.h"
#include "HAL/PlatformMisc.h"
#include "HAL/PlatformTime.h"
#include "Misc/Paths.h"

#include "ThirdParty/rob-twophase/face.h"
#include "ThirdParty/rob-twophase/move.h"
#include "ThirdParty/rob-twophase/cubie.h"
#include "ThirdParty/rob-twophase/coord.h"
#include "ThirdParty/rob-twophase/sym.h"
#include "ThirdParty/rob-twophase/prun.h"
#include "ThirdParty/rob-twophase/solve.h"

#include <atomic>
#include <mutex>

static std::atomic_bool bSolverInitialized{false};
static std::once_flag CoreTablesInitOnceFlag;
static std::once_flag InitOnceFlag;
static std::mutex SolverInvocationMutex;

namespace
{
	int32 DetermineSolverThreadCount()
	{
		const int32 AvailablePhysicalCores = FMath::Max(1, FPlatformMisc::NumberOfCores());
		return FMath::Clamp(AvailablePhysicalCores - 2, 1, 8);
	}

	int32 DetermineSolverSplitCount(const int32 ThreadCount)
	{
		// Single-split search leaves short classic-cube solves badly under-distributed on
		// multicore hosts. Split phase-1 roots more aggressively so guidance/search work
		// can saturate the available worker pool.
		return FMath::Clamp(ThreadCount * 3, 1, move::COUNT1);
	}

	bool TryResolveMoveIndex(const FString& MoveString, int32& OutMoveIndex)
	{
		const FString TrimmedMove = MoveString.TrimStartAndEnd();
		for (int32 MoveIndex = 0; MoveIndex < move::COUNT; ++MoveIndex)
		{
			if (TrimmedMove.Equals(UTF8_TO_TCHAR(move::names[MoveIndex].c_str()), ESearchCase::IgnoreCase))
			{
				OutMoveIndex = MoveIndex;
				return true;
			}
		}

		return false;
	}

	void EnsureCoreTablesInitialized()
	{
		std::call_once(CoreTablesInitOnceFlag, []()
		{
			face::init();
			move::init();
		});
	}
}

static void DoInitSolver()
{
	const FString SavedDirectory = FPaths::ConvertRelativePathToFull(FPaths::ProjectSavedDir());
	if (!IFileManager::Get().MakeDirectory(*SavedDirectory, true))
	{
		UE_LOG(
			LogTemp,
			Error,
			TEXT("HyperTwistSolver: Could not create solver cache directory '%s'."),
			*SavedDirectory);
		return;
	}

	const FString PruningTablePath = FPaths::Combine(
		SavedDirectory,
		UTF8_TO_TCHAR(prun::SAVE.c_str()));
	const bool bExistingTable = IFileManager::Get().FileSize(*PruningTablePath) > 0;
	const double InitializationStartedAt = FPlatformTime::Seconds();
	UE_LOG(
		LogTemp,
		Display,
		TEXT("HyperTwistSolver: %s pruning table at '%s'."),
		bExistingTable ? TEXT("Loading") : TEXT("Generating first-run"),
		*PruningTablePath);

	EnsureCoreTablesInitialized();
	coord::init();
	sym::init();
	const bool bPruningTablesReady = prun::init(
		true,
		TCHAR_TO_UTF8(*PruningTablePath));
	bSolverInitialized.store(bPruningTablesReady, std::memory_order_release);
	if (bPruningTablesReady)
	{
		UE_LOG(
			LogTemp,
			Display,
			TEXT("HyperTwistSolver: Pruning table ready in %.2f seconds (%lld bytes)."),
			FPlatformTime::Seconds() - InitializationStartedAt,
			IFileManager::Get().FileSize(*PruningTablePath));
	}
	else
	{
		UE_LOG(
			LogTemp,
			Error,
			TEXT("HyperTwistSolver: Pruning table failed after %.2f seconds."),
			FPlatformTime::Seconds() - InitializationStartedAt);
	}
}

bool UHyperTwistSolverLibrary::IsSolverInitialized()
{
	return bSolverInitialized.load(std::memory_order_acquire);
}

bool UHyperTwistSolverLibrary::InitializeSolver()
{
	std::call_once(InitOnceFlag, DoInitSolver);
	return bSolverInitialized.load(std::memory_order_acquire);
}

TArray<FString> UHyperTwistSolverLibrary::SolveClassicState(
	const FString& FaceletString,
	int32 TimeLimitMs,
	int32 MaxLength,
	int32 NumSolutions
)
{
	const std::lock_guard<std::mutex> SolverLock(SolverInvocationMutex);
	TArray<FString> Result;

	if (!InitializeSolver())
	{
		UE_LOG(LogTemp, Error, TEXT("HyperTwistSolver: Failed to initialize solver"));
		return Result;
	}

	std::string Facelets(TCHAR_TO_UTF8(*FaceletString));
	cubie::cube Cube;
	int Err = face::to_cubie(Facelets, Cube);
	if (Err != 0)
	{
		UE_LOG(LogTemp, Error, TEXT("HyperTwistSolver: Invalid facelet string, error code %d"), Err);
		return Result;
	}

	if (cubie::check_cube(Cube) != 0)
	{
		UE_LOG(LogTemp, Error, TEXT("HyperTwistSolver: Cube state is not solvable"));
		return Result;
	}

	const int32 ThreadCount = DetermineSolverThreadCount();
	const int32 SplitCount = DetermineSolverSplitCount(ThreadCount);
	solve::Engine Solver(ThreadCount, TimeLimitMs, NumSolutions, MaxLength, SplitCount);
	Solver.prepare();

	std::vector<std::vector<int>> Solutions;
	Solver.solve(Cube, Solutions);
	Solver.finish();

	if (Solutions.empty())
	{
		UE_LOG(LogTemp, Warning, TEXT("HyperTwistSolver: No solution found within time limit"));
		return Result;
	}

	// Return the shortest solution
	const std::vector<int>& Best = Solutions[0];
	for (int Move : Best)
	{
		if (Move >= 0 && Move < move::COUNT)
		{
			Result.Add(FString(UTF8_TO_TCHAR(move::names[Move].c_str())));
		}
	}

	return Result;
}

FString UHyperTwistSolverLibrary::GetMoveName(int32 MoveIndex)
{
	if (MoveIndex >= 0 && MoveIndex < move::COUNT)
	{
		return FString(UTF8_TO_TCHAR(move::names[MoveIndex].c_str()));
	}
	return FString();
}

int32 UHyperTwistSolverLibrary::GetMoveCount()
{
	return move::COUNT;
}

bool UHyperTwistSolverLibrary::VerifyFaceletString(const FString& FaceletString)
{
	if (FaceletString.Len() != 54)
	{
		return false;
	}

	EnsureCoreTablesInitialized();
	std::string Facelets(TCHAR_TO_UTF8(*FaceletString));
	cubie::cube Cube;
	return face::to_cubie(Facelets, Cube) == 0;
}

bool UHyperTwistSolverLibrary::VerifySolution(
	const FString& FaceletString,
	const TArray<FString>& Moves
)
{
	const std::lock_guard<std::mutex> SolverLock(SolverInvocationMutex);
	if (!InitializeSolver() || !VerifyFaceletString(FaceletString))
	{
		return false;
	}

	std::string Facelets(TCHAR_TO_UTF8(*FaceletString));
	cubie::cube Cube;
	if (face::to_cubie(Facelets, Cube) != 0 || cubie::check_cube(Cube) != 0)
	{
		return false;
	}

	for (const FString& MoveString : Moves)
	{
		int32 MoveIndex = INDEX_NONE;
		if (!TryResolveMoveIndex(MoveString, MoveIndex))
		{
			UE_LOG(LogTemp, Warning, TEXT("HyperTwistSolver: Unrecognized move '%s' while verifying solution"), *MoveString);
			return false;
		}

		cubie::cube NextCube;
		cubie::mul(Cube, move::cubes[MoveIndex], NextCube);
		Cube = NextCube;
	}

	return Cube == cubie::SOLVED_CUBE;
}
