#include "Subsystems/RRDADataAssetSubsystem.h"

#include "RefinedRDApi.h"
#include "DataAssets/RRDADirtDataAsset.h"
#include "Engine/AssetManager.h"
#include "AssetRegistry/AssetRegistryModule.h"
#include "DataAssets/RRDABoilerDataAsset.h"
#include "DataAssets/RRDACoolerDataAsset.h"
#include "DataAssets/RRDADataAssetDSModAdapter.h"
#include "DataAssets/RRDAHeaterDataAsset.h"
#include "DataAssets/RRDATurbineDataAsset.h"
#include "Logging/StructuredLog.h"

URRDADataAssetSubsystem::URRDADataAssetSubsystem() {}

URRDADataAssetSubsystem* URRDADataAssetSubsystem::Get(UObject* Context) {
	if (IsValid(Context)) {
		return Context->GetWorld()->GetGameInstance()->GetSubsystem<URRDADataAssetSubsystem>();
	}
	return nullptr;
}

URRDADataAssetSubsystem* URRDADataAssetSubsystem::GetChecked(UObject* Context) {
	URRDADataAssetSubsystem* Subsystem = Get(Context);
	fgcheck(Subsystem);
	return Subsystem;
}

void URRDADataAssetSubsystem::Initialize(FSubsystemCollectionBase& Collection) {
	Super::Initialize(Collection);

	FAssetRegistryModule& AssetRegistryModule = FModuleManager::LoadModuleChecked<FAssetRegistryModule>(FName("AssetRegistry"));
	IAssetRegistry& AssetRegistry = AssetRegistryModule.Get();

	// Must wait until all assets are discovered before populating list of assets.
	if (AssetRegistry.IsLoadingAssets()) {
		AssetRegistry.OnFilesLoaded().AddUObject(this, &URRDADataAssetSubsystem::StartScanForDataAssets);
	}
	else {
		StartScanForDataAssets();
	}
}

void URRDADataAssetSubsystem::Deinitialize() {
	Super::Deinitialize();

	mEnabledDataAssets.Empty();
	mDisabledDataAssets.Empty();
	mDirtAssets.Empty();
	mTurbineAssets.Empty();
	mBoilerAssets.Empty();
	mCoolerAssets.Empty();
	mHeaterAssets.Empty();
	mGeneratorAssets.Empty();
	mDSBuildingData.Empty();
}

void URRDADataAssetSubsystem::StartScanForDataAssets() {
	mDisabledDataAssets.Empty();
	mEnabledDataAssets.Empty();

	ReCacheBoilerDataAssets();
	ReCacheCoolerDataAssets();
	ReCacheDirtDataAssets();
	ReCacheHeaterDataAssets();
	ReCacheTurbineDataAssets();
	ReCacheGeneratorDataAssets();

	mDSBuildingData.Empty();
	ReCacheDSModAdapterDataAssets();
	ReCacheDSAdapterCoversDataAssets();
}

void URRDADataAssetSubsystem::ReCacheDirtDataAssets() {
	mDirtAssets.Empty();
	TSet<URRDADirtDataAsset*> DataAssets;
	if (FindAllDataAssetsOfClass(DataAssets)) {
		for (URRDADirtDataAsset* DataAsset : DataAssets) {
			mDirtAssets.Add(DataAsset->mItem, DataAsset);
		}
	}
	UE_LOG(LogRRDApi, Log, TEXT("ReCacheDirtDataAssets: %d"), mDirtAssets.Num());
}

void URRDADataAssetSubsystem::ReCacheTurbineDataAssets() {
	mTurbineAssets.Empty();
	TSet<URRDATurbineDataAsset*> DataAssets;
	int32 Added = 0;
	if (FindAllDataAssetsOfClass(DataAssets)) {
		for (URRDATurbineDataAsset* DataAsset : DataAssets) {
			if (!mTurbineAssets.Contains(DataAsset->mTier)) {
				mTurbineAssets.Add(DataAsset->mTier, FRRDADataHolderTurbineData());
			}

			fgcheckf(DataAsset->mItem.ItemClass, TEXT("TurbineDataAsset %s has no ItemClass"), *DataAsset->GetPathName());

			mTurbineAssets[DataAsset->mTier].Data.Add(DataAsset->mItem.ItemClass, DataAsset);
			Added++;
		}
	}

	UE_LOG(LogRRDApi, Log, TEXT("ReCacheTurbineDataAssets: %d"), Added);
}

