#include "HyperTwistSimulation/HyperTwistMagic120CellRuntimeLibrary.h"

#include "Generated/HyperTwistMagic120CellPermutationData.inl"
#include "JsonObjectConverter.h"
#include "Misc/Base64.h"
#include "Misc/Compression.h"

namespace HyperTwistMagic120CellRuntimeLibraryInternal
{
	constexpr int32 CellCount = 120;
	constexpr int32 StickersPerCell = 63;
	constexpr int32 TwistStickerCount = StickersPerCell - 1;
	constexpr int32 StateSize = CellCount * StickersPerCell;
	constexpr int32 MoveCount = CellCount * TwistStickerCount;
	constexpr int32 HeaderSize = 24;
	constexpr uint8 RawMagic[8] = {'H', 'T', 'M', '1', '2', '0', 'P', '1'};

	uint16 ReadUInt16(const uint8* Bytes)
	{
		return static_cast<uint16>(Bytes[0])
			| (static_cast<uint16>(Bytes[1]) << 8);
	}

	uint32 ReadUInt32(const uint8* Bytes)
	{
		return static_cast<uint32>(Bytes[0])
			| (static_cast<uint32>(Bytes[1]) << 8)
			| (static_cast<uint32>(Bytes[2]) << 16)
			| (static_cast<uint32>(Bytes[3]) << 24);
	}

	struct FPermutationTable
	{
		TArray<uint8> Raw;
		TArray<int32> MoveOffsets;
		FString Error;
		bool bValid = false;
	};

	FPermutationTable BuildPermutationTable()
	{
		FPermutationTable Table;
		FString Encoded;
		Encoded.Reserve(
			HyperTwistMagic120CellGenerated::CompressedPermutationTableSize
			* 4 / 3 + 8);
		for (const TCHAR* Chunk :
			 HyperTwistMagic120CellGenerated::CompressedPermutationTableBase64Chunks)
		{
			Encoded.Append(Chunk);
		}

		TArray<uint8> Compressed;
		if (!FBase64::Decode(Encoded, Compressed)
			|| Compressed.Num()
				!= HyperTwistMagic120CellGenerated::CompressedPermutationTableSize)
		{
			Table.Error = TEXT("compressed permutation table failed Base64 validation");
			return Table;
		}

		Table.Raw.SetNumUninitialized(
			HyperTwistMagic120CellGenerated::RawPermutationTableSize);
		if (!FCompression::UncompressMemory(
				NAME_Zlib,
				Table.Raw.GetData(),
				Table.Raw.Num(),
				Compressed.GetData(),
				Compressed.Num()))
		{
			Table.Raw.Reset();
			Table.Error = TEXT("compressed permutation table failed zlib validation");
			return Table;
		}

		if (Table.Raw.Num() < HeaderSize
			|| FMemory::Memcmp(Table.Raw.GetData(), RawMagic, UE_ARRAY_COUNT(RawMagic)) != 0
			|| ReadUInt32(Table.Raw.GetData() + 8) != CellCount
			|| ReadUInt32(Table.Raw.GetData() + 12) != StickersPerCell
			|| ReadUInt32(Table.Raw.GetData() + 16) != MoveCount
			|| ReadUInt32(Table.Raw.GetData() + 20) != StateSize)
		{
			Table.Raw.Reset();
			Table.Error = TEXT("permutation table header is invalid");
			return Table;
		}

		Table.MoveOffsets.Reserve(MoveCount);
		int32 Cursor = HeaderSize;
		for (int32 MoveIndex = 0; MoveIndex < MoveCount; ++MoveIndex)
		{
			if (Cursor + static_cast<int32>(sizeof(uint16)) > Table.Raw.Num())
			{
				Table.Raw.Reset();
				Table.MoveOffsets.Reset();
				Table.Error = TEXT("permutation table ended before a move header");
				return Table;
			}
			Table.MoveOffsets.Add(Cursor);
			const int32 PairCount = ReadUInt16(Table.Raw.GetData() + Cursor);
			Cursor += sizeof(uint16);
			const int32 PairBytes = PairCount * sizeof(uint16) * 2;
			if (PairCount < 1
				|| PairCount > StateSize
				|| Cursor + PairBytes > Table.Raw.Num())
			{
				Table.Raw.Reset();
				Table.MoveOffsets.Reset();
				Table.Error = FString::Printf(
					TEXT("permutation table move %d is malformed"),
					MoveIndex);
				return Table;
			}

			for (int32 PairIndex = 0; PairIndex < PairCount; ++PairIndex)
			{
				const uint8* PairBytesAddress =
					Table.Raw.GetData() + Cursor + PairIndex * 4;
				if (ReadUInt16(PairBytesAddress) >= StateSize
					|| ReadUInt16(PairBytesAddress + 2) >= StateSize)
				{
					Table.Raw.Reset();
					Table.MoveOffsets.Reset();
					Table.Error = FString::Printf(
						TEXT("permutation table move %d references an invalid slot"),
						MoveIndex);
					return Table;
				}
			}
			Cursor += PairBytes;
		}

		if (Cursor != Table.Raw.Num())
		{
			Table.Raw.Reset();
			Table.MoveOffsets.Reset();
			Table.Error = TEXT("permutation table contains trailing bytes");
			return Table;
		}
		Table.bValid = true;
		return Table;
	}

