#include "Systems/DescentSonar.h"
#include "Audio/DescentAudio.h"
#include "Core/DescentMath.h"
#include "Core/DescentSimSubsystem.h"
#include "Core/Tuning/DescentTune_World.h"
#include "Debug/DescentCommandUtil.h"
#include "Debug/DescentStateRegistry.h"
#include "Player/DescentCharacter.h"
#include "Player/DescentPlayerController.h"
#include "Player/DescentTools.h"
#include "Systems/DescentCraft.h"
#include "Systems/DescentHorror.h"
#include "Systems/DescentSave.h"
#include "UI/DescentHudSubsystem.h"
#include "UI/DescentUiSubsystem.h"
#include "Vehicles/DescentCarrier.h"
#include "Vehicles/KazahTender.h"
#include "Vehicles/TenderFlight.h"
#include "World/DescentTerrainSubsystem.h"
#include "World/DescentWorldDef.h"
#include "DescentUnits.h"
#include "TheDescent.h"

#include "Engine/World.h"
#include "HAL/IConsoleManager.h"
#include "Misc/Compression.h"

namespace DS = Descent::Sonar;
namespace WT = Descent::Tune::World;

namespace Descent::Sonar
{
	FVector FibonacciDir(int32 I, int32 N)
	{
		const double Z = 1.0 - 2.0 * (I + 0.5) / FMath::Max(1, N);
		const double R = FMath::Sqrt(FMath::Max(0.0, 1.0 - Z * Z)), Phi = I * 2.399963229728653;   // the golden angle
		return FVector(R * FMath::Cos(Phi), R * FMath::Sin(Phi), Z);
	}

	bool CellOf(const FVector& P, FIntVector& Out)
	{
		Out.X = FMath::FloorToInt32((P.X - WT::ClaimX0) / CellXY);
		Out.Y = FMath::FloorToInt32((P.Y - WT::ClaimY0) / CellXY);
		Out.Z = FMath::FloorToInt32((P.Z - WT::ClaimZ0) / CellZ);
		return Out.X >= 0 && Out.X < MapNx && Out.Y >= 0 && Out.Y < MapNy && Out.Z >= 0 && Out.Z < MapNz;
	}
}

// ---- the survey map -----------------------------------------------------------------------------------------------------

UDescentMapSubsystem* UDescentMapSubsystem::Get(const UWorld* World)
{
	return World ? World->GetSubsystem<UDescentMapSubsystem>() : nullptr;
}

void UDescentMapSubsystem::OnWorldBeginPlay(UWorld& InWorld)
{
	Super::OnWorldBeginPlay(InWorld);
	Known.Init(false, DS::MapNx * DS::MapNy * DS::MapNz);
	Visited.Init(false, DS::MapNx * DS::MapNy);
	if (UDescentSimSubsystem* Sim = UDescentSimSubsystem::Get(&InWorld))
	{
		FDescentSimRegisterOptions Opt;
		Opt.Owner = TEXT("Map");
		Opt.SubOrder = 20;
		Opt.RateHz = 4.0;
		StepHandle = Sim->RegisterStep(EDescentSimStep::Sonar, [this](const FDescentSimFrame& F) { Step(F.Dt); }, Opt);
	}
	if (UDescentSaveSubsystem* Save = UDescentSaveSubsystem::Get(&InWorld))
	{
		// the survey bitset, compressed: the set cells as indices
		Save->OnCapture.AddWeakLambda(this, [this](UDescentSaveGame* S)
		{
			TArray<int32> Idx;
			for (TConstSetBitIterator<> It(Known); It; ++It) { Idx.Add(It.GetIndex()); }
			Idx.Add(-1);
			for (TConstSetBitIterator<> It(Visited); It; ++It) { Idx.Add(It.GetIndex()); }
			S->MapSurvey.SetNumUninitialized(Idx.Num() * sizeof(int32));
			FMemory::Memcpy(S->MapSurvey.GetData(), Idx.GetData(), S->MapSurvey.Num());
		});
		Save->OnApply.AddWeakLambda(this, [this](const UDescentSaveGame* S)
		{
			Known.Init(false, DS::MapNx * DS::MapNy * DS::MapNz);
			Visited.Init(false, DS::MapNx * DS::MapNy);
			KnownCells = VisitedCells = 0;
			const int32 N = S->MapSurvey.Num() / sizeof(int32);
			const int32* Idx = reinterpret_cast<const int32*>(S->MapSurvey.GetData());
			bool bSecond = false;
			for (int32 I = 0; I < N; ++I)
			{
				if (Idx[I] < 0) { bSecond = true; continue; }
				if (!bSecond && Known.IsValidIndex(Idx[I])) { Known[Idx[I]] = true; ++KnownCells; }
				if (bSecond && Visited.IsValidIndex(Idx[I])) { Visited[Idx[I]] = true; ++VisitedCells; }
			}
		});
	}
}

