#include "iwa/addons/imgui/addon.hpp"

#include <imgui.h>
#include <backends/imgui_impl_sdl2.h>
#include <backends/imgui_impl_vulkan.h>

namespace iwa
{
namespace
{
ImGuiAddon gImguiAddon;

PFN_vkVoidFunction imguiLoaderCallback(const char* functionName, void* userData)
{
    return static_cast<Instance*>(userData)->getVkHandle().getProcAddr(functionName);
}

ImGuiKey keyToImGui(const KeyCode keyCode)
{
    switch (keyCode)
    {
        //<editor-fold desc="Many cases">
        case KeyCode::TAB: return ImGuiKey_Tab;
        case KeyCode::LEFT: return ImGuiKey_LeftArrow;
        case KeyCode::RIGHT: return ImGuiKey_RightArrow;
        case KeyCode::UP: return ImGuiKey_UpArrow;
        case KeyCode::DOWN: return ImGuiKey_DownArrow;
        case KeyCode::PAGEUP: return ImGuiKey_PageUp;
        case KeyCode::PAGEDOWN: return ImGuiKey_PageDown;
        case KeyCode::HOME: return ImGuiKey_Home;
        case KeyCode::END: return ImGuiKey_End;
        case KeyCode::INSERT: return ImGuiKey_Insert;
        case KeyCode::DELETE: return ImGuiKey_Delete;
        case KeyCode::BACKSPACE: return ImGuiKey_Backspace;
        case KeyCode::SPACE: return ImGuiKey_Space;
        case KeyCode::RETURN: return ImGuiKey_Enter;
        case KeyCode::ESCAPE: return ImGuiKey_Escape;
        case KeyCode::QUOTE: return ImGuiKey_Apostrophe;
        case KeyCode::COMMA: return ImGuiKey_Comma;
        case KeyCode::MINUS: return ImGuiKey_Minus;
        case KeyCode::PERIOD: return ImGuiKey_Period;
        case KeyCode::SLASH: return ImGuiKey_Slash;
        case KeyCode::SEMICOLON: return ImGuiKey_Semicolon;
        case KeyCode::EQUALS: return ImGuiKey_Equal;
        case KeyCode::LEFTBRACKET: return ImGuiKey_LeftBracket;
        case KeyCode::BACKSLASH: return ImGuiKey_Backslash;
        case KeyCode::RIGHTBRACKET: return ImGuiKey_RightBracket;
        case KeyCode::BACKQUOTE: return ImGuiKey_GraveAccent;
        case KeyCode::CAPSLOCK: return ImGuiKey_CapsLock;
        case KeyCode::SCROLLLOCK: return ImGuiKey_ScrollLock;
        case KeyCode::NUMLOCKCLEAR: return ImGuiKey_NumLock;
        case KeyCode::PRINTSCREEN: return ImGuiKey_PrintScreen;
        case KeyCode::PAUSE: return ImGuiKey_Pause;
        case KeyCode::KP_0: return ImGuiKey_Keypad0;
        case KeyCode::KP_1: return ImGuiKey_Keypad1;
        case KeyCode::KP_2: return ImGuiKey_Keypad2;
        case KeyCode::KP_3: return ImGuiKey_Keypad3;
        case KeyCode::KP_4: return ImGuiKey_Keypad4;
        case KeyCode::KP_5: return ImGuiKey_Keypad5;
        case KeyCode::KP_6: return ImGuiKey_Keypad6;
        case KeyCode::KP_7: return ImGuiKey_Keypad7;
        case KeyCode::KP_8: return ImGuiKey_Keypad8;
        case KeyCode::KP_9: return ImGuiKey_Keypad9;
        case KeyCode::KP_PERIOD: return ImGuiKey_KeypadDecimal;
        case KeyCode::KP_DIVIDE: return ImGuiKey_KeypadDivide;
        case KeyCode::KP_MULTIPLY: return ImGuiKey_KeypadMultiply;
        case KeyCode::KP_MINUS: return ImGuiKey_KeypadSubtract;
        case KeyCode::KP_PLUS: return ImGuiKey_KeypadAdd;
        case KeyCode::KP_ENTER: return ImGuiKey_KeypadEnter;
        case KeyCode::KP_EQUALS: return ImGuiKey_KeypadEqual;
        case KeyCode::LCTRL: return ImGuiKey_LeftCtrl;
        case KeyCode::LSHIFT: return ImGuiKey_LeftShift;
        case KeyCode::LALT: return ImGuiKey_LeftAlt;
        case KeyCode::LGUI: return ImGuiKey_LeftSuper;
        case KeyCode::RCTRL: return ImGuiKey_RightCtrl;
        case KeyCode::RSHIFT: return ImGuiKey_RightShift;
        case KeyCode::RALT: return ImGuiKey_RightAlt;
        case KeyCode::RGUI: return ImGuiKey_RightSuper;
        case KeyCode::APPLICATION: return ImGuiKey_Menu;
        case KeyCode::_0: return ImGuiKey_0;
        case KeyCode::_1: return ImGuiKey_1;
        case KeyCode::_2: return ImGuiKey_2;
        case KeyCode::_3: return ImGuiKey_3;
        case KeyCode::_4: return ImGuiKey_4;
        case KeyCode::_5: return ImGuiKey_5;
        case KeyCode::_6: return ImGuiKey_6;
        case KeyCode::_7: return ImGuiKey_7;
        case KeyCode::_8: return ImGuiKey_8;
        case KeyCode::_9: return ImGuiKey_9;
        case KeyCode::A: return ImGuiKey_A;
        case KeyCode::B: return ImGuiKey_B;
        case KeyCode::C: return ImGuiKey_C;
        case KeyCode::D: return ImGuiKey_D;
        case KeyCode::E: return ImGuiKey_E;
        case KeyCode::F: return ImGuiKey_F;
        case KeyCode::G: return ImGuiKey_G;
        case KeyCode::H: return ImGuiKey_H;
        case KeyCode::I: return ImGuiKey_I;
        case KeyCode::J: return ImGuiKey_J;
        case KeyCode::K: return ImGuiKey_K;
        case KeyCode::L: return ImGuiKey_L;
        case KeyCode::M: return ImGuiKey_M;
        case KeyCode::N: return ImGuiKey_N;
        case KeyCode::O: return ImGuiKey_O;
        case KeyCode::P: return ImGuiKey_P;
        case KeyCode::Q: return ImGuiKey_Q;
        case KeyCode::R: return ImGuiKey_R;
        case KeyCode::S: return ImGuiKey_S;
        case KeyCode::T: return ImGuiKey_T;
        case KeyCode::U: return ImGuiKey_U;
        case KeyCode::V: return ImGuiKey_V;
        case KeyCode::W: return ImGuiKey_W;
        case KeyCode::X: return ImGuiKey_X;
        case KeyCode::Y: return ImGuiKey_Y;
        case KeyCode::Z: return ImGuiKey_Z;
        case KeyCode::F1: return ImGuiKey_F1;
        case KeyCode::F2: return ImGuiKey_F2;
        case KeyCode::F3: return ImGuiKey_F3;
        case KeyCode::F4: return ImGuiKey_F4;
        case KeyCode::F5: return ImGuiKey_F5;
        case KeyCode::F6: return ImGuiKey_F6;
        case KeyCode::F7: return ImGuiKey_F7;
        case KeyCode::F8: return ImGuiKey_F8;
        case KeyCode::F9: return ImGuiKey_F9;
        case KeyCode::F10: return ImGuiKey_F10;
        case KeyCode::F11: return ImGuiKey_F11;
        case KeyCode::F12: return ImGuiKey_F12;
        case KeyCode::F13: return ImGuiKey_F13;
        case KeyCode::F14: return ImGuiKey_F14;
        case KeyCode::F15: return ImGuiKey_F15;
        case KeyCode::F16: return ImGuiKey_F16;
        case KeyCode::F17: return ImGuiKey_F17;
        case KeyCode::F18: return ImGuiKey_F18;
        case KeyCode::F19: return ImGuiKey_F19;
        case KeyCode::F20: return ImGuiKey_F20;
        case KeyCode::F21: return ImGuiKey_F21;
        case KeyCode::F22: return ImGuiKey_F22;
        case KeyCode::F23: return ImGuiKey_F23;
        case KeyCode::F24: return ImGuiKey_F24;
        case KeyCode::AC_BACK: return ImGuiKey_AppBack;
        case KeyCode::AC_FORWARD: return ImGuiKey_AppForward;
        default: return ImGuiKey_None;
            //</editor-fold>
    }
}

ImGuiMouseButton mouseButtonToImGui(const MouseButton button)
{
    switch (button)
    {
        case MouseButton::LEFT:
            return ImGuiMouseButton_Left;
        case MouseButton::RIGHT:
            return ImGuiMouseButton_Right;
        case MouseButton::MIDDLE:
            return ImGuiMouseButton_Middle;
        case MouseButton::EXTRA_1:
            return ImGuiMouseButton_Middle + 1;
        case MouseButton::EXTRA_2:
            return ImGuiMouseButton_Middle + 2;
    }
    return -1;
}
}

void ImGuiAddon::init(const AddonInitArgs& args)
{
    (void) args;

    IMGUI_CHECKVERSION();
    ImGui::CreateContext();
    ImGui::StyleColorsDark();
}

void ImGuiAddon::cleanup()
{
    ImGui::DestroyContext();
}

mijin::Task<> ImGuiAddon::c_createResources(const ImguiCreateResourcesArgs& args)
{
    const unsigned QUEUED_FRAMES = 3; // TODO: is this okay?
    const unsigned MAX_IMAGES = 256;
    Device& device = *args.swapchain.getOwner();

    mDescriptorPool = device.createChild<DescriptorPool>(DescriptorPoolCreationArgs{
        .flags = vk::DescriptorPoolCreateFlagBits::eFreeDescriptorSet,
        .maxSets = MAX_IMAGES,
        .poolSizes = {
            vk::DescriptorPoolSize{.type = vk::DescriptorType::eCombinedImageSampler, .descriptorCount = MAX_IMAGES}
        }
    });
    registerObjectDestructionHandler(args.swapchain, [this, device = device.getPointer()]
    {
        // destroy the descriptor pool when the window is destroyed so we don't keep the device alive
        device->getVkHandle().waitIdle();
        ImGui_ImplVulkan_Shutdown();
        ImGui_ImplSDL2_Shutdown();
        mDescriptorPool = nullptr;
        mWidgets.clear();
    });

    if (!ImGui_ImplSDL2_InitForVulkan(args.swapchain.getWindow()->getSDLWindow())) {
        throw std::runtime_error("Error initializing ImGui for SDL2.");
    }

    ImGui_ImplVulkan_InitInfo initInfo{
        .Instance = device.getOwner()->getVkHandle(),
        .PhysicalDevice = device.getVkPhysicalDevice(),
        .Device = device.getVkHandle(),
        .QueueFamily = device.getDeviceInfo().graphicsQueueFamily,
        .Queue = device.getGraphicsQueue(),
        .PipelineCache = VK_NULL_HANDLE,
        .DescriptorPool = mDescriptorPool->getVkHandle(),
        .Subpass = 0,
        .MinImageCount = QUEUED_FRAMES,
        .ImageCount = QUEUED_FRAMES,
        .MSAASamples = VK_SAMPLE_COUNT_1_BIT,
        .UseDynamicRendering = true,
        .ColorAttachmentFormat = static_cast<VkFormat>(args.format),
        .Allocator = nullptr,
        .CheckVkResultFn = nullptr,
    };

    const bool success = ImGui_ImplVulkan_LoadFunctions(imguiLoaderCallback, device.getOwner())
        && ImGui_ImplVulkan_Init(&initInfo, VK_NULL_HANDLE);
    if (!success) {
        throw std::runtime_error("Error initializing ImGui for Vulkan.");
    }

    // setup input
    args.swapchain.getWindow()->keyChanged.connect(*this, &ImGuiAddon::handleKeyChanged);
    args.swapchain.getWindow()->mouseButtonChanged.connect(*this, &ImGuiAddon::handleMouseButtonChanged);
    args.swapchain.getWindow()->mouseMoved.connect(*this, &ImGuiAddon::handleMouseMoved);
    args.swapchain.getWindow()->mouseScrolled.connect(*this, &ImGuiAddon::handleMouseScrolled);
    args.swapchain.getWindow()->textEntered.connect(*this, &ImGuiAddon::handleTextEntered);

    // first frame
    beginFrame();
    co_return;
}

void ImGuiAddon::renderFrame(vk::CommandBuffer cmdBuffer)
{
    for (const std::unique_ptr<ImGuiWidget>& widget : mWidgets) {
        widget->draw();
    }
    ImGui::Render();
    ImGui_ImplVulkan_RenderDrawData(ImGui::GetDrawData(), cmdBuffer);
    beginFrame();
}

void ImGuiAddon::removeWidget(ImGuiWidget* widget)
{
    auto it = std::ranges::find_if(mWidgets, [widget](const std::unique_ptr<ImGuiWidget>& widgetPtr)
    {
        return widgetPtr.get() == widget;
    });
    mWidgets.erase(it);
}

void ImGuiAddon::beginFrame() noexcept
{
    ImGui_ImplVulkan_NewFrame();
    ImGui_ImplSDL2_NewFrame();
    ImGui::NewFrame();
}

void ImGuiAddon::handleKeyChanged(const KeyEvent& event)
{
    ImGui::GetIO().AddKeyEvent(keyToImGui(event.keyCode), event.down);
}

void ImGuiAddon::handleMouseButtonChanged(const MouseButtonEvent& event)
{
    ImGui::GetIO().AddMouseButtonEvent(mouseButtonToImGui(event.button), event.down);
}

void ImGuiAddon::handleMouseMoved(const MouseMoveEvent& event)
{
    ImGui::GetIO().AddMousePosEvent(static_cast<float>(event.absoluteX), static_cast<float>(event.absoluteY));
}

void ImGuiAddon::handleMouseScrolled(const MouseWheelEvent& event)
{
    ImGui::GetIO().AddMouseWheelEvent(static_cast<float>(event.relativeX), static_cast<float>(event.relativeY));
}

void ImGuiAddon::handleTextEntered(const TextInputEvent& event)
{
    ImGui::GetIO().AddInputCharactersUTF8(event.text.c_str());
}

ImGuiAddon& ImGuiAddon::get() noexcept
{
    return gImguiAddon;
}
} // namespace iwa