	const FPermutationTable& GetPermutationTable()
	{
		static const FPermutationTable Table = BuildPermutationTable();
		return Table;
	}

	FHyperTwistPuzzleDefinitionRef BuildDefinition()
	{
		FHyperTwistPuzzleDefinitionRef Definition;
		Definition.PuzzleId = TEXT("polychoron/magic120cell/full-color");
		Definition.PuzzleFamily = EHyperTwistPuzzleFamily::Other;
		Definition.Dimension = 4;
		Definition.DefinitionVersion = TEXT("2026.07");
		Definition.NotationProfile = TEXT("magic120cell-cell-sticker-axis-v1");
		Definition.SizeVector = {CellCount, StickersPerCell};
		Definition.Variant = TEXT("exact-7560-facelet-permutation");
		return Definition;
	}

	bool IsSupportedDefinition(const FHyperTwistPuzzleDefinitionRef& Definition)
	{
		return Definition.IsStructurallyValid()
			&& Definition.PuzzleId == TEXT("polychoron/magic120cell/full-color")
			&& Definition.PuzzleFamily == EHyperTwistPuzzleFamily::Other
			&& Definition.Dimension == 4
			&& Definition.SizeVector.Num() == 2
			&& Definition.SizeVector[0] == CellCount
			&& Definition.SizeVector[1] == StickersPerCell;
	}

	bool IsSolvedState(const FHyperTwistMagic120CellRuntimeState& State)
	{
		if (State.StickerColorIndices.Num() != StateSize)
		{
			return false;
		}
		for (int32 SlotIndex = 0; SlotIndex < StateSize; ++SlotIndex)
		{
			if (State.StickerColorIndices[SlotIndex]
				!= SlotIndex / StickersPerCell)
			{
				return false;
			}
		}
		return true;
	}

	bool IsStateShapeValid(const FHyperTwistMagic120CellRuntimeState& State)
	{
		if (State.StateProfile != TEXT("magic120cell-7560-facelet-permutation-v1")
			|| !IsSupportedDefinition(State.Definition)
			|| State.StickerColorIndices.Num() != StateSize
			|| State.AppliedMoveCount < 0)
		{
			return false;
		}

		int32 ColorCounts[CellCount] = {};
		for (const int32 ColorIndex : State.StickerColorIndices)
		{
			if (ColorIndex < 0 || ColorIndex >= CellCount)
			{
				return false;
			}
			++ColorCounts[ColorIndex];
		}
		for (const int32 ColorCount : ColorCounts)
		{
			if (ColorCount != StickersPerCell)
			{
				return false;
			}
		}
		return true;
	}

	int32 GetMoveIndex(const FHyperTwistMagic120CellTurnRequest& Request)
	{
		return Request.CellIndex * TwistStickerCount + Request.StickerIndex - 1;
	}

	FString BuildNotation(const FHyperTwistMagic120CellTurnRequest& Request)
	{
		return FString::Printf(
			TEXT("C%03d/S%02d%s"),
			Request.CellIndex + 1,
			Request.StickerIndex,
			Request.bInverse ? TEXT("'") : TEXT(""));
	}
}

bool FHyperTwistMagic120CellRuntimeState::IsStructurallyValid() const
{
	using namespace HyperTwistMagic120CellRuntimeLibraryInternal;
	return IsStateShapeValid(*this)
		&& bIsSolved == IsSolvedState(*this);
}