void UDescentMapSubsystem::Deinitialize()
{
	if (UDescentSimSubsystem* Sim = UDescentSimSubsystem::Get(GetWorld())) { Sim->Unregister(StepHandle); }
	Super::Deinitialize();
}

void UDescentMapSubsystem::Mark(const FVector& P)
{
	FIntVector C;
	if (!DS::CellOf(P, C)) { return; }
	const int32 I = (C.Z * DS::MapNy + C.Y) * DS::MapNx + C.X;
	if (!Known[I]) { Known[I] = true; ++KnownCells; }
}

bool UDescentMapSubsystem::IsKnown(const FVector& P) const
{
	FIntVector C;
	return DS::CellOf(P, C) && Known[(C.Z * DS::MapNy + C.Y) * DS::MapNx + C.X];
}

void UDescentMapSubsystem::Visit(const FVector& P)
{
	FIntVector C;
	if (!DS::CellOf(P, C)) { return; }
	const int32 I = C.Y * DS::MapNx + C.X;
	if (!Visited[I]) { Visited[I] = true; ++VisitedCells; }
}

void UDescentMapSubsystem::Step(double Dt)
{
	UWorld* World = GetWorld();
	const ADescentPlayerController* PC = World ? Cast<ADescentPlayerController>(World->GetFirstPlayerController()) : nullptr;
	if (PC) { Visit(PC->RefPointM()); }
	if (const ADescentCarrier* C = UDescentCarrierSubsystem::GetCarrier(World)) { Visit(C->GetActorLocation() / 100.0); }
	WindowTimer -= Dt;
	if (WindowTimer <= 0.0) { WindowTimer = 1.0; FillWindow(); }
}

