#include "ArtifactTracker.h"
#include "BookCheck.h"
#include "EventListener.h"
#include "Util.h"

namespace ArtifactTracker
{
	RE::TESGlobal* notifyNewArtifact;
	RE::TESBoundObject* cellContainer;
	RE::BGSListForm* listNew;
	RE::BGSListForm* listStored;
	RE::BGSListForm* listFound;
	RE::BGSListForm* persistentStorage;
	RE::BGSKeyword* homeKeyword;
	bool bAtHome;
	std::unordered_map<RE::FormID, RE::TESObjectMISC*> validMisc;
	std::unordered_map<RE::FormID, RE::TESObjectBOOK*> validBooks;

	void Init()
	{
		EventListener::Install();

		const auto dataHandler = RE::TESDataHandler::GetSingleton();

		notifyNewArtifact = dataHandler->LookupForm<RE::TESGlobal>(0x809, "Artifact Tracker.esp");
		cellContainer = dataHandler->LookupForm(0x800, "Artifact Tracker.esp")->As<RE::TESBoundObject>();

		listNew = dataHandler->LookupForm<RE::BGSListForm>(0x803, "Artifact Tracker.esp");
		listStored = dataHandler->LookupForm<RE::BGSListForm>(0x805, "Artifact Tracker.esp");
		listFound = dataHandler->LookupForm<RE::BGSListForm>(0x806, "Artifact Tracker.esp");
		persistentStorage = dataHandler->LookupForm<RE::BGSListForm>(0x807, "Artifact Tracker.esp");

		homeKeyword = dataHandler->LookupForm<RE::BGSKeyword>(0xFC1A3, "Skyrim.esm");

		if (!cellContainer) {
			SKSE::log::warn("cellContainer is empty");
		}

		if (!persistentStorage) {
			SKSE::log::warn("persistentStorage is empty");
		}

		// Preloading item lists
		const RE::BGSKeyword* recipeKeyword = dataHandler->LookupForm<RE::BGSKeyword>(0xF5CB0, "Skyrim.esm");  // VendorItemRecipe

		for (const auto& form : dataHandler->GetFormArray<RE::TESObjectBOOK>()) {
			if (form && !form->TeachesSpell() && (!form->HasKeyword(recipeKeyword) || BookCheck::IsBook(form))) {
				validBooks[form->formID] = form;
			}
		}

		RE::BGSListForm* excludeKeywords = dataHandler->LookupForm<RE::BGSListForm>(0x801, "Artifact Tracker.esp");  // ETR_ExcludeMiscKeywords

		for (const auto& form : dataHandler->GetFormArray<RE::TESObjectMISC>()) {
			if (form->GetPlayable() && !form->IsGold() && !form->IsLockpick() && !form->HasKeywordInList(excludeKeywords, false)) {
				validMisc[form->formID] = form;
			}
		}
	}

	bool IsArtifact(RE::TESForm* a_form)
	{
		switch (a_form->GetFormType()) {
			case RE::FormType::Armor:
				return a_form->GetPlayable();
			case RE::FormType::Weapon:
				return a_form->GetPlayable() && a_form->formID != 0x000001F4;
			case RE::FormType::Book:
				return validBooks.contains(a_form->formID);
			case RE::FormType::Misc:
				return validMisc.contains(a_form->formID);
			default:
				return false;
		}
	}

	bool SetHomeLocation(RE::BGSLocation* a_location = NULL)
	{
		bAtHome = a_location && a_location->HasKeyword(homeKeyword);
		return bAtHome;
	}

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