bool FHyperTwistMagic120CellTurnRequest::IsStructurallyValid() const
{
	return CellIndex >= 0
		&& CellIndex <
			HyperTwistMagic120CellRuntimeLibraryInternal::CellCount
		&& StickerIndex >= 1
		&& StickerIndex <
			HyperTwistMagic120CellRuntimeLibraryInternal::StickersPerCell;
}

bool FHyperTwistMagic120CellProjectedCell::IsStructurallyValid() const
{
	using namespace HyperTwistMagic120CellRuntimeLibraryInternal;
	if (CellIndex < 0
		|| CellIndex >= CellCount
		|| StickerColorIndices.Num() != StickersPerCell
		|| RepresentativeColorIndex < 0
		|| RepresentativeColorIndex >= CellCount
		|| DisplacedStickerCount < 0
		|| DisplacedStickerCount > StickersPerCell
		|| bCellSolved != (DisplacedStickerCount == 0))
	{
		return false;
	}
	for (const int32 ColorIndex : StickerColorIndices)
	{
		if (ColorIndex < 0 || ColorIndex >= CellCount)
		{
			return false;
		}
	}
	return true;
}

bool FHyperTwistMagic120CellProjection::IsStructurallyValid() const
{
	using namespace HyperTwistMagic120CellRuntimeLibraryInternal;
	if (ProjectionProfile
			!= TEXT("magic120cell-120-cell-7560-facelet-projection-v1")
		|| Cells.Num() != CellCount)
	{
		return false;
	}

	TBitArray<> SeenCells(false, CellCount);
	int32 ProjectedStickerCount = 0;
	for (const FHyperTwistMagic120CellProjectedCell& Cell : Cells)
	{
		if (!Cell.IsStructurallyValid() || SeenCells[Cell.CellIndex])
		{
			return false;
		}
		SeenCells[Cell.CellIndex] = true;
		ProjectedStickerCount += Cell.StickerColorIndices.Num();
	}
	return ProjectedStickerCount == StateSize;
}

FHyperTwistPuzzleDefinitionRef
UHyperTwistMagic120CellRuntimeLibrary::MakePuzzleDefinition()
{
	return HyperTwistMagic120CellRuntimeLibraryInternal::BuildDefinition();
}

FHyperTwistMagic120CellRuntimeState
UHyperTwistMagic120CellRuntimeLibrary::BuildSolvedState()
{
	using namespace HyperTwistMagic120CellRuntimeLibraryInternal;
	FHyperTwistMagic120CellRuntimeState State;
	State.Definition = MakePuzzleDefinition();
	State.StickerColorIndices.Reserve(StateSize);
	for (int32 SlotIndex = 0; SlotIndex < StateSize; ++SlotIndex)
	{
		State.StickerColorIndices.Add(SlotIndex / StickersPerCell);
	}
	State.AppliedMoveCount = 0;
	State.bIsSolved = true;
	return State;
}

FHyperTwistPuzzleState
UHyperTwistMagic120CellRuntimeLibrary::BuildPuzzleStateEnvelope(
	const FHyperTwistMagic120CellRuntimeState& RuntimeState
)
{
	const FHyperTwistMagic120CellRuntimeState SafeState =
		RuntimeState.IsStructurallyValid()
			? RuntimeState
			: BuildSolvedState();
	FHyperTwistPuzzleState PuzzleState;
	PuzzleState.Definition = SafeState.Definition;
	PuzzleState.StateEncodingKind = EHyperTwistStateEncodingKind::Facelet;
	PuzzleState.StateEncoding.EncodingProfile = SafeState.StateProfile;
	PuzzleState.StateEncoding.PayloadJson =
		SerializeRuntimeStateToJson(SafeState);
	PuzzleState.OrientationFrame.Reference =
		TEXT("magic120cell-canonical-cell-sticker-frame-v1");
	PuzzleState.bIsSolved = SafeState.bIsSolved;
	PuzzleState.Source = EHyperTwistStateSource::Runtime;
	return PuzzleState;
}

