#pragma once

#include "Util.h"

#include <cctype>

// Three corrections to the Scaleform scale mode menus are loaded with, all applied from the same
// BSScaleformManager::LoadMovie hook:
//
//   1. Restore the modes used before 1.6.1130. That build switched nearly every menu to
//      kExactFit, which stretches the movie across the whole viewport instead of preserving its
//      aspect ratio. Enderal's menus - including the SkyUI-derived widescreen versions shipped
//      with them - are authored for the old modes, so on anything that is not 16:9 they come out
//      distorted. Only 1.6.1130 and later need this; see g_restoreScaleModes.
//   2. Stop kNoBorder from cropping menus on displays wider than 16:9. That one predates
//      1.6.1130 and is needed on every runtime, so it is not version gated; see ClampToViewport.
//   3. For a table menu whose movie file was replaced by a re-author on a stage wider than
//      16:9, override both of the above with a mode derived from the stage itself; see
//      FitWideStage. Corrections 1 and 2 assume the 1280x720 authoring every vanilla, SkyUI and
//      Enderal movie uses, and they break the wide re-authors at most aspect ratios.
//
// The table below is what 1.6.640 passes to LoadMovie, plus BookMenu and GiftMenu, which 1.6.1130
// left on kNoBorder and which need the clamp. Menus that want no help (Book, Console, CreditsMenu,
// CursorMenu, FaderMenu, StartMenu) are absent and keep whatever the caller asked for. It doubles
// as the list of menus this patch is allowed to touch at all, so on older runtimes - where the
// engine already passes these very modes - a third-party menu loaded with kNoBorder is left alone.
//
// Enderal's own UI.OpenCustomMenu movies are in the table too. SKSE's CustomMenu hardcodes
// kNoBorder for every custom menu on every runtime (skse64/CustomMenu.cpp), and this hook sees
// those loads as well - the movie name from Papyrus is passed to LoadMovie verbatim. Without the
// clamp the wheel menu at 32:9 renders at 2x with the stage cropped top and bottom. For these
// entries both halves of the patch collapse into the clamp: the caller already passes kNoBorder,
// so restoring the table mode is a no-op on every runtime.
//
// The mode-restoration approach follows Widescreen Scale Removed by SkyHorizon
// (https://www.nexusmods.com/skyrimspecialedition/mods/136793). The table is re-derived from the
// 1.6.640/1.6.1170 binaries (see CLAUDE.md for the method) and the implementation is original.
namespace MenuAspectRatioFix
{
	using ScaleModeType = RE::GFxMovieView::ScaleModeType;

	// True on the runtimes that stretch menus to fill the viewport, i.e. 1.6.1130 and later. Below
	// that the engine already asks for the modes in the table, so only the clamp has work to do.
	inline bool g_restoreScaleModes = false;

	// Engine call sites pass fixed literals, but custom menu names arrive through BSFixedString,
	// which interns case-insensitively and hands out whichever caller's casing got interned first -
	// Enderal itself opens the hero menu as both "00E_heromenu" (Papyrus) and "00e_heromenu"
	// (HeroMenuPatch). The lookup therefore must not be case sensitive.
	struct CaseInsensitiveLess
	{
		using is_transparent = void;

		bool operator()(std::string_view a_lhs, std::string_view a_rhs) const
		{
			return std::ranges::lexicographical_compare(a_lhs, a_rhs, [](unsigned char a_l, unsigned char a_r) {
				return std::tolower(a_l) < std::tolower(a_r);
			});
		}
	};