void UDescentMapSubsystem::FillWindow()
{
	UWorld* World = GetWorld();
	UDescentUiSubsystem* Ui = UDescentUiSubsystem::Get(World);
	if (!Ui || Ui->IsDemo() || Ui->GetWindow() != EDescentWindow::Pda || Ui->GetTab() != TEXT("map")) { return; }
	const ADescentPlayerController* PC = Cast<ADescentPlayerController>(World->GetFirstPlayerController());
	const ADescentCarrier* Carrier = UDescentCarrierSubsystem::GetCarrier(World);
	const FVector Ref = PC ? PC->RefPointM() : FVector::ZeroVector;
	const double RefDepth = Descent::DepthOfZ(Ref.Z);
	// the slice of depth that is drawn
	double D0 = 0.0, D1 = 1e9;
	FDescentMapViewModel& M = Ui->Models().Map;
	if (SliceMode == 0) { D0 = RefDepth - 60.0; D1 = RefDepth + 60.0; M.SliceText = FString::Printf(TEXT("AUTO  %d m +/- 60"), FMath::RoundToInt32(RefDepth)); }
	else if (SliceMode == 2) { D0 = SliceDepthM - 40.0; D1 = SliceDepthM + 40.0; M.SliceText = FString::Printf(TEXT("%d m +/- 40"), FMath::RoundToInt32(SliceDepthM)); }
	else { M.SliceText = TEXT("ALL"); }
	const int32 Z1 = FMath::Clamp(FMath::FloorToInt32((Descent::DatumM - D0 - WT::ClaimZ0) / DS::CellZ), 0, DS::MapNz - 1);   // shallow end (top)
	const int32 Z0 = FMath::Clamp(FMath::FloorToInt32((Descent::DatumM - D1 - WT::ClaimZ0) / DS::CellZ), 0, DS::MapNz - 1);
	const int32 Step = 4;
	M.GridW = DS::MapNx / Step;
	M.GridH = DS::MapNy / Step;
	M.Cells.SetNumUninitialized(M.GridW * M.GridH);
	for (int32 Gy = 0; Gy < M.GridH; ++Gy)
	{
		for (int32 Gx = 0; Gx < M.GridW; ++Gx)
		{
			float V = -1.f;
			int32 Top = -1;
			bool bVisited = false;
			for (int32 Y = Gy * Step; Y < (Gy + 1) * Step; ++Y)
			{
				for (int32 X = Gx * Step; X < (Gx + 1) * Step; ++X)
				{
					bVisited |= Visited[Y * DS::MapNx + X];
					for (int32 Z = Z1; Z >= Z0 && Z > Top; --Z) { if (Known[(Z * DS::MapNy + Y) * DS::MapNx + X]) { Top = Z; break; } }
				}
			}
			if (Top >= 0)
			{
				const double Depth = Descent::DatumM - (WT::ClaimZ0 + (Top + 0.5) * DS::CellZ);
				V = static_cast<float>(FMath::Clamp((Depth - 1740.0) / 1460.0, 0.0, 1.0));
			}
			else if (bVisited) { V = 2.f; }
			M.Cells[Gy * M.GridW + Gx] = V;
		}
	}
	auto Norm = [](const FVector& P) { return FVector2f(static_cast<float>((P.X - WT::ClaimX0) / (WT::ClaimX1 - WT::ClaimX0)), static_cast<float>((P.Y - WT::ClaimY0) / (WT::ClaimY1 - WT::ClaimY0))); };
	M.You = Norm(Ref);
	M.Carrier = Carrier ? Norm(Carrier->GetActorLocation() / 100.0) : M.You;
	M.TargetText = TEXT("TARGET STRATUM  3,105 m (Tethys survey)");
	M.Markers.Reset();
	auto AddMarker = [&M, &Norm](const FString& Label, const FVector& P, const FLinearColor& Color, bool bRing)
	{
		FDescentMapMarker Mk;
		Mk.Label = Label;
		const FVector2f N = Norm(P);
		Mk.X = N.X; Mk.Y = N.Y; Mk.Color = Color; Mk.bRing = bRing;
		M.Markers.Add(Mk);
	};
	if (const UDescentToolsSubsystem* Tools = UDescentToolsSubsystem::Get(World)) { for (int32 I = 0; I < Tools->Markers.Num(); ++I) { AddMarker(FString::Printf(TEXT("BEACON %d"), I + 1), Tools->Markers[I], FLinearColor(1.0f, 0.7f, 0.2f), false); } }
	if (const UDescentHorrorSubsystem* Horror = UDescentHorrorSubsystem::Get(World)) { for (int32 I = 0; I < Horror->NumBeacons(); ++I) { if (Horror->IsBeaconFound(I)) { AddMarker(FString::Printf(TEXT("DEAD BEACON %d"), I + 1), Horror->BeaconPosM(I), FLinearColor(1.0f, 0.25f, 0.2f), false); } } }
	if (const UDescentSonarSubsystem* Sonar = UDescentSonarSubsystem::Get(World)) { for (const UDescentSonarSubsystem::FContact& K : Sonar->GetContacts()) { if (K.Kind == 2) { AddMarker(K.Label, K.P, FLinearColor(1.0f, 0.2f, 0.15f), true); } } }
	// place names appear once the place is known or visited
	for (const Descent::World::FPlace& Pl : Descent::World::FDescentWorldDef::Get().Places)
	{
		FIntVector C;
		if (IsKnown(Pl.PosM) || IsKnown(Pl.LookAtM) || (DS::CellOf(Pl.PosM, C) && Visited[C.Y * DS::MapNx + C.X])) { AddMarker(Pl.Name.ToString().ToUpper(), Pl.PosM, FLinearColor(0.75f, 0.85f, 0.9f), false); }
	}
	Ui->NotifyChanged();
}

// ---- sonar ----------------------------------------------------------------------------------------------------------------

UDescentSonarSubsystem* UDescentSonarSubsystem::Get(const UWorld* World)
{
	return World ? World->GetSubsystem<UDescentSonarSubsystem>() : nullptr;
}

