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

namespace ArtifactTracker
{
	bool g_bLoaded;
	bool g_bHomeContainer;
	RE::TESBoundObject* g_cellContainer;
	RE::BGSListForm* g_listNew;
	RE::BGSListForm* g_listStored;
	RE::BGSListForm* g_listFound;
	RE::BGSListForm* g_persistentStorage;
	RE::BGSKeyword* g_homeKeyword;
	std::unordered_map<RE::FormID, RE::TESForm*> g_artifactMap;
	std::unordered_map<RE::FormID, RE::TESObjectREFR*> g_persistentMap;
	RE::TESObjectREFR* g_cellStorage;

	void Init()
	{
		g_bLoaded = false;

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

		if (!dataHandler) {
			SKSE::log::error("Failed to call RE::TESDataHandler::GetSingleton()");
			return;
		}

		g_cellContainer = dataHandler->LookupForm(0x800, "Artifact Tracker.esp")->As<RE::TESBoundObject>(); // ETR_CellStorageContainer

		g_listNew = dataHandler->LookupForm<RE::BGSListForm>(0x803, "Artifact Tracker.esp"); // ETR_ItemsNew
		g_listStored = dataHandler->LookupForm<RE::BGSListForm>(0x805, "Artifact Tracker.esp"); // ETR_ItemsStored
		g_listFound = dataHandler->LookupForm<RE::BGSListForm>(0x806, "Artifact Tracker.esp"); // ETR_ItemsFound
		g_persistentStorage = dataHandler->LookupForm<RE::BGSListForm>(0x807, "Artifact Tracker.esp"); // ETR_PersistentStorageList

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

		const auto recipeKeyword = dataHandler->LookupForm<RE::BGSKeyword>(0xF5CB0, "Skyrim.esm"); // VendorItemRecipe
		const auto excludeKeywords = dataHandler->LookupForm<RE::BGSListForm>(0x801, "Artifact Tracker.esp"); // ETR_ExcludeMiscKeywords

		if (!g_cellContainer || !g_listNew || !g_listStored || !g_listFound || !g_persistentStorage || !g_homeKeyword || !recipeKeyword || !excludeKeywords) {
			SKSE::log::warn("Failed to load data from Artifact Tracker.esp");
			RE::DebugMessageBox("Failed to load data from Artifact Tracker.esp, the mod is disabled.");
			return;
		}

		// Preloading item lists

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

		for (const auto& form : dataHandler->GetFormArray(RE::FormType::Misc)) {
			if (form->GetPlayable() && !form->HasKeywordInList(excludeKeywords, false)) {
				g_artifactMap[form->formID] = form;
			}
		}
		g_artifactMap.erase(0xA); // Lockpick
		g_artifactMap.erase(0xF); // Gold

		for (const auto& form : dataHandler->GetFormArray(RE::FormType::Weapon)) {
			if (form->GetPlayable()) {
				g_artifactMap[form->formID] = form;
			}
		}
		g_artifactMap.erase(0x1F4); // Unarmed

		for (const auto& form : dataHandler->GetFormArray(RE::FormType::Armor)) {
			if (form->GetPlayable()) {
				g_artifactMap[form->formID] = form;
			}
		}
		
		OnGameLoad();
		EventListener::Install();

		g_bLoaded = true;
	}

	bool IsArtifact(RE::TESForm* a_form)
	{
		if (!a_form) {
			return false;
		}

		const auto formType = a_form->GetFormType();

		if (formType == RE::FormType::Armor || formType == RE::FormType::Weapon || formType == RE::FormType::Book || formType == RE::FormType::Misc) {
			return g_artifactMap.contains(a_form->formID);
		}

		return false;
	}

	RE::TESForm* GetArtifactByID(RE::FormID a_formID)
	{
		if (!a_formID) {
			return nullptr;
		}
		
		const auto it = g_artifactMap.find(a_formID);
		return it != g_artifactMap.end() ? it->second : nullptr;
	}