	std::optional<ScaleModeType> GetScaleMode(const char* a_fileName)
	{
		static const std::map<std::string_view, ScaleModeType, CaseInsensitiveLess> modes{
			{ "FavoritesMenu"sv, ScaleModeType::kShowAll },
			{ "HUDMenu"sv, ScaleModeType::kShowAll },
			{ "LevelUpMenu"sv, ScaleModeType::kShowAll },
			{ "LoadWaitSpinner"sv, ScaleModeType::kShowAll },
			{ "LoadingMenu"sv, ScaleModeType::kShowAll },
			{ "Map"sv, ScaleModeType::kShowAll },
			{ "Quest_Journal"sv, ScaleModeType::kShowAll },
			{ "SafeZone"sv, ScaleModeType::kShowAll },
			{ "SleepWaitMenu"sv, ScaleModeType::kShowAll },
			{ "Titles"sv, ScaleModeType::kShowAll },
			{ "TutorialMenu"sv, ScaleModeType::kShowAll },
			{ "TweenMenu"sv, ScaleModeType::kShowAll },

			{ "BarterMenu"sv, ScaleModeType::kNoBorder },
			{ "BookMenu"sv, ScaleModeType::kNoBorder },
			{ "ContainerMenu"sv, ScaleModeType::kNoBorder },
			{ "CraftingMenu"sv, ScaleModeType::kNoBorder },
			{ "DialogueMenu"sv, ScaleModeType::kNoBorder },
			{ "GiftMenu"sv, ScaleModeType::kNoBorder },
			{ "InventoryMenu"sv, ScaleModeType::kNoBorder },
			{ "LockpickingMenu"sv, ScaleModeType::kNoBorder },
			{ "MagicMenu"sv, ScaleModeType::kNoBorder },
			{ "MessageBox"sv, ScaleModeType::kNoBorder },
			{ "RaceSex_menu"sv, ScaleModeType::kNoBorder },

			// The one menu with no direct LoadMovie call site: its ctor queues a task through the
			// deferred-load wrapper, and the task's Run() calls the hooked LoadMovie. kNoBorder up
			// to 1.6.640, kExactFit from 1.6.1130 (checked 1.5.97, 1.6.640, 1.6.117x, 1.7.99, 1.7.104).
			{ "StatsMenu"sv, ScaleModeType::kNoBorder },
			{ "TrainingMenu"sv, ScaleModeType::kNoBorder },

			// Replaces CraftingMenu, https://www.nexusmods.com/skyrimspecialedition/mods/81409
			{ "ConstructibleObjectMenu"sv, ScaleModeType::kNoBorder },

			// Enderal's UI.OpenCustomMenu movies, loaded kNoBorder by SKSE's CustomMenu. The keys
			// are the names Papyrus passes; every shipped movie is listed (uidyemenu.psc references
			// "dyemenu", but no such movie ships).
			{ "00e_heromenu"sv, ScaleModeType::kNoBorder },
			{ "listmenu"sv, ScaleModeType::kNoBorder },
			{ "magicmenuext"sv, ScaleModeType::kNoBorder },
			{ "selectionmenu"sv, ScaleModeType::kNoBorder },
			{ "statssheetmenu"sv, ScaleModeType::kNoBorder },
			{ "textentrymenu"sv, ScaleModeType::kNoBorder },
			{ "uilib/uilib_1_listmenu"sv, ScaleModeType::kNoBorder },
			{ "uilib/uilib_1_textinputmenu"sv, ScaleModeType::kNoBorder },
			{ "wheelmenu"sv, ScaleModeType::kNoBorder }
		};

		const auto it = modes.find(a_fileName);
		return it != modes.end() ? std::optional{ it->second } : std::nullopt;
	}

	// The movie gets the whole framebuffer as its viewport, so kNoBorder scales a 1280x720 menu
	// until it covers the full width and crops whatever no longer fits: at 32:9 that is a 2x zoom
	// with half the menu off screen. Past 16:9 it can only ever crop, so hand out kShowAll instead -
	// that keeps the 16:9 scale and merely leaves more of the stage surroundings visible, which is
	// exactly what SkyUI's widescreen menus stretch their bars over (they cover up to 32:9).
	// Narrower than 16:9 kNoBorder is still the better answer: it hides the strips above and below
	// the stage that kShowAll would expose.
	ScaleModeType ClampToViewport(ScaleModeType a_mode)
	{
		if (a_mode != ScaleModeType::kNoBorder) {
			return a_mode;
		}

		const auto state = RE::BSGraphics::State::GetSingleton();
		if (!state || state->screenHeight == 0) {
			return a_mode;
		}

		// Every vanilla menu is authored on a 1280x720 stage
		constexpr float stageAspect = 1280.0f / 720.0f;
		const float viewportAspect = static_cast<float>(state->screenWidth) / static_cast<float>(state->screenHeight);

		return viewportAspect > stageAspect ? ScaleModeType::kShowAll : a_mode;
	}

