#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/RRDAGeneratorDataAsset.h"
#include "DataAssets/RRDAHeaterDataAsset.h"
#include "DataAssets/RRDATurbineDataAsset.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();
}

void URRDADataAssetSubsystem::StartScanForDataAssets()
{
	UAssetManager& Manager = UAssetManager::Get();
	Manager.LoadPrimaryAssetsWithType(FPrimaryAssetType("RRDADataAsset"));
	
	mDisabledDataAssets.Empty();
	mEnabledDataAssets.Empty();

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

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;
	if(FindAllDataAssetsOfClass(DataAssets))
	{
		for (URRDATurbineDataAsset* DataAsset : DataAssets)
		{
			if(!mTurbineAssets.Contains(DataAsset->mTier)) {
				mTurbineAssets.Add(DataAsset->mTier, TMap<TSubclassOf<UFGItemDescriptor>, URRDATurbineDataAsset*>());	
			}

			if(!DataAsset->mItem.ItemClass) {
				UE_LOG(LogRRDApi, Fatal, TEXT("TurbineDataAsset %s has no ItemClass"), *DataAsset->GetPathName());
				continue;
			}

			mTurbineAssets[DataAsset->mTier].Add(DataAsset->mItem.ItemClass, DataAsset);
		}
	}
	
	UE_LOG(LogRRDApi, Log, TEXT("ReCacheTurbineDataAssets: %d"), mTurbineAssets.Num());
}

void URRDADataAssetSubsystem::ReCacheBoilerDataAssets() {
	mBoilerAssets.Empty();
	TSet<URRDABoilerDataAsset*> DataAssets;
	if(FindAllDataAssetsOfClass(DataAssets))
	{
		for (URRDABoilerDataAsset* DataAsset : DataAssets)
		{
			if(!mBoilerAssets.Contains(DataAsset->mTier)) {
				mBoilerAssets.Add(DataAsset->mTier, TMap<TSubclassOf<UFGItemDescriptor>, URRDABoilerDataAsset*>());	
			}

			if(!DataAsset->mInput.ItemClass) {
				UE_LOG(LogRRDApi, Fatal, TEXT("BoilerDataAsset %s has no ItemClass"), *DataAsset->GetPathName());
				continue;
			}

			mBoilerAssets[DataAsset->mTier].Add(DataAsset->mInput.ItemClass, DataAsset);
		}
	}
	
	UE_LOG(LogRRDApi, Log, TEXT("ReCacheBoilerDataAssets: %d"), mTurbineAssets.Num());
}

void URRDADataAssetSubsystem::ReCacheCoolerDataAssets() {
	mCoolerAssets.Empty();
	TSet<URRDACoolerDataAsset*> DataAssets;
	if(FindAllDataAssetsOfClass(DataAssets))
	{
		for (URRDACoolerDataAsset* DataAsset : DataAssets)
		{
			if(!mCoolerAssets.Contains(DataAsset->mType)) {
				mCoolerAssets.Add(DataAsset->mType, TMap<TSubclassOf<UFGItemDescriptor>, URRDACoolerDataAsset*>());	
			}

			if(!DataAsset->mItem) {
				UE_LOG(LogRRDApi, Fatal, TEXT("CoolerDataAsset %s has no ItemClass"), *DataAsset->GetPathName());
				continue;
			}

			mCoolerAssets[DataAsset->mType].Add(DataAsset->mItem, DataAsset);
		}
	}
	
	UE_LOG(LogRRDApi, Log, TEXT("ReCacheCoolerDataAssets: %d"), mTurbineAssets.Num());
}

void URRDADataAssetSubsystem::ReCacheHeaterDataAssets() {
	mHeaterAssets.Empty();
	TSet<URRDAHeaterDataAsset*> DataAssets;
	if(FindAllDataAssetsOfClass(DataAssets))
	{
		for (URRDAHeaterDataAsset* DataAsset : DataAssets)
		{
			if(!mHeaterAssets.Contains(DataAsset->mType)) {
				mHeaterAssets.Add(DataAsset->mType, TMap<TSubclassOf<UFGItemDescriptor>, URRDAHeaterDataAsset*>());	
			}

			if(!DataAsset->mInput.ItemClass) {
				UE_LOG(LogRRDApi, Fatal, TEXT("HeaterDataAsset %s has no ItemClass"), *DataAsset->GetPathName());
				continue;
			}

			mHeaterAssets[DataAsset->mType].Add(DataAsset->mInput.ItemClass, DataAsset);
		}
	}
	
	UE_LOG(LogRRDApi, Log, TEXT("ReCacheHeaterDataAssets: %d"), mTurbineAssets.Num());
}

void URRDADataAssetSubsystem::ReCacheGeneratorDataAssets()
{
	mHeaterAssets.Empty();
	TSet<URRDAGeneratorDataAsset*> DataAssets;
	if(FindAllDataAssetsOfClass(DataAssets))
	{
		for (URRDAGeneratorDataAsset* DataAsset : DataAssets)
		{
			mGeneratorAssets.Add(DataAsset->mTier, DataAsset);
		}
	}
	
	UE_LOG(LogRRDApi, Log, TEXT("ReCacheHeaterDataAssets: %d"), mTurbineAssets.Num());
}

int32 URRDADataAssetSubsystem::GetAllDirtItems(TArray<TSubclassOf<UFGItemDescriptor>>& Items) const
{
	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;
}

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

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

	return true;
}

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

	Items.Empty();
	for (TTuple<TSubclassOf<UFGItemDescriptor>, URRDATurbineDataAsset*> TurbineAsset : mTurbineAssets[Tier])
	{
		Items.AddUnique(TurbineAsset.Key);
		if(TurbineAsset.Value->mWasteItem.ItemClass) {
			Items.AddUnique(TurbineAsset.Value->mWasteItem.ItemClass);
		}
	}

	return true;
}

TArray<URRDATurbineDataAsset*> URRDADataAssetSubsystem::GetAllTurbineAssets(int32 Tier) const {
	if(!mTurbineAssets.Contains(Tier)) {
		return TArray<URRDATurbineDataAsset*>();
	}

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

URRDATurbineDataAsset* URRDADataAssetSubsystem::GetDefaultTurbineAsset(int32 Tier) const {
	return GetAllTurbineAssets(Tier)[0];
}

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

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


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

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

	return true;
}

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

URRDACoolerDataAsset* URRDADataAssetSubsystem::GetDefaultCoolerAsset(ERRDACoolerType Type) const
{
	return GetAllCoolerAssets(Type)[0];
}

TArray<URRDACoolerDataAsset*> URRDADataAssetSubsystem::GetAllCoolerAssets(ERRDACoolerType Type) const
{
	if(!mCoolerAssets.Contains(Type)) return TArray<URRDACoolerDataAsset*>();

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

	return Assets;
}