	void OnGameLoad()
	{
		#ifdef _DEBUG
		SKSE::log::info("OnGameLoad");
		#endif

		g_persistentMap.clear();
		g_persistentStorage->ForEachForm([&](RE::TESForm& a_exform) {
			if (&a_exform) {
				g_persistentMap[a_exform.formID] = a_exform.As<RE::TESObjectREFR>();
			}
			return true;
		});
	}
	
	void SetContainerMode(bool bOpening)
	{
		if (bOpening) {

			const auto refr = RE::TESObjectREFR::LookupByHandle(RE::ContainerMenu::GetTargetRefHandle());

			g_bHomeContainer = IsHome()
				&& refr
				&& refr.get()->GetParentCell() == RE::PlayerCharacter::GetSingleton()->GetParentCell()
				&& !g_persistentMap.contains(refr.get()->formID);

			#ifdef _DEBUG
			if (g_bHomeContainer) {
				RE::DebugNotification("Delayed processing enabled");
			}
			#endif

		} else if (g_bHomeContainer) {
			g_bHomeContainer = false;
			SyncCellStorage();
		}
	}

	bool IsHome()
	{
		return (bool)g_cellStorage;
	}

	bool ToggleHomeMode(RE::TESObjectREFR* cellStorage)
	{
		if (cellStorage) {
			g_bHomeContainer = false;
			g_cellStorage = cellStorage;
			RE::UI::GetSingleton()->AddEventSink<RE::MenuOpenCloseEvent>(EventListener::GetSingleton());
			#ifdef _DEBUG
			SKSE::log::info("Home mode ON");
			#endif
			return true;
		} else if (g_cellStorage) {
			g_bHomeContainer = false;
			g_cellStorage = nullptr;
			RE::UI::GetSingleton()->RemoveEventSink<RE::MenuOpenCloseEvent>(EventListener::GetSingleton());
			#ifdef _DEBUG
			SKSE::log::info("Home mode OFF");
			#endif
		}
		return false;
	}

