#include "bsp/UQLightgridUpdaterComponent.h"
#include "bsp/AQLightgridHolder.h"
#include "Engine/World.h"
#include "EngineUtils.h"
#include "Components/DecalComponent.h"
#include "Components/BillboardComponent.h"

#include "Kismet/KismetMathLibrary.h"

static int cellIndexFor(FQLightgridMeta meta, int gridX, int gridY, int gridZ);
static FColor ambientColorForCell(UDataTable* cells, int cell);
static FVector colorToVector(FColor color);

UQLightgridUpdaterComponent::UQLightgridUpdaterComponent(const FObjectInitializer& initializer) : Super(initializer)
{
	PrimaryComponentTick.bCanEverTick = true;
	lightgridParamName = FName("lightgrid");
}

void UQLightgridUpdaterComponent::BeginPlay()
{
	Super::BeginPlay();

	// overwrite our meta key with the level-appropriate key
	for (TActorIterator<AQLightgridHolder> iter(GetWorld()); iter; ++iter)
	{
		AQLightgridHolder* holder = *iter;
		if (!holder)
			continue;

		this->metaRowIdx = holder->lightgridKey;

		// override meta too, if we need to.
		if(holder->lightgridMetaTable)
			this->lightgridMeta = holder->lightgridMetaTable;
		break;
	}

	if(!lightgridMeta)
	{
		UE_LOG(LogTemp, Warning, TEXT("UQLightgridUpdaterComponent::BeginPlay() - Failed to find lightgrid meta table! Set it in the QLightgridHolder, subclass this per-project, or specify it in the constructor of the actor who owns it."));
		this->DestroyComponent();
		return;
	}

	FQLightgridMeta* meta = lightgridMeta->FindRow<FQLightgridMeta>(FName(FString::FromInt(metaRowIdx)), FString(""));
	if (!meta)
	{
		UE_LOG(LogTemp, Warning, TEXT("UQLightgridUpdaterComponent::BeginPlay() - Failed to find meta data for lightgrid!"));

		this->DestroyComponent();
		return;
	}

	// set cell table from meta reference
	FSoftObjectPath cellTablePath = meta->cellTableRef;
	
	cellTablePath.TryLoad();
	this->lightgridCells = Cast<UDataTable>(cellTablePath.ResolveObject());
	if (!this->lightgridCells)
	{
		UE_LOG(LogTemp, Warning, TEXT("UQLightgridUpdaterComponent::BeginPlay() - Failed to find cell table for lightgrid!"));

		this->DestroyComponent();
		return;
	}

	// if the owner is not movable, we can just set the material parameter once and be done with it.
	PrimaryComponentTick.bCanEverTick = GetOwner()->IsRootComponentMovable();
	
	if(!PrimaryComponentTick.bCanEverTick)
		updateMaterials(GetOwner(), FVector::Zero());
}

void UQLightgridUpdaterComponent::updateMaterials(AActor* actor, FVector colorOverride)
{
	TArray<UMeshComponent*> meshes;

	FVector ambient = colorOverride;
	if (colorOverride == FVector::ZeroVector)
	{
		FVector loc = actor->GetActorLocation();
		loc.Z += actor->GetSimpleCollisionHalfHeight();
		ambient = getAmbientColorForLocation(loc);
	}

	// set material parameter on all meshes
	actor->GetComponents<UMeshComponent>(meshes);
	for (UMeshComponent* mesh : meshes)
	{
		FVector color = getAmbientColorForLocation(mesh->GetComponentLocation());
		mesh->SetVectorParameterValueOnMaterials(lightgridParamName, color);
	}

	// also do it for decals
	TArray<UDecalComponent*> decals;
	actor->GetComponents<UDecalComponent>(decals);
	for (UDecalComponent* decal : decals)
	{
		UMaterialInterface* mat = decal->GetDecalMaterial();
		if (!mat)
			continue;

		UMaterialInstanceDynamic* dynamicMat = Cast<UMaterialInstanceDynamic>(mat);
		if (!dynamicMat)
			dynamicMat = decal->CreateDynamicMaterialInstance();

		dynamicMat->SetVectorParameterValue(lightgridParamName, ambient);
	}

	// and for any primitive (e.g., billboards)
	TArray<UPrimitiveComponent*> primitives;
	actor->GetComponents<UPrimitiveComponent>(primitives);
	for (UPrimitiveComponent* primitive : primitives)
		for (int i = 0; i < primitive->GetNumMaterials(); i++)
		{
			UMaterialInterface* mat = primitive->GetMaterial(i);
			if (!mat)
				continue;

			UMaterialInstanceDynamic* dynamicMat = Cast<UMaterialInstanceDynamic>(mat);
			if (!dynamicMat)
				dynamicMat = primitive->CreateDynamicMaterialInstance(i, primitive->GetMaterial(i));
			 
			dynamicMat->SetVectorParameterValue(lightgridParamName, ambient);
		}	

	// recursively do this.
	for (AActor* child : actor->Children)
		updateMaterials(child, ambient);

	// for attached actors too.
	TArray<AActor*> attached;
	actor->GetAttachedActors(attached);
	for (AActor* child : attached)
		updateMaterials(child, ambient);
}