FHyperTwistMagic120CellTurnResult
UHyperTwistMagic120CellRuntimeLibrary::ApplyTurn(
	const FHyperTwistMagic120CellRuntimeState& State,
	const FHyperTwistMagic120CellTurnRequest& Request
)
{
	using namespace HyperTwistMagic120CellRuntimeLibraryInternal;
	FHyperTwistMagic120CellTurnResult Result;
	Result.State = State;
	if (!State.IsStructurallyValid())
	{
		Result.Warnings.Add(TEXT("invalid-magic120cell-runtime-state"));
		return Result;
	}
	if (!Request.IsStructurallyValid())
	{
		Result.Warnings.Add(TEXT("invalid-magic120cell-turn-request"));
		return Result;
	}

	const FPermutationTable& Table = GetPermutationTable();
	if (!Table.bValid)
	{
		Result.Warnings.Add(FString::Printf(
			TEXT("magic120cell-permutation-table-unavailable: %s"),
			*Table.Error));
		return Result;
	}

	const int32 MoveIndex = GetMoveIndex(Request);
	const int32 MoveOffset = Table.MoveOffsets[MoveIndex];
	const int32 PairCount =
		ReadUInt16(Table.Raw.GetData() + MoveOffset);
	const uint8* PairData =
		Table.Raw.GetData() + MoveOffset + sizeof(uint16);
	Result.State.StickerColorIndices = State.StickerColorIndices;
	for (int32 PairIndex = 0; PairIndex < PairCount; ++PairIndex)
	{
		const uint8* Pair = PairData + PairIndex * 4;
		const int32 Destination = ReadUInt16(Pair);
		const int32 Source = ReadUInt16(Pair + 2);
		if (Request.bInverse)
		{
			Result.State.StickerColorIndices[Source] =
				State.StickerColorIndices[Destination];
		}
		else
		{
			Result.State.StickerColorIndices[Destination] =
				State.StickerColorIndices[Source];
		}
	}

	Result.State.AppliedMoveCount = State.AppliedMoveCount + 1;
	Result.State.bIsSolved = IsSolvedState(Result.State);
	Result.AppliedNotation = BuildNotation(Request);
	Result.MovedStickerCount = PairCount;
	Result.bApplied = true;
	Result.bExactStateUpdate = Result.State.IsStructurallyValid();
	if (!Result.bExactStateUpdate)
	{
		Result.Warnings.Add(TEXT("magic120cell-state-became-invalid"));
	}
	return Result;
}

FHyperTwistMagic120CellScrambleResult
UHyperTwistMagic120CellRuntimeLibrary::GenerateScramble(
	const int32 MoveCount,
	const int32 RandomSeed
)
{
	using namespace HyperTwistMagic120CellRuntimeLibraryInternal;
	FHyperTwistMagic120CellScrambleResult Result;
	if (MoveCount < 1 || MoveCount > 1000)
	{
		Result.Warnings.Add(TEXT("invalid-magic120cell-scramble-request"));
		return Result;
	}
	FString TableError;
	if (!IsPermutationTableReady(TableError))
	{
		Result.Warnings.Add(FString::Printf(
			TEXT("magic120cell-permutation-table-unavailable: %s"),
			*TableError));
		return Result;
	}

	FRandomStream Random(RandomSeed);
	Result.State = BuildSolvedState();
	Result.Moves.Reserve(MoveCount);
	Result.AppliedNotation.Reserve(MoveCount);
	FHyperTwistMagic120CellTurnRequest PreviousRequest;
	bool bHasPreviousRequest = false;
	for (int32 MoveIndex = 0; MoveIndex < MoveCount; ++MoveIndex)
	{
		FHyperTwistMagic120CellTurnRequest Request;
		do
		{
			Request.CellIndex = Random.RandRange(0, CellCount - 1);
			Request.StickerIndex = Random.RandRange(1, StickersPerCell - 1);
			Request.bInverse = Random.RandRange(0, 1) == 1;
		}
		while (bHasPreviousRequest
			&& Request.CellIndex == PreviousRequest.CellIndex
			&& Request.StickerIndex == PreviousRequest.StickerIndex);

		const FHyperTwistMagic120CellTurnResult Turn =
			ApplyTurn(Result.State, Request);
		if (!Turn.bApplied || !Turn.bExactStateUpdate)
		{
			Result.Warnings = Turn.Warnings;
			Result.Warnings.Add(TEXT("magic120cell-scramble-turn-failed"));
			Result.State = FHyperTwistMagic120CellRuntimeState();
			Result.Moves.Reset();
			Result.AppliedNotation.Reset();
			return Result;
		}

		Result.State = Turn.State;
		Result.Moves.Add(Request);
		Result.AppliedNotation.Add(Turn.AppliedNotation);
		PreviousRequest = Request;
		bHasPreviousRequest = true;
	}

	Result.bExactStateUpdate = Result.State.IsStructurallyValid();
	Result.bGenerated = Result.bExactStateUpdate
		&& Result.Moves.Num() == MoveCount
		&& !Result.State.bIsSolved;
	if (!Result.bGenerated)
	{
		Result.Warnings.Add(TEXT("magic120cell-scramble-result-invalid"));
	}
	return Result;
}