	bool IsValidContainer(RE::TESObjectREFR* a_ref)
	{
		if (!a_ref) {
			return false;
		}

		const auto baseObj = a_ref->GetBaseObject();

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

	void OnCellEnter(RE::FormID a_formID)
	{
		RE::TESObjectCELL* cell = RE::TESForm::LookupByID<RE::TESObjectCELL>(a_formID);
		RE::BGSLocation* location = cell ? cell->GetLocation() : nullptr;

		if (!cell || !location || !cell->IsInteriorCell() || !location->HasKeyword(g_homeKeyword)) {
			if (IsHome()) {
				RE::ScriptEventSourceHolder::GetSingleton()->RemoveEventSink<RE::TESCellFullyLoadedEvent>(EventListener::GetSingleton());
				ToggleHomeMode(nullptr);
			}
			return;
		}

		RE::TESObjectREFR* cellStorage = nullptr;

		g_persistentStorage->ForEachForm([&](RE::TESForm& a_form) {
			const auto refr = a_form.As<RE::TESObjectREFR>();
			if (refr && refr->GetParentCell()->formID == a_formID) {
				cellStorage = refr;
				return false;
			}
			return true;
		});

		ToggleHomeMode(cellStorage);

		if (!cellStorage) {
			RE::ScriptEventSourceHolder::GetSingleton()->AddEventSink<RE::TESCellFullyLoadedEvent>(EventListener::GetSingleton());
		}
	}

	void OnCellEnter(RE::BGSLocation* location, RE::TESObjectCELL* cell)
	{
		if (!location || !cell->IsInteriorCell() || cell != RE::PlayerCharacter::GetSingleton()->GetParentCell() || !location->HasKeyword(g_homeKeyword)) {
			ToggleHomeMode(nullptr);
			return;
		}

		RE::TESObjectREFR* cellStorage = nullptr;

		for (const auto& a_ref : cell->references) {
			if (a_ref.get()->GetBaseObject()->formID == g_cellContainer->formID) {
				cellStorage = a_ref.get();
				break;
			}
		}

		if (cellStorage) {
			if (!g_persistentStorage->HasForm(cellStorage)) {
				g_persistentStorage->AddForm(cellStorage);
				g_persistentMap[cellStorage->formID] = cellStorage;
			}
			ToggleHomeMode(cellStorage);
			return;
		}

		#ifdef _DEBUG
		SKSE::log::info("Adding new storage in {}", cell->GetName());
		#endif

		cellStorage = RE::PlayerCharacter::GetSingleton()->PlaceObjectAtMe(g_cellContainer, true).get();

		if (cellStorage) {
			cellStorage->Disable();
			g_persistentStorage->AddForm(cellStorage);
			g_persistentMap[cellStorage->formID] = cellStorage;
			ToggleHomeMode(cellStorage);
		} else {
			SKSE::log::error("Failed to create cell storage in OnCellEnter");
			ToggleHomeMode(nullptr);
		}
	}

	void SyncCellStorage()
	{
		if (!IsHome()) {
			#ifdef _DEBUG
			SKSE::log::info("SyncCellStorage called while not at home");
			#endif
			return;
		}

		#ifdef _DEBUG
		SKSE::log::info("Running SyncCellStorage");
		#endif

		std::unordered_set<RE::FormID> cellItems;

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

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

			if (IsValidContainer(a_ref.get())) {
				if (g_cellContainer->formID == baseObj->formID || baseObj->formID == 0xDC9E7 || g_persistentMap.contains(a_ref->formID)) {  // 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.insert(item->formID);
						if (IsArtifact(item)) {
							if (inv.find(item) == inv.end()) {
								g_cellStorage->AddObjectToContainer(item, nullptr, 1, nullptr);
							}
							if (!g_listStored->HasForm(item)) {
								ListRemoveItem(g_listNew, item);
								ListRemoveItem(g_listFound, item);
								g_listStored->AddForm(item);
							}
						}
					}
				}

				continue;
			}

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

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

			cellItems.insert(baseObj->formID);

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

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

			if (!g_listStored->HasForm(baseObj)) {
				ListRemoveItem(g_listNew, baseObj);
				ListRemoveItem(g_listFound, baseObj);
				g_listStored->AddForm(baseObj);
			}
		}

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

