artifact-tracker/Source/DLL/src/Functions/ArtifactTracker.h

#pragma once

namespace Papyrus::ArtifactTracker
{
	inline bool is_artifact(RE::TESForm* a_form, RE::TESForm* a_excludeForm = NULL)
	{
		const auto formType = a_form->GetFormType();

		if (formType != RE::FormType::Armor && formType != RE::FormType::Weapon && formType != RE::FormType::Book && formType != RE::FormType::Misc) {
			return false;
		}

		if (Papyrus::ObjectReference::inv_util::is_excluded(a_form, a_excludeForm)) {
			return false;
		}

		return true;
	}

	inline RE::TESObjectREFR* GetCellStorage(RE::StaticFunctionTag*,
		RE::TESObjectREFR* a_ref,
		RE::BGSListForm* a_refList,
		RE::TESBoundObject* a_objectToCreate,
		bool a_autoCreate = true)
	{
		RE::TESObjectREFR* result = NULL;

		if (!a_ref || !a_refList || !a_objectToCreate) {
			return result;
		}

		RE::TESObjectCELL* cell = a_ref->GetParentCell();

		a_refList->ForEachForm([&result, &cell, &a_objectToCreate](RE::TESForm& a_exform) {
			const auto ref = a_exform.As<RE::TESObjectREFR>();

			if (ref && ref->GetParentCell() == cell && ref->GetBaseObject()->formID == a_objectToCreate->formID) {
				result = ref;
				return false;
			}

			return true;
		});

		if (!result && a_autoCreate) {
			result = a_ref->PlaceObjectAtMe(a_objectToCreate, true).get();
			result->Disable();
			a_refList->AddForm(result);
		}

		return result;
	}

	inline bool HasRefInCell(VM* a_vm, StackID a_stackID, RE::StaticFunctionTag*,
		RE::TESBoundObject* a_item,
		RE::TESObjectCELL* a_cell = NULL,
		bool a_checkInContainers = true,
		RE::BGSListForm* a_excludeList = NULL)
	{
		bool bResult = false;

		if (!a_cell) {
			a_cell = RE::PlayerCharacter::GetSingleton()->GetParentCell();
		}

		a_cell->ForEachReference([&](RE::TESObjectREFR& a_ref) {
			const auto baseObj = a_ref.GetBaseObject();

			if (a_item->formID == baseObj->formID) {
				if (a_ref.IsDisabled() || a_ref.IsMarkedForDeletion()) {
					return true;
				}
				if (a_excludeList && a_excludeList->HasForm(a_ref.formID)) {
					return true;
				}
				bResult = true;
				return false;
			} else if (a_checkInContainers) {
				if (baseObj->formType == RE::FormType::Container || (baseObj->formType == RE::FormType::NPC && !a_ref.IsDisabled() && baseObj->As<RE::TESNPC>()->GetRace()->formID == 0x0010760A)) {
					if (a_excludeList && a_excludeList->HasForm(a_ref.formID)) {
						return true;
					}
					const auto inv = a_ref.GetInventory([&](RE::TESBoundObject& a_object) -> bool {
						return a_item->formID == a_object.formID;
					});
					const auto it = inv.find(a_item);
					if (it != inv.end() ? it->second.first : 0) {
						bResult = true;
						return false;
					}
				}
			}

			return true;
		});

		return bResult;
	}

	inline void SyncCellStorage(VM* a_vm, StackID a_stackID, RE::StaticFunctionTag*,
		RE::TESObjectREFR* a_cellStorage,
		RE::TESForm* a_excludeForm = NULL)
	{
		if (!a_cellStorage) {
			return;
		}

		std::unordered_map<RE::FormID, bool> cellItems;

		const auto storageBase = a_cellStorage->GetBaseObject();
		const auto cell = a_cellStorage->GetParentCell();
		const auto inv = a_cellStorage->GetInventory();

		for (const auto a_ref : cell->references) {
			const auto baseObj = a_ref->GetBaseObject();

			if (baseObj->formType == RE::FormType::Container || (baseObj->formType == RE::FormType::NPC && !a_ref->IsDisabled() && baseObj->As<RE::TESNPC>()->GetRace()->formID == 0x0010760A)) {
				if (a_ref->GetBaseObject()->formID == storageBase->formID) {
					continue;
				}

				const auto contInv = a_ref->GetInventory();

				for (const auto& [item, data] : contInv) {
					if (cellItems.contains(item->formID)) {
						continue;
					}
					const auto& [count, entry] = data;
					if (count > 0) {
						cellItems[item->formID] = true;
						if (inv.find(item) == inv.end()) {
							if (is_artifact(item, a_excludeForm)) {
								a_cellStorage->AddObjectToContainer(item, nullptr, 1, nullptr);
							}
						}
					}
				}

				continue;
			}
			
			if (a_ref->IsDisabled() || a_ref->IsMarkedForDeletion()) {
				continue;
			}

			if (cellItems.contains(baseObj->formID)) {
				continue;
			}

			if (!is_artifact(baseObj, a_excludeForm)) {
				continue;
			}
			
			cellItems[baseObj->formID] = true;

			if (inv.find(baseObj) == inv.end()) {
				a_cellStorage->AddObjectToContainer(baseObj, nullptr, 1, nullptr);
			}
		}

		for (const auto& [item, data] : inv) {
			const auto& [count, entry] = data;
			if (count > 0 && !cellItems.contains(item->formID)) {
				a_cellStorage->RemoveItem(item, count, RE::ITEM_REMOVE_REASON::kRemove, nullptr, nullptr);
			}
		}

		cellItems.clear();
	}

	inline void Bind(VM& a_vm)
	{
		BIND(GetCellStorage);
		logger::info("Registered GetCellStorage"sv);
		BIND(HasRefInCell);
		logger::info("Registered HasRefInCell"sv);
		BIND(SyncCellStorage);
		logger::info("Registered SyncCellStorage"sv);
	}
}