FHyperTwistMagic120CellProjectionBuildResult
UHyperTwistMagic120CellRuntimeLibrary::BuildProjection(
	const FHyperTwistMagic120CellRuntimeState& State
)
{
	using namespace HyperTwistMagic120CellRuntimeLibraryInternal;
	FHyperTwistMagic120CellProjectionBuildResult Result;
	if (!State.IsStructurallyValid())
	{
		Result.Warnings.Add(TEXT("invalid-magic120cell-runtime-state"));
		return Result;
	}

	Result.Projection.Cells.Reserve(CellCount);
	for (int32 CellIndex = 0; CellIndex < CellCount; ++CellIndex)
	{
		FHyperTwistMagic120CellProjectedCell Cell;
		Cell.CellIndex = CellIndex;
		Cell.StickerColorIndices.Reserve(StickersPerCell);
		int32 ForeignColorCounts[CellCount] = {};
		for (int32 StickerIndex = 0;
			 StickerIndex < StickersPerCell;
			 ++StickerIndex)
		{
			const int32 ColorIndex =
				State.StickerColorIndices[
					CellIndex * StickersPerCell + StickerIndex];
			Cell.StickerColorIndices.Add(ColorIndex);
			if (ColorIndex != CellIndex)
			{
				++Cell.DisplacedStickerCount;
				++ForeignColorCounts[ColorIndex];
			}
		}

		Cell.RepresentativeColorIndex = CellIndex;
		int32 StrongestForeignColorCount = 0;
		for (int32 ColorIndex = 0; ColorIndex < CellCount; ++ColorIndex)
		{
			if (ForeignColorCounts[ColorIndex] > StrongestForeignColorCount)
			{
				StrongestForeignColorCount = ForeignColorCounts[ColorIndex];
				Cell.RepresentativeColorIndex = ColorIndex;
			}
		}
		Cell.bCellSolved = Cell.DisplacedStickerCount == 0;
		Result.Projection.Cells.Add(MoveTemp(Cell));
	}

	Result.bProjected = Result.Projection.IsStructurallyValid();
	Result.bExactProjection = Result.bProjected;
	if (!Result.bProjected)
	{
		Result.Warnings.Add(TEXT("magic120cell-projection-invalid"));
	}
	return Result;
}

int32 UHyperTwistMagic120CellRuntimeLibrary::GetTurnOrderForSticker(
	const int32 StickerIndex
)
{
	if (StickerIndex >= 1 && StickerIndex <= 12)
	{
		return 5;
	}
	if (StickerIndex >= 13 && StickerIndex <= 42)
	{
		return 2;
	}
	if (StickerIndex >= 43 && StickerIndex <= 62)
	{
		return 3;
	}
	return 0;
}

FString UHyperTwistMagic120CellRuntimeLibrary::SerializeRuntimeStateToJson(
	const FHyperTwistMagic120CellRuntimeState& RuntimeState
)
{
	FString Json;
	FJsonObjectConverter::UStructToJsonObjectString(
		FHyperTwistMagic120CellRuntimeState::StaticStruct(),
		&RuntimeState,
		Json,
		0,
		0);
	return Json;
}

bool UHyperTwistMagic120CellRuntimeLibrary::
	TryDeserializeRuntimeStateFromJson(
		const FString& Json,
		FHyperTwistMagic120CellRuntimeState& OutRuntimeState
	)
{
	FHyperTwistMagic120CellRuntimeState ParsedState;
	if (Json.IsEmpty()
		|| !FJsonObjectConverter::JsonObjectStringToUStruct(
			Json,
			&ParsedState,
			0,
			0)
		|| !ParsedState.IsStructurallyValid())
	{
		return false;
	}
	OutRuntimeState = MoveTemp(ParsedState);
	return true;
}

bool UHyperTwistMagic120CellRuntimeLibrary::IsPermutationTableReady(
	FString& OutError
)
{
	const HyperTwistMagic120CellRuntimeLibraryInternal::FPermutationTable& Table =
		HyperTwistMagic120CellRuntimeLibraryInternal::GetPermutationTable();
	OutError = Table.Error;
	return Table.bValid;
}
