#pragma once

#if defined(SKSE_SUPPORT_XBYAK)
namespace Xbyak
{
	class CodeGenerator;
}
#endif

namespace SKSE
{
	class Trampoline
	{
	public:
		using deleter_type = std::function<void(void* a_mem, std::size_t a_size)>;

		Trampoline() = default;
		Trampoline(const Trampoline&) = delete;

		Trampoline(Trampoline&& a_rhs) noexcept { move_from(std::move(a_rhs)); }

		explicit Trampoline(std::string_view a_name) :
			_name(a_name)
		{}

		~Trampoline() { release(); }

		Trampoline& operator=(const Trampoline&) = delete;

		Trampoline& operator=(Trampoline&& a_rhs) noexcept
		{
			if (this != std::addressof(a_rhs)) {
				move_from(std::move(a_rhs));
			}
			return *this;
		}

		void create(std::size_t a_size) { return create(a_size, nullptr); }
		void create(std::size_t a_size, void* a_module);

		void set_trampoline(void* a_trampoline, std::size_t a_size, deleter_type a_deleter = {})
		{
			auto trampoline = static_cast<std::byte*>(a_trampoline);
			if (trampoline) {
				constexpr auto INT3 = static_cast<int>(0xCC);
				std::memset(trampoline, INT3, a_size);
			}

			release();

			_deleter = std::move(a_deleter);
			_data = trampoline;
			_capacity = a_size;
			_size = 0;

			log_stats();
		}

		[[nodiscard]] void* allocate(std::size_t a_size)
		{
			auto result = do_allocate(a_size);
			log_stats();
			return result;
		}

#ifdef SKSE_SUPPORT_XBYAK
		[[nodiscard]] void* allocate(Xbyak::CodeGenerator& a_code);
#endif

		template <class T>
		[[nodiscard]] T* allocate()
		{
			return static_cast<T*>(allocate(sizeof(T)));
		}

		[[nodiscard]] constexpr std::size_t empty() const noexcept { return _capacity == 0; }
		[[nodiscard]] constexpr std::size_t capacity() const noexcept { return _capacity; }
		[[nodiscard]] constexpr std::size_t allocated_size() const noexcept { return _size; }
		[[nodiscard]] constexpr std::size_t free_size() const noexcept { return _capacity - _size; }

		template <std::size_t N>
		std::uintptr_t write_branch(std::uintptr_t a_src, std::uintptr_t a_dst)
		{
			std::uint8_t data = 0;
			if constexpr (N == 5) {
				// E9 cd
				// JMP rel32
				data = 0xE9;
			} else if constexpr (N == 6) {
				// FF /4
				// JMP r/m64
				data = 0x25;
			} else {
				static_assert(false && N, "invalid branch size");
			}

			return write_branch<N>(a_src, a_dst, data);
		}

		template <std::size_t N, class F>
		std::uintptr_t write_branch(std::uintptr_t a_src, F a_dst)
		{
			return write_branch<N>(a_src, stl::unrestricted_cast<std::uintptr_t>(a_dst));
		}

		template <std::size_t N>
		std::uintptr_t write_call(std::uintptr_t a_src, std::uintptr_t a_dst)
		{
			std::uint8_t data = 0;
			if constexpr (N == 5) {
				// E8 cd
				// CALL rel32
				data = 0xE8;
			} else if constexpr (N == 6) {
				// FF /2
				// CALL r/m64
				data = 0x15;
			} else {
				static_assert(false && N, "invalid call size");
			}

			return write_branch<N>(a_src, a_dst, data);
		}

		template <std::size_t N, class F>
		std::uintptr_t write_call(std::uintptr_t a_src, F a_dst)
		{
			return write_call<N>(a_src, stl::unrestricted_cast<std::uintptr_t>(a_dst));
		}

	private:
		[[nodiscard]] void* do_create(std::size_t a_size, std::uintptr_t a_address);
		[[nodiscard]] void* do_allocate(std::size_t a_size);

		void write_5branch(std::uintptr_t a_src, std::uintptr_t a_dst, std::uint8_t a_opcode);
		void write_6branch(std::uintptr_t a_src, std::uintptr_t a_dst, std::uint8_t a_modrm);

		template <std::size_t N>
		[[nodiscard]] std::uintptr_t write_branch(std::uintptr_t a_src, std::uintptr_t a_dst, std::uint8_t a_data)
		{
			const auto isNop = *reinterpret_cast<std::int8_t*>(a_src) == 0x90;
			const auto disp = reinterpret_cast<std::int32_t*>(a_src + N - 4);
			const auto nextOp = a_src + N;
			const auto func = isNop ? 0 : nextOp + *disp;

			if constexpr (N == 5) {
				write_5branch(a_src, a_dst, a_data);
			} else if constexpr (N == 6) {
				write_6branch(a_src, a_dst, a_data);
			} else {
				static_assert(false && N, "invalid branch size");
			}

			return func;
		}

		void move_from(Trampoline&& a_rhs)
		{
			_5branches = std::move(a_rhs._5branches);
			_6branches = std::move(a_rhs._6branches);
			_name = std::move(a_rhs._name);

			_deleter = std::move(a_rhs._deleter);

			_data = a_rhs._data;
			a_rhs._data = nullptr;

			_capacity = a_rhs._capacity;
			a_rhs._capacity = 0;

			_size = a_rhs._size;
			a_rhs._size = 0;
		}

		void log_stats() const;

		[[nodiscard]] bool in_range(std::ptrdiff_t a_disp) const
		{
			constexpr auto min = (std::numeric_limits<std::int32_t>::min)();
			constexpr auto max = (std::numeric_limits<std::int32_t>::max)();

			return min <= a_disp && a_disp <= max;
		}

		void release()
		{
			if (_data && _deleter) {
				_deleter(_data, _capacity);
			}

			_5branches.clear();
			_6branches.clear();
			_data = nullptr;
			_capacity = 0;
			_size = 0;
		}

		std::map<std::uintptr_t, std::byte*> _5branches;
		std::map<std::uintptr_t, std::byte*> _6branches;
		std::string                          _name{ "Default Trampoline"sv };
		deleter_type                         _deleter;
		std::byte*                           _data{ nullptr };
		std::size_t                          _capacity{ 0 };
		std::size_t                          _size{ 0 };
	};

	Trampoline& GetTrampoline();
}