void UDescentSonarSubsystem::OnWorldBeginPlay(UWorld& InWorld)
{
	Super::OnWorldBeginPlay(InWorld);
	if (UDescentSimSubsystem* Sim = UDescentSimSubsystem::Get(&InWorld))
	{
		FDescentSimRegisterOptions Opt;
		Opt.Owner = TEXT("Sonar");
		StepHandle = Sim->RegisterStep(EDescentSimStep::Sonar, [this](const FDescentSimFrame& F) { Step(F.Dt); }, Opt);
		LateHandle = Sim->RegisterLateStep(EDescentLateStep::HudModel, [this](const FDescentSimFrame&) { WriteHud(); }, Opt);
	}
}

void UDescentSonarSubsystem::Deinitialize()
{
	if (UDescentSimSubsystem* Sim = UDescentSimSubsystem::Get(GetWorld())) { Sim->Unregister(StepHandle); Sim->Unregister(LateHandle); }
	Super::Deinitialize();
}

int32 UDescentSonarSubsystem::PointsAlive() const
{
	int32 N = 0;
	for (const FPoint& P : Points) { N += (Now >= P.Born && Now < P.Born + DS::PointLifeSec) ? 1 : 0; }
	return N;
}

int32 UDescentSonarSubsystem::RaysPending() const
{
	int32 N = 0;
	for (const FPing& P : Active) { N += P.Rays - P.Next; }
	return N;
}

bool UDescentSonarSubsystem::Ping(const FVector& OriginM, double RangeM, int32 Rays)
{
	FPing P;
	P.O = OriginM;
	P.Range = RangeM;
	P.Rays = Rays;
	P.T0 = Now;
	P.Rot = FQuat(Rng.VRand(), Rng.FRandRange(0.0, 6.283));   // a random rotation per ping
	Active.Add(P);
	++Pings;
	LastOriginM = OriginM;
	LastRangeM = RangeM;
	if (UDescentAudioSubsystem* Audio = UDescentAudioSubsystem::Get(GetWorld())) { Audio->OneShot(TEXT("ping")); }
	if (UDescentHorrorSubsystem* Horror = UDescentHorrorSubsystem::Get(GetWorld())) { Horror->OnSonarPing(OriginM, RangeM); }
	return true;
}

bool UDescentSonarSubsystem::PingFromPlayer(FString* OutWhyNot)
{
	UWorld* World = GetWorld();
	ADescentPlayerController* PC = World ? Cast<ADescentPlayerController>(World->GetFirstPlayerController()) : nullptr;
	auto No = [OutWhyNot](const TCHAR* Why) { if (OutWhyNot) { *OutWhyNot = Why; } return false; };
	if (!PC) { return No(TEXT("no player")); }
	if (CooldownLeft > 0.0) { return No(TEXT("the sonar is recharging")); }
	if (PC->GetMode() == EDescentMode::Sub && PC->GetTender())
	{
		AKazahTender* T = PC->GetTender();
		const Descent::Tender::FTenderStats St = T->Flight->Stats();
		if (T->Flight->PowerPct <= 0.0) { return No(TEXT("no power")); }
		T->Flight->PowerPct = FMath::Max(0.0, T->Flight->PowerPct - 100.0 * 1.5 / St.Capacity);
		CooldownLeft = DS::TenderCooldownSec;
		return Ping(T->GetActorLocation() / 100.0 + T->GetActorForwardVector() * 1.5, T->Flight->Upgrades.bSonarModule ? DS::TenderModuleRangeM : DS::TenderRangeM, DS::TenderRays);
	}
	if (PC->GetMode() == EDescentMode::Helm)
	{
		const ADescentCarrier* C = UDescentCarrierSubsystem::GetCarrier(World);
		if (!C) { return No(TEXT("no carrier")); }
		CooldownLeft = DS::HelmCooldownSec;
		return Ping(C->ToWorldUU(FVector(40.0, 0.0, -10.0)) / 100.0, DS::HelmRangeM, DS::HelmRays);
	}
	if (PC->GetMode() == EDescentMode::Swim)
	{
		const UDescentCraftSubsystem* Craft = UDescentCraftSubsystem::Get(World);
		if (!Craft || Craft->Count(TEXT("hand_sonar")) <= 0) { return No(TEXT("no handheld sonar (fabricator, EXPLORATION)")); }
		CooldownLeft = DS::HandCooldownSec;
		return Ping(PC->RefPointM(), DS::HandRangeM, DS::HandRays);
	}
	return No(TEXT("sonar works from a Tender, the helm, or swimming with the handheld sonar"));
}

