#include "World/DescentWorldSlice.h"
#include "World/DescentRoutes.h"
#include "World/DescentTerrainSubsystem.h"
#include "World/DescentWorldDef.h"
#include "Core/Tuning/DescentTune_World.h"
#include "Debug/DescentStateRegistry.h"

#include "Async/ParallelFor.h"
#include "Engine/World.h"
#include "HAL/FileManager.h"
#include "ImageUtils.h"
#include "Misc/FileHelper.h"

namespace WTune = Descent::Tune::World;

namespace Descent::World
{
	FColor SliceOreColor(int32 OreType)
	{
		return OreGlowColor(static_cast<EOreType>(OreType)).ToFColorSRGB();
	}

	bool WriteSlicePng(UWorld* World, TCHAR Axis, double Coord, const FString& Stem, double Mpp, FString& OutPathOrError)
	{
		if (!FDescentStateRegistry::IsValidFileStem(Stem))
		{
			OutPathOrError = FString::Printf(TEXT("'%s' is not a valid file name (use letters, digits, _ - . only)"), *Stem);
			return false;
		}
		Axis = FChar::ToUpper(Axis);
		if (Axis != TEXT('X') && Axis != TEXT('Y') && Axis != TEXT('Z')) { OutPathOrError = TEXT("axis must be X, Y or Z"); return false; }
		if (Mpp < 0.25 || Mpp > 20.0) { OutPathOrError = TEXT("metres per pixel must be between 0.25 and 20"); return false; }

		const FDescentWorldDef& Def = FDescentWorldDef::Get();
		const FVector Lo(WTune::ClaimX0, WTune::ClaimY0, WTune::ClaimZ0), Hi(WTune::ClaimX1, WTune::ClaimY1, WTune::ClaimZ1);
		// Horizontal axis index (0 X, 1 Y, 2 Z) and vertical axis index of the image; the plane's own axis is `Fixed`.
		int32 Fixed, Horiz, Vert;
		bool bFlipVert;
		if (Axis == TEXT('Z')) { Fixed = 2; Horiz = 1; Vert = 0; bFlipVert = true; }        // plan: Y right, X up
		else if (Axis == TEXT('Y')) { Fixed = 1; Horiz = 0; Vert = 2; bFlipVert = true; }   // X right, Z up
		else { Fixed = 0; Horiz = 1; Vert = 2; bFlipVert = true; }                          // Y right, Z up
		const double H0 = Lo[Horiz], H1 = Hi[Horiz], V0 = Lo[Vert], V1 = Hi[Vert];
		const int32 W = FMath::CeilToInt((H1 - H0) / Mpp), H = FMath::CeilToInt((V1 - V0) / Mpp);
		TArray<FColor> Pix;
		Pix.SetNumUninitialized(W * H);
		TArray<uint8> Rock;
		Rock.SetNumZeroed(W * H);

		auto PointOf = [&](int32 Px, int32 Py)
		{
			FVector P;
			P[Fixed] = Coord;
			P[Horiz] = H0 + (Px + 0.5) * Mpp;
			P[Vert] = bFlipVert ? V1 - (Py + 0.5) * Mpp : V0 + (Py + 0.5) * Mpp;
			return P;
		};

		ParallelFor(H, [&](int32 Py)
		{
			TMap<FIntVector, TSharedPtr<FPrimList>> Cache;
			for (int32 Px = 0; Px < W; ++Px)
			{
				const FVector P = PointOf(Px, Py);
				const FIntVector Cell(FMath::FloorToInt(P.X / WTune::GenQuickCell), FMath::FloorToInt(P.Y / WTune::GenQuickCell), FMath::FloorToInt(P.Z / WTune::GenQuickCell));
				TSharedPtr<FPrimList>* F = Cache.Find(Cell);
				if (!F)
				{
					TSharedPtr<FPrimList> L = MakeShared<FPrimList>();
					Def.Sdf().Gather((FVector(Cell) + FVector(0.5)) * WTune::GenQuickCell, WTune::GenQuickCell * 0.87 + 2.0, *L);
					F = &Cache.Add(Cell, L);
				}
				const bool bRock = Def.Sdf().GenRaw(P, **F) > 0.0;
				Rock[Px + Py * W] = bRock ? 1 : 0;
			}
		});
		for (int32 I = 0; I < W * H; ++I)
		{
			Pix[I] = Rock[I] ? FColor(255, 255, 255, 255) : FColor(0, 0, 0, 255);
		}

		// Rock carved away: where a modified chunk now says water but the untouched field said rock.
		if (UDescentTerrainSubsystem* T = UDescentTerrainSubsystem::Get(World))
		{
			for (int32 Py = 0; Py < H; ++Py)
			{
				for (int32 Px = 0; Px < W; ++Px)
				{
					if (!Rock[Px + Py * W]) { continue; }
					const FVector P = PointOf(Px, Py);
					FTerrainChunk* C = T->FindChunk(UDescentTerrainSubsystem::KeyOfPoint(P));
					if (C && C->bModified && T->Density(P) <= 0.0) { Pix[Px + Py * W] = FColor(128, 128, 128, 255); }
				}
			}
		}

		auto Plot = [&](const FVector& P, double RadiusM, FColor Col)
		{
			if (FMath::Abs(P[Fixed] - Coord) > RadiusM + 0.5) { return; }
			const double Slice = FMath::Sqrt(FMath::Max(0.0, RadiusM * RadiusM - FMath::Square(P[Fixed] - Coord)));
			const double Rpx = FMath::Max(1.0, Slice / Mpp);
			const int32 Cx = FMath::RoundToInt((P[Horiz] - H0) / Mpp - 0.5);
			const int32 Cy = FMath::RoundToInt(bFlipVert ? (V1 - P[Vert]) / Mpp - 0.5 : (P[Vert] - V0) / Mpp - 0.5);
			for (int32 Dy = -FMath::CeilToInt(Rpx); Dy <= FMath::CeilToInt(Rpx); ++Dy)
			{
				for (int32 Dx = -FMath::CeilToInt(Rpx); Dx <= FMath::CeilToInt(Rpx); ++Dx)
				{
					if (Dx * Dx + Dy * Dy > Rpx * Rpx) { continue; }
					const int32 X = Cx + Dx, Y = Cy + Dy;
					if (X >= 0 && Y >= 0 && X < W && Y < H) { Pix[X + Y * W] = Col; }
				}
			}
		};

		// Ore veins: points along each capsule that lie within the plane's reach, drawn in the ore colour.
		for (const FDeposit& D : Def.Deposits)
		{
			if (FMath::Min(D.A[Fixed], D.B[Fixed]) - D.R > Coord || FMath::Max(D.A[Fixed], D.B[Fixed]) + D.R < Coord) { continue; }
			const FColor Col = SliceOreColor(static_cast<int32>(D.Type));
			for (int32 K = 0; K <= 8; ++K) { Plot(FMath::Lerp(D.A, D.B, K / 8.0), D.R * 1.5, Col); }
		}
		// Carrier route axes: the three routes' waypoint lines and the well axes, in red.
		const FColor Red(255, 40, 40, 255);
		for (int32 R = 1; R <= FDescentRoutes::NumRoutes; ++R)
		{
			const TArray<FRouteWaypoint>& Wp = FDescentRoutes::Waypoints(R);
			for (int32 I = 1; I < Wp.Num(); ++I)
			{
				const int32 N = FMath::Max(1, FMath::CeilToInt(FVector::Dist(Wp[I - 1].P, Wp[I].P) / 2.0));
				for (int32 K = 0; K <= N; ++K) { Plot(FMath::Lerp(Wp[I - 1].P, Wp[I].P, static_cast<double>(K) / N), 2.5, Red); }
			}
		}
		for (const FWell& Wl : Def.Wells)
		{
			for (double Z = Wl.Z0; Z <= Wl.Z1; Z += 2.0) { Plot(FVector(Wl.X, Wl.Y, Z), 2.5, Red); }
		}

		TArray64<uint8> Png;
		FImageUtils::PNGCompressImageArray(W, H, TArrayView64<const FColor>(Pix.GetData(), Pix.Num()), Png);
		const FString Path = FDescentStateRegistry::GetShotDir() / (Stem + TEXT(".png"));
		IFileManager::Get().MakeDirectory(*FPaths::GetPath(Path), true);
		if (Png.Num() == 0 || !FFileHelper::SaveArrayToFile(Png, *Path))
		{
			OutPathOrError = FString::Printf(TEXT("could not write %s"), *Path);
			return false;
		}
		OutPathOrError = Path;
		return true;
	}
}