#include "Systems/DescentAutomation.h"
#include "Core/Tuning/DescentTune_Cargo.h"
#include "Player/DescentCharacter.h"
#include "Systems/DescentCargo.h"
#include "UI/DescentHudSubsystem.h"
#include "Vehicles/DescentCarrier.h"
#include "DescentUnits.h"
#include "TheDescent.h"

#include "Components/StaticMeshComponent.h"
#include "Engine/World.h"

namespace ATune = Descent::Tune::Cargo;

namespace Descent::Cargo
{
	bool ChooseMove(TConstArrayView<FSocketView> Sockets, int32& OutFrom, int32& OutTo)
	{
		OutFrom = OutTo = INDEX_NONE;
		int32 BestFromRank = MAX_int32;
		for (int32 F = 0; F < Sockets.Num(); ++F)
		{
			const FSocketView& From = Sockets[F];
			if (From.BoxCount <= 0 || From.Rank >= 2 || From.Rank >= BestFromRank) { continue; }   // outputs and buffers only; the lowest rank wins
			int32 BestTo = INDEX_NONE;
			for (int32 T = 0; T < Sockets.Num(); ++T)
			{
				const FSocketView& To = Sockets[T];
				if (T == F || To.bOccupied || To.Rank <= From.Rank || To.Capacity < From.BoxCount) { continue; }
				if (BestTo == INDEX_NONE || To.Rank > Sockets[BestTo].Rank) { BestTo = T; }
			}
			if (BestTo != INDEX_NONE) { OutFrom = F; OutTo = BestTo; BestFromRank = From.Rank; }
		}
		return OutFrom != INDEX_NONE;
	}

	void TwoLinkIK(double Dist, double Height, double L1, double L2, double& OutShoulderDeg, double& OutElbowDeg)
	{
		const double D = FMath::Clamp(FMath::Sqrt(Dist * Dist + Height * Height), FMath::Abs(L1 - L2) + 0.01, L1 + L2 - 0.01);
		const double Interior = FMath::Acos(FMath::Clamp((L1 * L1 + L2 * L2 - D * D) / (2.0 * L1 * L2), -1.0, 1.0));
		const double Lift = FMath::Acos(FMath::Clamp((L1 * L1 + D * D - L2 * L2) / (2.0 * L1 * D), -1.0, 1.0));
		OutShoulderDeg = FMath::RadiansToDegrees(FMath::Atan2(Height, Dist) + Lift);
		OutElbowDeg = -FMath::RadiansToDegrees(UE_DOUBLE_PI - Interior);
	}
}

// ---- the rack ---------------------------------------------------------------------------------------------------------

void ADescentBoxRack::InitKit(const Descent::Build::FPlacement& P, FName InKit, int32 InId)
{
	KitId = InKit;
	PlaceId = InId;
	Place = P;
	// two shelves and a back plate, so the boxes on it stay visible
	const FLinearColor Steel(0.14f, 0.15f, 0.16f);
	Body = AddPart(TEXT("Body"), FVector(0, 0, 0.03), FVector(P.BodySizeM.X, P.BodySizeM.Y, 0.06), Steel, FLinearColor::Black, 0.f, true);
	AddPart(TEXT("Shelf"), FVector(0, 0, 0.52), FVector(P.BodySizeM.X, P.BodySizeM.Y, 0.04), Steel);
	AddPart(TEXT("Back"), FVector(0, -0.5 * P.BodySizeM.Y + 0.02, 0.5), FVector(P.BodySizeM.X, 0.04, 1.0), Steel, FLinearColor::Black, 0.f, true);
	if (UDescentInteractionRegistry* Reg = UDescentInteractionRegistry::Get(GetWorld())) { Reg->Register(this); }
	if (UDescentCargoSubsystem* Cargo = UDescentCargoSubsystem::Get(GetWorld())) { Cargo->Racks.AddUnique(this); }
}

void ADescentBoxRack::EndPlay(const EEndPlayReason::Type Reason)
{
	if (UDescentCargoSubsystem* Cargo = UDescentCargoSubsystem::Get(GetWorld())) { Cargo->Racks.Remove(this); }
	Super::EndPlay(Reason);
}

FVector ADescentBoxRack::SocketLocalM(int32 Slot) const
{
	return LocalM(FVector((Slot & 1) ? 0.35 : -0.35, 0.0, (Slot & 2) ? 0.74 : 0.26));
}