void URRDADataAssetSubsystem::ReCacheBoilerDataAssets() {
	mBoilerAssets.Empty();
	TSet<URRDABoilerDataAsset*> DataAssets;
	int32 Added = 0;
	if (FindAllDataAssetsOfClass(DataAssets)) {
		for (URRDABoilerDataAsset* DataAsset : DataAssets) {
			if (!mBoilerAssets.Contains(DataAsset->mTier)) {
				mBoilerAssets.Add(DataAsset->mTier, FRRDADataHolderBoilerData());
			}

			fgcheckf(DataAsset->mInput.ItemClass, TEXT("BoilerDataAsset %s has no ItemClass"), *DataAsset->GetPathName());

			mBoilerAssets[DataAsset->mTier].Data.Add(DataAsset->mInput.ItemClass, DataAsset);
			Added++;
		}
	}

	UE_LOG(LogRRDApi, Log, TEXT("ReCacheBoilerDataAssets: %d"), Added);
}

void URRDADataAssetSubsystem::ReCacheCoolerDataAssets() {
	mCoolerAssets.Empty();
	TSet<URRDACoolerDataAsset*> DataAssets;
	int32 Added = 0;
	if (FindAllDataAssetsOfClass(DataAssets)) {
		for (URRDACoolerDataAsset* DataAsset : DataAssets) {
			if (!mCoolerAssets.Contains(DataAsset->mType)) {
				mCoolerAssets.Add(DataAsset->mType, FRRDADataHolderCoolerData());
			}

			fgcheckf(DataAsset->mItem, TEXT("CoolerDataAsset %s has no ItemClass"), *DataAsset->GetPathName());

			mCoolerAssets[DataAsset->mType].Data.Add(DataAsset->mItem, DataAsset);
			Added++;
		}
	}

	UE_LOG(LogRRDApi, Log, TEXT("ReCacheCoolerDataAssets: %d"), Added);
}

void URRDADataAssetSubsystem::ReCacheHeaterDataAssets() {
	mHeaterAssets.Empty();
	TSet<URRDAHeaterDataAsset*> DataAssets;
	int32 Added = 0;
	if (FindAllDataAssetsOfClass(DataAssets)) {
		for (URRDAHeaterDataAsset* DataAsset : DataAssets) {
			if (!mHeaterAssets.Contains(DataAsset->mType)) {
				mHeaterAssets.Add(DataAsset->mType, FRRDADataHolderHeaterData());
			}

			fgcheckf(DataAsset->mInput.ItemClass, TEXT("HeaterDataAsset %s has no ItemClass"), *DataAsset->GetPathName());

			mHeaterAssets[DataAsset->mType].Data.Add(DataAsset->mInput.ItemClass, DataAsset);
			Added++;
		}
	}

	UE_LOG(LogRRDApi, Log, TEXT("ReCacheHeaterDataAssets: %d"), Added);
}

void URRDADataAssetSubsystem::ReCacheGeneratorDataAssets() {
	mGeneratorAssets.Empty();
	TSet<URRDAGeneratorDataAsset*> DataAssets;
	if (FindAllDataAssetsOfClass(DataAssets)) {
		for (URRDAGeneratorDataAsset* DataAsset : DataAssets) {
			mGeneratorAssets.Add(DataAsset->mTier, DataAsset);
		}
	}

	UE_LOG(LogRRDApi, Log, TEXT("ReCacheGeneratorDataAssets: %d"), mGeneratorAssets.Num());
}

void URRDADataAssetSubsystem::ReCacheDSModAdapterDataAssets() {
	TSet<URRDADataAssetDSModAdapter*> DataAssets;
	if (FindAllDataAssetsOfClass(DataAssets)) {
		for (URRDADataAssetDSModAdapter* DataAsset : DataAssets) {
			FRRDADataHolderDSBuildingData* existingData = mDSBuildingData.Find(DataAsset->mBuildingClass);
			if (existingData == nullptr) {
				FRRDADataHolderDSBuildingData newData = FRRDADataHolderDSBuildingData();
				newData.ModAdapter = DataAsset;

				mDSBuildingData.Add(DataAsset->mBuildingClass, newData);
			}
			else {
				existingData->ModAdapter = DataAsset;
			}
		}
	}
	UE_LOGFMT(LogRRDApi, Log, "ReCacheDSModAdapterDataAssets: {0}", mDSBuildingData.Num());
}