FVector UQLightgridUpdaterComponent::getAmbientColorForLocation(FVector loc)
{
	// lightgrid is stored in Quake coordinates, let's convert ours to that.
	// this is done here rather than in the export process, because it's a real pain to rearrange every grid cell.
	loc.Y *= -1;

	FQLightgridMeta* meta = lightgridMeta->FindRow<FQLightgridMeta>(FName(FString::FromInt(metaRowIdx)), FString(""));

	// ref: https://github.com/id-Software/Quake-III-Arena/blob/dbe4ddb10315479fc00086f08e25d968b4b43c49/code/renderer/tr_light.c#L145-L157
	// ref: https://github.com/id-Software/Quake-III-Arena/blob/master/code/renderer/tr_bsp.c#L1653-L1664
	
	// lightGridOrigin[i] = w->lightGridSize[i] * ceil( wMins[i] / w->lightGridSize[i] );
	// lightOrigin = loc - lightGridOrigin
	// v = lightOrigin * lightGridInverseSize;
	// pos = floor(v);
	// frac = v - pos;

	FVector lightGridSize = FVector(64, 64, 128);
	FVector lightGridInverseSize = FVector(1.0f / lightGridSize.X, 1.0f / lightGridSize.Y, 1.0f / lightGridSize.Z);

	FVector lightGridOrigin;
	lightGridOrigin.X = lightGridSize.X * FMath::CeilToFloat(meta->mapMins[0] / lightGridSize.X);
	lightGridOrigin.Y = lightGridSize.Y * FMath::CeilToFloat(meta->mapMins[1] / lightGridSize.Y);
	lightGridOrigin.Z = lightGridSize.Z * FMath::CeilToFloat(meta->mapMins[2] / lightGridSize.Z);

	FVector lightOrigin = loc - lightGridOrigin;
	FVector v = lightOrigin * lightGridInverseSize;
	
	// if not interpolating, just grab the cell we're in.
	if (!useTrilerp)
	{
		int gridX = FMath::FloorToInt(v.X);
		int gridY = FMath::FloorToInt(v.Y);
		int gridZ = FMath::FloorToInt(v.Z);

		int cellIdx = cellIndexFor(*meta, gridX, gridY, gridZ);
		
		return colorToVector(ambientColorForCell(lightgridCells, cellIndexFor(*meta, gridX, gridY, gridZ)));
	}

	// otherwise, trilerp.
	if (trilerpSteps > 0)
	{
		v.X = FMath::RoundToFloat(v.X * trilerpSteps) / trilerpSteps;
		v.Y = FMath::RoundToFloat(v.Y * trilerpSteps) / trilerpSteps;
		v.Z = FMath::RoundToFloat(v.Z * trilerpSteps) / trilerpSteps;
	}

	return trilerpCellColors(v, *meta);
}