		if (!a_ref || !a_refList || !a_objectToCreate) {
			SKSE::log::warn("Invalid arguments in GetCellStorage");
			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);
		}

		if (!result) {
			SKSE::log::warn("Failed to find or create cell storage in GetCellStorage");
		}

		return result;
	}

	void SyncCellStorage(RE::TESObjectREFR* a_cellStorage, RE::BGSListForm* a_excludeContainers)
	{
		if (!a_cellStorage || !a_excludeContainers) {
			SKSE::log::warn("Invalid arguments in SyncCellStorage");
			return;
		}

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

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

		for (const RE::NiPointer<RE::TESObjectREFR> 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_excludeContainers->HasForm(a_ref->formID) || baseObj->formID == 0xDC9E7) {  // skip persistent and PlayerBookShelfContainer
					continue;
				}

				const auto contInv = a_ref->GetInventory([&](RE::TESBoundObject& a_object) -> bool {
					return !cellItems.contains(a_object.formID);
				});

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

				continue;
			}

			if (a_ref->IsMarkedForDeletion()) {
				SKSE::log::info("found marked for deletion");
			}

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

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

			cellItems[baseObj->formID] = true;

			if (!IsArtifact(baseObj)) {
				continue;
			}

			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();
	}

	void OnItemPickup(RE::TESBoundObject* form)
	{
		RE::TESObjectREFR* cellStorage = GetCellStorage(RE::PlayerCharacter::GetSingleton(), persistentStorage, cellContainer, true);

		SyncCellStorage(cellStorage, persistentStorage);

		if (!RefHasItem(cellStorage, form) && !GetItemCountInList(persistentStorage, form)) {
			RemoveListItem(listStored, form);
			listFound->AddForm(form);
		}
	}

	void OnContainerChanged(const RE::TESContainerChangedEvent* a_event)
	{
		if (a_event->newContainer == 0x14) {
			return;
			RE::TESBoundObject* form = RE::TESForm::LookupByID<RE::TESBoundObject>(a_event->baseObj);

			if (!form || !IsArtifact(form)) {
				return;
			}

			if (listFound->HasForm(form) || listStored->HasForm(form)) {
				return;
			}

			if (listNew->HasForm(form)) {

				RemoveListItem(listNew, form);

				listFound->AddForm(form);

				if (notifyNewArtifact->value) {
					const auto itemName = form->GetName();
					char* notificationText = new char[strlen("New artifact acquired: ") + strlen(itemName) + 1];
					strcpy(notificationText, "New artifact acquired: ");
					strcat(notificationText, itemName);
					RE::DebugNotification(notificationText);
					delete[] notificationText;
				}

			}

		} else if (a_event->oldContainer == 0x14) {

			RE::TESBoundObject* form = RE::TESForm::LookupByID<RE::TESBoundObject>(a_event->baseObj);

			SKSE::log::info("{}", form->GetName());
			if (!form || !IsArtifact(form)) {
				return;
			}

			if (!a_event->newContainer) {
				if (bAtHome && a_event->reference) {  // dropped at home
					RE::TESObjectREFR* cellStorage = GetCellStorage(RE::PlayerCharacter::GetSingleton(), persistentStorage, cellContainer, true);
					if (!RefHasItem(cellStorage, form)) {
						cellStorage->AddObjectToContainer(form, nullptr, 1, nullptr);
					}
					if (listFound->HasForm(form)) {
						RemoveListItem(listFound, form);
					}
					listStored->AddForm(form);
					RE::DebugNotification("added to stored");
					return;
				}

				if (listStored->HasForm(form)) {
					return;
				}

				if (!RefHasItem(RE::PlayerCharacter::GetSingleton(), form) && !FollowersHaveItem(form)) {
					if (listFound->HasForm(form)) {
						RemoveListItem(listFound, form);
					}
					listNew->AddForm(form);
					RE::DebugNotification("added to new");
				}

				return;
			}

			bool bPersistent = false;
			for (const auto& ref : persistentStorage->forms) {
				if (ref && ref->formID == a_event->newContainer) {
					bPersistent = true;
					break;
				}
			}

			if (!bPersistent && !bAtHome) {
				if (!RefHasItem(RE::PlayerCharacter::GetSingleton(), form) && !FollowersHaveItem(form->As<RE::TESBoundObject>())) {
					if (listFound->HasForm(form)) {
						RemoveListItem(listFound, form);
					}
					listNew->AddForm(form);
					RE::DebugNotification("added to new");
				}
				return;
			}

			if (listStored->HasForm(form)) {
				return;
			}

			if (!bPersistent && bAtHome) {
				RE::TESObjectREFR* targetContainer = RE::TESForm::LookupByID<RE::TESObjectREFR>(a_event->newContainer);

				if (!targetContainer || targetContainer->GetParentCell() != RE::PlayerCharacter::GetSingleton()->GetParentCell()) {
					return;
				}

				if (targetContainer->GetBaseObject()->Is(RE::FormType::NPC) && targetContainer->GetBaseObject()->As<RE::TESNPC>()->GetRace()->formID != 0x0010760A) {
					return;
				}

				RE::TESObjectREFR* cellStorage = GetCellStorage(RE::PlayerCharacter::GetSingleton(), persistentStorage, cellContainer, true);
				cellStorage->AddObjectToContainer(form->As<RE::TESBoundObject>(), nullptr, 1, nullptr);
				RE::DebugNotification("updated cellStorage");
			}

			if (listFound->HasForm(form)) {
				RemoveListItem(listFound, form);
			}

			listStored->AddForm(form);
			RE::DebugNotification("added to stored");
		}
	}
}