void URRDADataAssetSubsystem::ReCacheDSAdapterCoversDataAssets() {
	TSet<URRDADataAssetDSAdapterCovers*> DataAssets;
	if (FindAllDataAssetsOfClass(DataAssets)) {
		for (URRDADataAssetDSAdapterCovers* DataAsset : DataAssets) {
			FRRDADataHolderDSBuildingData* existingData = mDSBuildingData.Find(DataAsset->mBuildingClass);
			if (existingData == nullptr) {
				FRRDADataHolderDSBuildingData newData = FRRDADataHolderDSBuildingData();
				newData.AdapterCovers = DataAsset;
				mDSBuildingData.Add(DataAsset->mBuildingClass, newData);
			}
			else {
				existingData->AdapterCovers = DataAsset;
			}
		}
	}
	UE_LOGFMT(LogRRDApi, Log, "ReCacheDSAdapterCoversDataAssets: {0}", mDSBuildingData.Num());
}

int32 URRDADataAssetSubsystem::GetAllDirtItems(TArray<TSubclassOf<UFGItemDescriptor>>& Items) const {
	if (!mDirtAssets.Num()) {
		UE_LOG(LogRRDApi, Error, TEXT("No Dirts found!"));
	}
	return mDirtAssets.GetKeys(Items);
}

URRDADirtDataAsset* URRDADataAssetSubsystem::GetDataForDirtItem(TSubclassOf<UFGItemDescriptor> Item) const {
	if (!Item) return nullptr;
	URRDADirtDataAsset* const* Desc = mDirtAssets.Find(Item);
	if (!Desc) return nullptr;
	return *Desc;
}

TArray<URRDADirtDataAsset*> URRDADataAssetSubsystem::GetAllDirtAssets() const {
	TArray<URRDADirtDataAsset*> Assets;
	mDirtAssets.GenerateValueArray(Assets);
	return Assets;
}

int32 URRDADataAssetSubsystem::GetAllDSBuildingClasses(TArray<TSubclassOf<AFGBuildableFactory>>& OutBuildingClasses) const {
	if (!mDSBuildingData.Num()) {
		UE_LOG(LogRRDApi, Error, TEXT("No DS Buildings found!"));
	}
	return mDSBuildingData.GetKeys(OutBuildingClasses);
}

URRDADataAssetDSModAdapter* URRDADataAssetSubsystem::GetModAdapterDataAssetForBuildingClass(TSubclassOf<AFGBuildableFactory> BuildingClass) const {
	if (!BuildingClass) return nullptr;

	const FRRDADataHolderDSBuildingData* data = mDSBuildingData.Find(BuildingClass);
	if (data == nullptr) return nullptr;

	URRDADataAssetDSModAdapter* DataAsset = data->ModAdapter;
	if (!DataAsset) return nullptr;

	return DataAsset;
}

TArray<URRDADataAssetDSModAdapter*> URRDADataAssetSubsystem::GetAllModAdapterDataAssets() const {
	TArray<URRDADataAssetDSModAdapter*> Assets;
	TArray<FRRDADataHolderDSBuildingData> buildingDatas;

	mDSBuildingData.GenerateValueArray(buildingDatas);

	for (FRRDADataHolderDSBuildingData& buildingData : buildingDatas) {
		Assets.AddUnique(buildingData.ModAdapter);
	}
	return Assets;
}

bool URRDADataAssetSubsystem::HasDSModAdapterForBuildingClass(TSubclassOf<AFGBuildableFactory> BuildingClass) {
	return IsValid(GetModAdapterDataAssetForBuildingClass(BuildingClass));
}

URRDADataAssetDSAdapterCovers* URRDADataAssetSubsystem::GetDSAdapterCoverDataAssetForBuildingClass(TSubclassOf<AFGBuildableFactory> BuildingClass) const {
	if (!BuildingClass) return nullptr;

	const FRRDADataHolderDSBuildingData* data = mDSBuildingData.Find(BuildingClass);
	if (data == nullptr) return nullptr;

	URRDADataAssetDSAdapterCovers* DataAsset = data->AdapterCovers;
	if (!DataAsset) return nullptr;

	return DataAsset;
}