FVector UQLightgridUpdaterComponent::trilerpCellColors(FVector v, FQLightgridMeta meta)
{
	FVector outputAmbient = FVector::ZeroVector;
	FVector outputDirectional = FVector::ZeroVector;

	// pos = floor(v)
	int pos[3];
	pos[0] = FMath::FloorToInt(v.X);
	pos[1] = FMath::FloorToInt(v.Y);
	pos[2] = FMath::FloorToInt(v.Z);

	// frac = v - pos
	FVector frac = v;
	frac.X -= pos[0];
	frac.Y -= pos[1];
	frac.Z -= pos[2];

	// clamp pos
	for(int i = 0; i < 3; i++)
		pos[i] = FMath::Clamp(pos[i], 0, meta.cellCounts[i] - 1);

	int gridStep[3];

	gridStep[0] = 1;
	gridStep[1] = meta.cellCounts[0];
	gridStep[2] = meta.cellCounts[0] * meta.cellCounts[1];

	int gridData = (pos[0] * gridStep[0]) + (pos[1] * gridStep[1]) + (pos[2] * gridStep[2]);
	float totalFactor = 0;

	for (int i = 0; i < 8; i++)
	{
		float factor = 1;
		int data = gridData;

		// this is a really fancy way of sampling all 8 cells around a center cell, by using this bit counter to ignore certain components based on the i value.
		for (int j = 0; j < 3; j++)
		{
			if (i & (1 << j))
			{
				factor *= frac[j];
				data += gridStep[j];
			}
			else
			{
				factor *= (1.0f - frac[j]);
			}

			FQLightgridCell* cell = lightgridCells->FindRow<FQLightgridCell>(FName(FString::FromInt(data)), FString());

			// importantly, don't sample cells for which there is no color. Helps prevent artifacts where a nearby cell is inside a wall, or outside the void.
			if(!cell || (cell->ambientColor.R + cell->ambientColor.G + cell->ambientColor.B) == 0)
				continue;

			totalFactor += factor;

			outputAmbient.X += factor * cell->ambientColor.R;
			outputAmbient.Y += factor * cell->ambientColor.G;
			outputAmbient.Z += factor * cell->ambientColor.B;

			outputDirectional.X += factor * cell->directionalColor.R;
			outputDirectional.Y += factor * cell->directionalColor.G;
			outputDirectional.Z += factor * cell->directionalColor.B;
		}
	}

	// in q3 this check is inverted, but here it totally screws it, with most light values being >1. This seems like it would only ever be required to normalize in the case where they _aren't_ 0-1.
	if (totalFactor <= 0 || totalFactor >= 1)
	{
		totalFactor = 1.0f / totalFactor;
		outputAmbient *= totalFactor;
		outputDirectional *= totalFactor;		
	}

	// normally the directional would be calculated per-frame based on direction to light, but in our case we just want a whole-model approximation
	// so we take part of the directional.
	FVector outputCombined = outputAmbient + (outputDirectional / 2);
	FColor ret = FColor(
		FMath::Clamp(outputCombined.X, 0, 255),
		FMath::Clamp(outputCombined.Y, 0, 255),
		FMath::Clamp(outputCombined.Z, 0, 255)
	);

	if (ret == FColor::Black)
		ret = ambientColorForCell(lightgridCells, cellIndexFor(meta, pos[0], pos[1], pos[2]));

	return colorToVector(ret);
}

void UQLightgridUpdaterComponent::TickComponent(float delta, enum ELevelTick tick, FActorComponentTickFunction* tickFunc)
{
	// if our attachment owner already has one of these components, do nothing.
	AActor* attachOwner = GetOwner()->GetAttachParentActor();
	while (attachOwner)
	{
		if (attachOwner->GetComponentByClass(UQLightgridUpdaterComponent::StaticClass()))
			return;
		attachOwner = attachOwner->GetAttachParentActor();
	}

	updateMaterials(GetOwner(), FVector::Zero());
}

int cellIndexFor(FQLightgridMeta meta, int gridX, int gridY, int gridZ)
{
	return gridZ * meta.cellCounts[0] * meta.cellCounts[1] + gridY * meta.cellCounts[0] + gridX;
}

FColor ambientColorForCell(UDataTable* cells, int cellIdx)
{
	FQLightgridCell* cell = cells->FindRow<FQLightgridCell>(FName(FString::FromInt(cellIdx)), FString());
	if (!cell)
		return FColor(0, 0, 0);

	// q3 uses directional (only if it's positive) as an addition to ambient
	// ref https://github.com/id-Software/Quake-III-Arena/blob/master/code/cgame/cg_players.c#L2192-L2219

	short r = FMath::Clamp(cell->ambientColor.R + cell->directionalColor.R, 0, 255);
	short g = FMath::Clamp(cell->ambientColor.G + cell->directionalColor.G, 0, 255);
	short b = FMath::Clamp(cell->ambientColor.B + cell->directionalColor.B, 0, 255);

	return FColor(r, g, b);
}

FVector colorToVector(FColor color)
{
	return UKismetMathLibrary::Conv_LinearColorToVector(UKismetMathLibrary::Conv_ColorToLinearColor(color));
}