int32 ADescentBoxRack::FreeSocket() const
{
	for (int32 I = 0; I < 4; ++I) { if (SlotBox[I] == 0 && !SlotReserved[I]) { return I; } }
	return INDEX_NONE;
}

int32 ADescentBoxRack::NumBoxes() const
{
	int32 N = 0;
	for (int32 I = 0; I < 4; ++I) { N += (SlotBox[I] != 0 || SlotReserved[I]) ? 1 : 0; }
	return N;
}

bool ADescentBoxRack::PutBox(ADescentOreBox* InBox, int32 Slot)
{
	ADescentCarrier* C = UDescentCarrierSubsystem::GetCarrier(GetWorld());
	if (!InBox || !C || Slot < 0 || Slot > 3 || SlotBox[Slot] != 0) { return false; }
	InBox->AttachToActor(C, FAttachmentTransformRules::KeepWorldTransform);
	InBox->SetActorLocationAndRotation(C->ToWorldUU(SocketLocalM(Slot)), GetActorRotation());
	InBox->Where = FString::Printf(TEXT("rack%d.%d"), PlaceId, Slot);
	SlotBox[Slot] = InBox->BoxId;
	SlotReserved[Slot] = false;
	return true;
}

bool ADescentBoxRack::CanPickUp(FString* OutWhyNot) const
{
	if (NumBoxes() > 0) { if (OutWhyNot) { *OutWhyNot = TEXT("The rack holds a box"); } return false; }
	return true;
}

FString ADescentBoxRack::GetPrompt() const
{
	const UDescentCargoSubsystem* Cargo = UDescentCargoSubsystem::Get(GetWorld());
	if (const ADescentOreBox* B = Cargo ? Cargo->Carried() : nullptr) { return FString::Printf(TEXT("Put %s on the rack"), *B->Label()); }
	return TEXT("Box Rack - hold E to pick up");
}

void ADescentBoxRack::Use(ADescentCharacter* User)
{
	UDescentCargoSubsystem* Cargo = UDescentCargoSubsystem::Get(GetWorld());
	if (!Cargo || !Cargo->Carried()) { return; }
	const int32 Slot = FreeSocket();
	if (Slot == INDEX_NONE)
	{
		if (UDescentHudSubsystem* Hud = UDescentHudSubsystem::Get(GetWorld())) { Hud->Model().PushToast(TEXT("The rack is full"), EDescentHudToastKind::Warn); }
		return;
	}
	PutBox(Cargo->ReleaseCarried(), Slot);
}

// ---- the arm ----------------------------------------------------------------------------------------------------------

void ADescentRobotArm::InitKit(const Descent::Build::FPlacement& P, FName InKit, int32 InId)
{
	KitId = InKit;
	PlaceId = InId;
	Place = P;
	Mk = InKit == TEXT("arm2") ? 2 : 1;
	ReachM = Mk == 2 ? 4.5 : 3.2;
	CycleSec = Mk == 2 ? 2.5 : 4.0;
	L1 = L2 = 0.56 * ReachM;
	const FLinearColor Yellow(0.75f, 0.42f, 0.03f), Dark(0.08f, 0.08f, 0.09f);
	const double Sign = P.Surface == Descent::Build::Ceiling ? -1.0 : 1.0;
	Body = AddPart(TEXT("Base"), P.BodyOffsetM, P.BodySizeM, Dark, FLinearColor::Black, 0.f, true);
	auto Joint = [this](const TCHAR* Name, USceneComponent* Parent, const FVector& OffsetM)
	{
		USceneComponent* J = NewObject<USceneComponent>(this, Name);
		J->SetMobility(EComponentMobility::Movable);
		J->SetupAttachment(Parent);
		J->SetRelativeLocation(Descent::MetresToUU(OffsetM));
		J->RegisterComponent();
		return J;
	};
	YawJoint = Joint(TEXT("YawJoint"), Root, FVector(0, 0, Sign * P.BodySizeM.Z));
	AddPart(TEXT("Turret"), FVector::ZeroVector, FVector(0.4, 0.4, 0.2), Yellow, FLinearColor::Black, 0.f, false, YawJoint);
	Shoulder = Joint(TEXT("Shoulder"), YawJoint, FVector(0, 0, Sign * 0.1));
	AddPart(TEXT("UpperArm"), FVector(0.5 * L1, 0, 0), FVector(L1, 0.16, 0.16), Yellow, FLinearColor::Black, 0.f, false, Shoulder);
	Elbow = Joint(TEXT("Elbow"), Shoulder, FVector(L1, 0, 0));
	AddPart(TEXT("Forearm"), FVector(0.5 * L2, 0, 0), FVector(L2, 0.12, 0.12), Yellow, FLinearColor::Black, 0.f, false, Elbow);
	Wrist = Joint(TEXT("Wrist"), Elbow, FVector(L2, 0, 0));
	Gripper = AddPart(TEXT("Gripper"), FVector::ZeroVector, FVector(0.2, 0.5, 0.08), Dark, FLinearColor::Black, 0.f, false, Wrist);
	RestM = LocalM(FVector(0.5 * ReachM, 0.0, Sign * (P.BodySizeM.Z + 0.6)));
	Pose(RestM);
	if (UDescentInteractionRegistry* Reg = UDescentInteractionRegistry::Get(GetWorld())) { Reg->Register(this); }
	if (UDescentCargoSubsystem* Cargo = UDescentCargoSubsystem::Get(GetWorld())) { Cargo->Arms.AddUnique(this); }
}