TArray<URRDADataAssetDSAdapterCovers*> URRDADataAssetSubsystem::GetAllDSAdapterCoverDataAssets() const {
	TArray<URRDADataAssetDSAdapterCovers*> Assets;
	TArray<FRRDADataHolderDSBuildingData> buildingDatas;

	mDSBuildingData.GenerateValueArray(buildingDatas);

	for (FRRDADataHolderDSBuildingData& buildingData : buildingDatas) {
		Assets.AddUnique(buildingData.AdapterCovers);
	}
	return Assets;
}

bool URRDADataAssetSubsystem::HasDSAdaperCoversForBuildingClass(TSubclassOf<AFGBuildableFactory> BuildingClass) {
	return IsValid(GetDSAdapterCoverDataAssetForBuildingClass(BuildingClass));
}

bool URRDADataAssetSubsystem::GetAllTurbineItems(TArray<TSubclassOf<UFGItemDescriptor>>& Items, int32 Tier) const {
	if (!mTurbineAssets.Contains(Tier)) {
		return false;
	}

	Items.Empty();
	mTurbineAssets[Tier].Data.GetKeys(Items);

	return true;
}

bool URRDADataAssetSubsystem::GetAllTurbineRelevantItems(TArray<TSubclassOf<UFGItemDescriptor>>& Items, int32 Tier) const {
	Items.Empty();
	for (URRDATurbineDataAsset* Asset : GetAllTurbineAssets(Tier)) {
		Items.AddUnique(Asset->mItem.ItemClass);
		if (Asset->mWasteItem.ItemClass) {
			Items.AddUnique(Asset->mWasteItem.ItemClass);
		}
	}
	return Items.Num() > 0;
}

TArray<URRDATurbineDataAsset*> URRDADataAssetSubsystem::GetAllTurbineAssets(int32 Tier) const {
	if (!mTurbineAssets.Contains(Tier)) {
		UE_LOG(LogRRDApi, Error, TEXT("No TurbineAssets found for Type: %d"), Tier);
		return TArray<URRDATurbineDataAsset*>();
	}

	TArray<URRDATurbineDataAsset*> Assets;
	mTurbineAssets[Tier].Data.GenerateValueArray(Assets);
	return Assets;
}

URRDATurbineDataAsset* URRDADataAssetSubsystem::GetDefaultTurbineAsset(int32 Tier) const {
	return GetAllTurbineAssets(Tier).Top();
}

URRDATurbineDataAsset* URRDADataAssetSubsystem::GetTurbineItemData(TSubclassOf<UFGItemDescriptor> Item, int32 Tier) const {
	if (!mTurbineAssets.Contains(Tier)) {
		return nullptr;
	}

	return mTurbineAssets[Tier].Data.FindRef(Item);
}


bool URRDADataAssetSubsystem::GetAllCoolerItems(TArray<TSubclassOf<UFGItemDescriptor>>& Items, ERRDACoolerType Type) const {
	if (!mCoolerAssets.Contains(Type)) return false;

	Items.Empty();
	mCoolerAssets[Type].Data.GetKeys(Items);

	return Items.Num() > 0;
}

URRDACoolerDataAsset* URRDADataAssetSubsystem::GetCoolerItemData(TSubclassOf<UFGItemDescriptor> Item, ERRDACoolerType Type) const {
	if (!mCoolerAssets.Contains(Type)) return nullptr;
	return mCoolerAssets[Type].Data.FindRef(Item);
}

URRDACoolerDataAsset* URRDADataAssetSubsystem::GetDefaultCoolerAsset(ERRDACoolerType Type) const {
	return GetAllCoolerAssets(Type).Top();
}

TArray<URRDACoolerDataAsset*> URRDADataAssetSubsystem::GetAllCoolerAssets(ERRDACoolerType Type) const {
	if (!mCoolerAssets.Contains(Type)) {
		UE_LOG(LogRRDApi, Error, TEXT("No CoolerAssets found for Type: %d"), Type);
		return TArray<URRDACoolerDataAsset*>();
	}

	TArray<URRDACoolerDataAsset*> Assets;
	mCoolerAssets[Type].Data.GenerateValueArray(Assets);

	return Assets;
}