		cellItems.clear();
	}

	void OnContainerChanged(const RE::TESContainerChangedEvent* a_event, RE::TESForm* form)
	{
		if (a_event->newContainer == 0x14) {
			
			if (a_event->oldContainer) {

				if (g_persistentMap.contains(a_event->oldContainer)) {

					// Items in persistent containers are marked as stored by definition, no need to check the list
					if (!RefListHasItem(g_persistentStorage, a_event->baseObj)) {
						ListRemoveItem(g_listStored, form);
						g_listFound->AddForm(form);
					}

					return;

				} else if (g_cellStorage) { // non-persistent container at home
				
					// g_bContainerMode is expected to be true, enabling processing after closing container.
					// Reachable by Loot Menu or a mod, retrieving items via a spell of after scanning containers in cell.
					// In extreme cases of mass retrieval, this can result in a few frames lag.

					#ifdef _DEBUG
					SKSE::log::info("Synchronous processing of a non-persistent container (moved to player)");
					#endif

					SyncCellStorage();

					return;
				}

			} else if (g_cellStorage) {
				// Items picked up at home are handled in the perk
				return;
			}

			// Instead of looking up in huge g_listNew, we check smaller g_listStored and g_listFound.
			// Mass retrieval of found items is rare, so we check the stored list first.
			if (g_listStored->HasForm(form) || g_listFound->HasForm(form)) {
				return;
			}

			// It's a new item, move it to found
			ListRemoveItem(g_listNew, form);
			g_listFound->AddForm(form);

			return;
		}

		if (a_event->oldContainer != 0x14) {
			return;
		}
		
		// Items moved from player's inventory

		if (!a_event->newContainer) { // no destination container

			if (g_cellStorage && a_event->reference) {  // dropped or placed on rack at home
				if (GetItemCount(g_cellStorage, form) <= 0) {
					#ifdef _DEBUG
					SKSE::log::info("Added dropped {} to cell storage", form->GetName());
					RE::DebugNotification("Adding to cell storage");
					#endif
					g_cellStorage->AddObjectToContainer(form->As<RE::TESBoundObject>(), nullptr, 1, nullptr);
				}
				ListRemoveItem(g_listFound, form);
				ListRemoveItem(g_listNew, form);
				g_listStored->AddForm(form);
				return;
			}

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

			// NB: During OnContainerChanged, InventoryChanges do not have the current change included yet
			if ((GetItemCount(RE::PlayerCharacter::GetSingleton(), form->formID) - a_event->itemCount <= 0) && !FollowersHaveItem(form)) {
				ListRemoveItem(g_listFound, form);
				ListRemoveItem(g_listNew, form);
				g_listNew->AddForm(form);
			}

			return;
		}

		if (g_persistentMap.contains(a_event->newContainer)) { // moved to a persistent container

			if (!g_listStored->HasForm(form)) {
				ListRemoveItem(g_listFound, form);
				g_listStored->AddForm(form);
			}

		} else if (g_cellStorage) { // stored at home in a non-persistent/non-registered container
		
			// g_bContainerMode is expected to be true, enabling processing after closing container.
			// Can be hit by autosorting mods. Most of them work with persistent containers, which should be added to the list of persistent containers.

			#ifdef _DEBUG
			SKSE::log::info("Synchronous processing of a non-persistent container (moved from player)");
			#endif

			const auto targetContainer = RE::TESForm::LookupByID<RE::TESObjectREFR>(a_event->newContainer);

			if (IsValidContainer(targetContainer)) {
				if (GetItemCount(g_cellStorage, form) <= 0) {
					g_cellStorage->AddObjectToContainer(form->As<RE::TESBoundObject>(), nullptr, 1, nullptr);
				}
			
				if (!g_listStored->HasForm(form)) {
					ListRemoveItem(g_listFound, form);
					g_listStored->AddForm(form);
				}
			}
			
		// NB: During OnContainerChanged, InventoryChanges do not have the current change included yet
		} else if (g_listFound->HasForm(form) && (GetItemCount(RE::PlayerCharacter::GetSingleton(), form->formID) - a_event->itemCount <= 0) && !FollowersHaveItem(form)) {
			ListRemoveItem(g_listFound, form);
			g_listNew->AddForm(form);
		}
	}

	void AddRefArtifactsToList(RE::TESForm* a_refOrList, RE::BGSListForm* a_targetList, RE::BGSListForm* a_excludeList)
	{
		if (!a_refOrList || !a_targetList) {
			SKSE::log::warn("Invalid arguments in AddRefArtifactsToList");
			return;
		}

		if (a_refOrList->Is(RE::FormType::FormList)) {
			a_refOrList->As<RE::BGSListForm>()->ForEachForm([&](RE::TESForm& a_exform) {
				const auto refrItem = a_exform.As<RE::TESObjectREFR>();
				if (refrItem) {
					AddRefArtifactsToList(refrItem, a_targetList, a_excludeList);
				}
				return true;
			});
			return;
		}

		const auto containerRef = a_refOrList->As<RE::TESObjectREFR>();

		if (!containerRef) {
			SKSE::log::warn("containerRef in AddRefArtifactsToList is not a reference");
			return;
		}

		const auto inv = containerRef->GetInventory([&](RE::TESBoundObject& a_exform) {
			return ArtifactTracker::IsArtifact(&a_exform) && (!a_excludeList || !a_excludeList->HasForm(&a_exform));
		});

		for (const auto& item : inv) {
			if (item.second.first > 0) {
				a_targetList->AddForm(item.first);
			}
		}
	}
}