	// Third-party UI packs replace table movies with re-authors on stages wider than 16:9 - the
	// Untarnished UI ports (for example "Vel'dun UI Enderal") ship 2560x720 movies for the
	// inventory family, the dialogue menu, the loading menu and the custom hero menu. Those
	// movies implement widescreen themselves, as a kNoBorder crop: the stage always renders at
	// the height-fit scale, so every element keeps one physical size, and the extra stage width
	// stays outside the viewport until the display is as wide as the stage (their loading menu
	// frame keeps a constant 676-stage-unit hole this way, aligned with the engine's 3D art at
	// every aspect ratio). Both modes the table can pick destroy that design: kShowAll must show
	// the whole stage, which letterboxes a 32:9 stage into a half-height band on a 16:9 display,
	// and kExactFit distorts it. So for these movies the mode comes from the stage instead -
	// kNoBorder up to the stage aspect, kShowAll beyond it, where kNoBorder would start cropping
	// again (the two coincide when viewport and stage aspect are equal, hence the loose
	// comparison). Like the clamp, this is not gated on g_restoreScaleModes: pre-1.6.1130 call
	// sites hand these movies an equally wrong mode, kShowAll for the loading menu for instance.
	//
	// The stage size is only readable from the loaded movie, so this runs after the original
	// LoadMovie - which applies its mode argument before returning - and writes the override to
	// the view directly.
	void FitWideStage(RE::GFxMovieView* a_view, const char* a_fileName)
	{
		const auto def = a_view->GetMovieDef();
		const auto state = RE::BSGraphics::State::GetSingleton();
		if (!def || !state || state->screenHeight == 0) {
			return;
		}

		const float stageAspect = def->GetWidth() / def->GetHeight();
		const float viewportAspect = static_cast<float>(state->screenWidth) / static_cast<float>(state->screenHeight);
		const auto mode = viewportAspect >= stageAspect - 0.005f ? ScaleModeType::kShowAll : ScaleModeType::kNoBorder;

		a_view->SetViewScaleMode(mode);

		// Menus reload on every open, so log each movie once
		static std::set<std::string, CaseInsensitiveLess> logged;
		if (logged.emplace(a_fileName).second) {
			logger::info("{} is authored wider than 16:9 ({:.0f}x{:.0f}), overriding its scale mode to {}", a_fileName, def->GetWidth(), def->GetHeight(), mode == ScaleModeType::kShowAll ? "kShowAll" : "kNoBorder");
		}
	}

	struct LoadMovie
	{
		static bool thunk(RE::BSScaleformManager* a_scaleformManager, RE::IMenu* a_menu, RE::GPtr<RE::GFxMovieView>& a_viewOut, const char* a_fileName, ScaleModeType a_mode, float a_backgroundAlpha)
		{
			std::optional<ScaleModeType> tableMode;

			if (a_fileName && a_fileName[0]) {
				tableMode = GetScaleMode(a_fileName);
				if (tableMode) {
					a_mode = ClampToViewport(g_restoreScaleModes ? *tableMode : a_mode);
				}
			}

			if (!func(a_scaleformManager, a_menu, a_viewOut, a_fileName, a_mode, a_backgroundAlpha)) {
				return false;
			}

			// The table still decides which menus may be touched at all
			if (tableMode && a_viewOut && MovieStageIsWiderThan16x9(a_viewOut.get())) {
				FitWideStage(a_viewOut.get(), a_fileName);
			}

			return true;
		}

		static inline REL::Relocation<decltype(thunk)> func;
	};

	void Install()
	{
		g_restoreScaleModes = REL::Module::get().version() >= REL::Version(1, 6, 1130, 0);

		// Checked: 1.5.97, 1.6.640, 1.6.659, 1.6.1130, 1.6.1170, 1.6.1179, 1.7.99, 1.7.104
		const auto target = REL::RelocationID(80302, 82325).address();

		// mov rax,rsp followed by one argument spill - 7 bytes of position-independent code, which
		// the hook relocates into the trampoline. The spilled argument differs between the runtime
		// families, so accept either and bail out rather than corrupt the function if a future
		// build changes the prologue again.
		constexpr std::uint8_t prologueAE[]{ 0x48, 0x8B, 0xC4, 0x48, 0x89, 0x50, 0x10 };  // mov [rax+10h],rdx
		constexpr std::uint8_t prologueSE[]{ 0x48, 0x8B, 0xC4, 0x4C, 0x89, 0x40, 0x18 };  // mov [rax+18h],r8

		static_assert(sizeof(prologueAE) == sizeof(prologueSE));

		const auto prologue = reinterpret_cast<const void*>(target);

		if (!target ||
			(std::memcmp(prologue, prologueAE, sizeof(prologueAE)) != 0 &&
				std::memcmp(prologue, prologueSE, sizeof(prologueSE)) != 0)) {
			logger::error("Unexpected BSScaleformManager::LoadMovie prologue, menu aspect ratio fix is disabled");
			return;
		}

		SKSE::AllocTrampoline(64);
		SKSE::stl::hook_function_prologue<LoadMovie, sizeof(prologueAE)>(target);

		logger::info("Applied menu aspect ratio fix{}", g_restoreScaleModes ? "" : " (viewport clamp only)");
	}
}