URRDAGeneratorDataAsset* URRDADataAssetSubsystem::GetGeneratorItemData(int32 Tier) const {
	UE_LOG(LogRRDApi, Log, TEXT("GetGeneratorItemData: %d"), Tier);
	if (!mGeneratorAssets.Contains(Tier)) return nullptr;
	return mGeneratorAssets.FindRef(Tier);
}

bool URRDADataAssetSubsystem::GetAllBoilerItems(TArray<TSubclassOf<UFGItemDescriptor>>& Items, int32 Tier) const {
	if (!mBoilerAssets.Contains(Tier)) return false;

	mBoilerAssets[Tier].Data.GetKeys(Items);

	return Items.Num() > 0;
}

URRDABoilerDataAsset* URRDADataAssetSubsystem::GetBoilerItemData(TSubclassOf<UFGItemDescriptor> Item, int32 Tier) const {
	if (!mBoilerAssets.Contains(Tier)) return nullptr;
	return mBoilerAssets[Tier].Data.FindRef(Item);
}

bool URRDADataAssetSubsystem::GetAllBoilerRelevantItems(TArray<TSubclassOf<UFGItemDescriptor>>& Items, int32 Tier) const {
	Items.Empty();
	for (URRDABoilerDataAsset* Asset : GetAllBoilerAssets(Tier)) {
		Items.AddUnique(Asset->mInput.ItemClass);
		if (Asset->mOutput.ItemClass) {
			Items.AddUnique(Asset->mOutput.ItemClass);
		}
	}
	return Items.Num() > 0;
}

URRDABoilerDataAsset* URRDADataAssetSubsystem::GetDefaultBoilerAsset(int32 Tier) const {
	return GetAllBoilerAssets(Tier).Top();
}

TArray<URRDABoilerDataAsset*> URRDADataAssetSubsystem::GetAllBoilerAssets(int32 Tier) const {
	if (!mBoilerAssets.Contains(Tier)) {
		UE_LOG(LogRRDApi, Error, TEXT("No BoilerAssets found for Type: %d"), Tier);
		return TArray<URRDABoilerDataAsset*>();
	}

	TArray<URRDABoilerDataAsset*> Assets;
	mBoilerAssets[Tier].Data.GenerateValueArray(Assets);

	return Assets;
}

bool URRDADataAssetSubsystem::GetAllHeaterItems(TArray<TSubclassOf<UFGItemDescriptor>>& Items, ERRDAHeaterType Type) const {
	if (!mHeaterAssets.Contains(Type)) return false;

	mHeaterAssets[Type].Data.GetKeys(Items);

	return Items.Num() > 0;
}

URRDAHeaterDataAsset* URRDADataAssetSubsystem::GetHeaterItemData(TSubclassOf<UFGItemDescriptor> Item, ERRDAHeaterType Type) const {
	if (!mHeaterAssets.Contains(Type)) return nullptr;

	return mHeaterAssets[Type].Data.FindRef(Item);
}

bool URRDADataAssetSubsystem::GetAllHeaterRelevantItems(TArray<TSubclassOf<UFGItemDescriptor>>& Items, ERRDAHeaterType Type) const {
	Items.Empty();
	for (URRDAHeaterDataAsset* HeaterAsset : GetAllHeaterAssets(Type)) {
		Items.AddUnique(HeaterAsset->mInput.ItemClass);
		if (HeaterAsset->mOutput.ItemClass) {
			Items.AddUnique(HeaterAsset->mOutput.ItemClass);
		}
	}
	return Items.Num() > 0;
}

URRDAHeaterDataAsset* URRDADataAssetSubsystem::GetDefaultHeaterAsset(ERRDAHeaterType Type) const {
	return GetAllHeaterAssets(Type).Top();
}

TArray<URRDAHeaterDataAsset*> URRDADataAssetSubsystem::GetAllHeaterAssets(ERRDAHeaterType Type) const {
	if (!mHeaterAssets.Contains(Type)) {
		UE_LOG(LogRRDApi, Error, TEXT("No HeaterAssets found for Type: %d"), Type);
		return TArray<URRDAHeaterDataAsset*>();
	}

	TArray<URRDAHeaterDataAsset*> Assets;
	mHeaterAssets[Type].Data.GenerateValueArray(Assets);

	return Assets;
}