void ADescentRobotArm::EndPlay(const EEndPlayReason::Type Reason)
{
	if (UDescentCargoSubsystem* Cargo = UDescentCargoSubsystem::Get(GetWorld())) { Cargo->Arms.Remove(this); }
	Super::EndPlay(Reason);
}

bool ADescentRobotArm::CanPickUp(FString* OutWhyNot) const
{
	if (IsMoving()) { if (OutWhyNot) { *OutWhyNot = TEXT("The arm holds a box"); } return false; }
	return true;
}

FString ADescentRobotArm::GetPrompt() const
{
	return FString::Printf(TEXT("Handling Arm Mk %s: %s - E power, hold E pick up"), Mk == 2 ? TEXT("II") : TEXT("I"), *Status);
}

void ADescentRobotArm::Use(ADescentCharacter* User)
{
	bPower = !bPower;
	if (UDescentHudSubsystem* Hud = UDescentHudSubsystem::Get(GetWorld())) { Hud->Model().PushToast(bPower ? TEXT("Arm on") : TEXT("Arm off (it finishes the box it holds)")); }
}

void ADescentRobotArm::Pose(const FVector& TargetLocalM)
{
	// the joint angles are visual only: two links in the vertical plane through the base and the target
	const FVector Rel = FRotator(0.0, -Place.YawDeg, 0.0).RotateVector(TargetLocalM - Place.PosM);
	const FVector S = Descent::UUToMetres(YawJoint->GetRelativeLocation() + Shoulder->GetRelativeLocation());
	const double Dist = FVector2D(Rel.X, Rel.Y).Size(), Height = Rel.Z - S.Z;
	double A = 0.0, B = 0.0;
	Descent::Cargo::TwoLinkIK(Dist, Height, L1, L2, A, B);
	YawJoint->SetRelativeRotation(FRotator(0.0, FMath::RadiansToDegrees(FMath::Atan2(Rel.Y, Rel.X)), 0.0));
	Shoulder->SetRelativeRotation(FRotator(A, 0.0, 0.0));
	Elbow->SetRelativeRotation(FRotator(B, 0.0, 0.0));
	Wrist->SetRelativeRotation(FRotator(-(A + B), 0.0, 0.0));   // the gripper stays level
}