void UDescentSonarSubsystem::AddContact(const FVector& Pm, double LifeSec, int32 Kind, const FString& Label)
{
	FContact K;
	K.P = Pm; K.Born = Now; K.Life = LifeSec; K.Kind = Kind; K.Label = Label;
	Contacts.Add(K);
}

void UDescentSonarSubsystem::Step(double Dt)
{
	Now += Dt;
	CooldownLeft = FMath::Max(0.0, CooldownLeft - Dt);
	UDescentTerrainSubsystem* Terrain = UDescentTerrainSubsystem::Get(GetWorld());
	UDescentMapSubsystem* Map = UDescentMapSubsystem::Get(GetWorld());
	// 420 rays per sim step across all active pings
	int32 Budget = DS::RaysPerStep;
	for (int32 A = 0; A < Active.Num() && Budget > 0 && Terrain; ++A)
	{
		FPing& P = Active[A];
		while (P.Next < P.Rays && Budget > 0)
		{
			const FVector Dir = P.Rot.RotateVector(DS::FibonacciDir(P.Next++, P.Rays));
			--Budget;
			const double D = Terrain->Raycast(P.O + Dir * 1.5, Dir, P.Range - 1.5, 0.9);
			if (D <= 0.0) { continue; }
			FPoint Pt;
			Pt.P = P.O + Dir * (D + 1.5);
			Pt.Dist = static_cast<float>(D + 1.5);
			Pt.Born = P.T0 + Pt.Dist / DS::RevealMps;   // the 170 m/s reveal
			if (Points.Num() < DS::MaxPoints) { Points.Add(Pt); } else { Points[Head] = Pt; Head = (Head + 1) % DS::MaxPoints; }
			++Hits;
			if (Map) { Map->Mark(Pt.P); }
		}
	}
	Active.RemoveAll([](const FPing& P) { return P.Next >= P.Rays; });
	Contacts.RemoveAll([this](const FContact& K) { return Now >= K.Born + K.Life; });
}

void UDescentSonarSubsystem::WriteHud()
{
	UDescentHudSubsystem* Hud = UDescentHudSubsystem::Get(GetWorld());
	if (!Hud) { return; }
	TArray<FDescentHudSonarPoint>& Out = Hud->Model().Sonar;
	Out.Reset();
	const int32 Stride = FMath::Max(1, Points.Num() / 9000);   // never more than 9,000 on screen
	for (int32 I = 0; I < Points.Num(); I += Stride)
	{
		const FPoint& P = Points[I];
		const double Age = Now - P.Born;
		if (Age < 0.0 || Age >= DS::PointLifeSec) { continue; }
		// it fades from 45 % of its life to the end; in its first second it flashes
		const double Life = Age / DS::PointLifeSec;
		float Alpha = static_cast<float>(Life < 0.45 ? 1.0 : 1.0 - (Life - 0.45) / 0.55);
		const float Far = FMath::Clamp(P.Dist / 140.f, 0.f, 1.f);
		FLinearColor Col = FMath::Lerp(FLinearColor(0.45f, 1.0f, 1.0f), FLinearColor(0.22f, 0.5f, 1.0f), Far);
		if (Age < 1.0) { Col = FMath::Lerp(FLinearColor::White, Col, static_cast<float>(Age)); }
		Col.A = Alpha * 0.9f;
		Out.Add({ P.P * 100.0, Col, FMath::Lerp(4.5f, 2.4f, Far) });
	}
	for (const FContact& K : Contacts)
	{
		const FLinearColor Col = K.Kind == 3 ? FLinearColor(1.0f, 0.7f, 0.2f) : FLinearColor(1.0f, 0.15f, 0.1f);
		Out.Add({ K.P * 100.0, Col, 9.f });
	}
}

// ---- commands and the state sections `sonar`, `map` (section 9.4) -----------------------------------------------------------

namespace
{
	void CmdSonarPing(const TArray<FString>& Args, UWorld* World)
	{
		UDescentSonarSubsystem* S = UDescentSonarSubsystem::Get(World);
		FString Why = TEXT("no sonar subsystem");
		if (S && Args.Num() >= 3)
		{
			TArray<double> N;
			if (Descent::Cmd::ParseNumbers(Args, 0, 3, N)) { S->Ping(FVector(N[0], N[1], N[2]), Args.Num() > 3 ? FCString::Atod(*Args[3]) : 160.0, Args.Num() > 4 ? FCString::Atoi(*Args[4]) : 2200); return; }
		}
		if (!S || !S->PingFromPlayer(&Why)) { Descent::Cmd::Fail(World, TEXT("Descent.Sonar.Ping"), *Why); }
	}
	void CmdMapSlice(const TArray<FString>& Args, UWorld* World)
	{
		UDescentMapSubsystem* M = UDescentMapSubsystem::Get(World);
		if (!M || Args.Num() < 1) { Descent::Cmd::Fail(World, TEXT("Descent.Map.Slice"), TEXT("usage: Descent.Map.Slice auto | all | <depth m>")); return; }
		if (Args[0].Equals(TEXT("auto"), ESearchCase::IgnoreCase)) { M->SliceMode = 0; }
		else if (Args[0].Equals(TEXT("all"), ESearchCase::IgnoreCase)) { M->SliceMode = 1; }
		else { M->SliceMode = 2; M->SliceDepthM = FCString::Atod(*Args[0]); }
		M->FillWindow();
	}
	FAutoConsoleCommandWithWorldAndArgs GCmdSonarPing(TEXT("Descent.Sonar.Ping"), TEXT("Descent.Sonar.Ping [x y z [range [rays]]] - ping from the player's vehicle (R), or from a world point"), FConsoleCommandWithWorldAndArgsDelegate::CreateStatic(&CmdSonarPing));
	FAutoConsoleCommandWithWorldAndArgs GCmdMapSlice(TEXT("Descent.Map.Slice"), TEXT("Descent.Map.Slice auto | all | <depth> - the depth slice of the survey map"), FConsoleCommandWithWorldAndArgsDelegate::CreateStatic(&CmdMapSlice));
}

static FDescentStateProviderRegistrar GSonarStateSection(TEXT("sonar"), [](UWorld* World, FJsonObject& Out)
{
	const UDescentSonarSubsystem* S = UDescentSonarSubsystem::Get(World);
	Out.SetBoolField(TEXT("available"), S != nullptr);
	if (!S) { return; }
	Out.SetNumberField(TEXT("points_alive"), S->PointsAlive());
	Out.SetNumberField(TEXT("pings"), S->Pings);
	Out.SetNumberField(TEXT("hits"), S->Hits);
	Out.SetNumberField(TEXT("rays_pending"), S->RaysPending());
	Out.SetNumberField(TEXT("contacts"), S->NumContacts());
	Out.SetNumberField(TEXT("cooldown_s"), Descent::Json::Round(S->CooldownLeft));
	Out.SetNumberField(TEXT("last_range_m"), S->LastRangeM);
});

static FDescentStateProviderRegistrar GMapStateSection(TEXT("map"), [](UWorld* World, FJsonObject& Out)
{
	const UDescentMapSubsystem* M = UDescentMapSubsystem::Get(World);
	Out.SetBoolField(TEXT("available"), M != nullptr);
	if (!M) { return; }
	Out.SetNumberField(TEXT("known_cells"), M->KnownCells);
	Out.SetNumberField(TEXT("visited_cells"), M->VisitedCells);
	Out.SetStringField(TEXT("slice"), M->SliceMode == 0 ? TEXT("auto") : (M->SliceMode == 1 ? TEXT("all") : TEXT("fixed")));
	TArray<TSharedPtr<FJsonValue>> Markers;
	if (const UDescentToolsSubsystem* Tools = UDescentToolsSubsystem::Get(World)) { for (const FVector& P : Tools->Markers) { Markers.Add(Descent::Json::Vec3(P)); } }
	Out.SetArrayField(TEXT("markers"), Markers);
});