void ADescentRobotArm::StepArm(double Dt)
{
	UDescentCargoSubsystem* Cargo = UDescentCargoSubsystem::Get(GetWorld());
	ADescentCarrier* C = UDescentCarrierSubsystem::GetCarrier(GetWorld());
	if (!Cargo || !C) { return; }
	if (IsMoving())
	{
		ADescentOreBox* B = Box.Get();
		if (!B) { CycleT = -1.0; return; }
		CycleT += Dt;
		const double F = FMath::Clamp(CycleFraction(), 0.0, 1.0);
		// reach and grip 25 %, swing 50 %, place 25 %
		const double Swing = FMath::Clamp((F - 0.25) / 0.5, 0.0, 1.0);
		FVector G = FMath::Lerp(FromM, ToM, Swing);
		G.Z += 0.7 * FMath::Sin(Swing * UE_DOUBLE_PI) + 0.24;   // the gripper sits on top of the box
		Pose(G);
		if (F >= 0.25) { B->SetActorLocation(C->ToWorldUU(G - FVector(0, 0, 0.24))); }
		if (F >= 1.0)
		{
			bool bPlaced = false;
			if (ADescentBoxRack* R = ToRack.Get()) { R->SlotReserved[ToSlot] = false; bPlaced = R->PutBox(B, ToSlot); }
			else if (ADescentRefinery* Ref = ToRefinery.Get())
			{
				Ref->Reserved = FMath::Max(0, Ref->Reserved - HeldCount);
				for (int32 I = 0; I < B->Count; ++I) { Ref->Queue.Add(B->OreType); }
				B->Destroy();
				bPlaced = true;
			}
			if (!bPlaced) { Cargo->DropBoxAt(B, ToM); }   // the target went away: the box goes onto the deck there
			++Moved;
			CycleT = -1.0;
			Box.Reset();
			ToRack.Reset();
			ToRefinery.Reset();
		}
		return;
	}
	Pose(RestM);
	if (!bPower) { Status = TEXT("OFF"); return; }

	// 1. sockets within reach: horizontal distance from the base at most the reach, at most 2 m above or below
	struct FSock { ADescentOreBox* Box = nullptr; ADescentBoxRack* Rack = nullptr; ADescentRefinery* Refinery = nullptr; int32 Slot = -1; FVector PosM = FVector::ZeroVector; };
	TArray<FSock> Socks;
	TArray<Descent::Cargo::FSocketView> Views;
	auto InReach = [this](const FVector& M) { return FVector2D(M.X - Place.PosM.X, M.Y - Place.PosM.Y).Size() <= ReachM && FMath::Abs(M.Z - Place.PosM.Z) <= 2.0 + (Place.Surface == Descent::Build::Ceiling ? 3.0 : 0.0); };
	for (int32 I = 0; I < ATune::OutfeedSlots; ++I)
	{
		const FVector M = Cargo->SlotLocalM(I);
		ADescentOreBox* B = Cargo->FindBox(Cargo->SlotBoxAt(I));
		if (!B || !InReach(M)) { continue; }
		Socks.Add({ B, nullptr, nullptr, I, M });
		Views.Add({ 0, true, B->Count, 0 });
	}
	for (ADescentBoxRack* R : Cargo->Racks)
	{
		if (!R) { continue; }
		for (int32 I = 0; I < 4; ++I)
		{
			const FVector M = R->SocketLocalM(I);
			if (!InReach(M)) { continue; }
			ADescentOreBox* B = Cargo->FindBox(R->SlotBox[I]);
			Socks.Add({ B, R, nullptr, I, M });
			Views.Add({ 1, B != nullptr || R->SlotReserved[I], B ? B->Count : 0, ATune::BoxSize });
		}
	}
	for (ADescentRefinery* Ref : Cargo->GetRefineries())
	{
		if (!Ref) { continue; }
		const FVector M = C->ToLocalM(Ref->HopperWorldUU());
		if (!InReach(M)) { continue; }
		Socks.Add({ nullptr, nullptr, Ref, 0, M });
		Views.Add({ 2, false, 0, ATune::RefineryQueueMax - Ref->Queue.Num() - Ref->Reserved });
	}
	if (Socks.Num() == 0) { Status = TEXT("IDLE - nothing in reach"); return; }
	// 2. the lowest-rank box some higher-rank socket accepts
	int32 From = INDEX_NONE, To = INDEX_NONE;
	if (!Descent::Cargo::ChooseMove(Views, From, To)) { Status = TEXT("IDLE - no free target"); return; }
	// 3. the cycle
	ADescentOreBox* B = Socks[From].Box;
	Cargo->ReleaseSockets(B->BoxId);
	B->Where = FString::Printf(TEXT("arm%d"), PlaceId);
	Box = B;
	HeldCount = B->Count;
	FromM = Socks[From].PosM;
	ToM = Socks[To].PosM;
	ToSlot = Socks[To].Slot;
	ToRack = Socks[To].Rack;
	ToRefinery = Socks[To].Refinery;
	if (ADescentBoxRack* R = ToRack.Get()) { R->SlotReserved[ToSlot] = true; }
	if (ADescentRefinery* Ref = ToRefinery.Get()) { Ref->Reserved += HeldCount; }
	CycleT = 0.0;
	Status = FString::Printf(TEXT("MOVING %s"), *B->Label());
}