#include "BuiltInIcons.h"

#include "ImGuiBackendBridge.h"
#include "Platform/Win32Utf8.h"
#include "StyleTokens.h"

#include <algorithm>
#include <cstdint>
#include <cstring>
#include <filesystem>
#include <string>
#include <wrl/client.h>

#include <stb_image.h>

namespace XCEngine {
namespace Editor {
namespace UI {

namespace {

using Microsoft::WRL::ComPtr;

struct BuiltInTexture {
    ImTextureID textureId = {};
    D3D12_CPU_DESCRIPTOR_HANDLE cpuHandle = {};
    D3D12_GPU_DESCRIPTOR_HANDLE gpuHandle = {};
    ComPtr<ID3D12Resource> texture;
    int width = 0;
    int height = 0;

    bool IsValid() const {
        return textureId != ImTextureID{} && texture != nullptr && width > 0 && height > 0;
    }
};

struct BuiltInIconState {
    ImGuiBackendBridge* backend = nullptr;
    BuiltInTexture folder;
    BuiltInTexture gameObject;
};

BuiltInIconState g_icons;

std::filesystem::path ResolveFolderIconPath() {
    const std::filesystem::path exeDir(Platform::GetExecutableDirectoryUtf8());
    return (exeDir / ".." / ".." / "resources" / "Icons" / "folder_icon.png").lexically_normal();
}

std::filesystem::path ResolveGameObjectIconPath() {
    const std::filesystem::path exeDir(Platform::GetExecutableDirectoryUtf8());
    return (exeDir / ".." / ".." / "resources" / "Icons" / "gameobject_icon.png").lexically_normal();
}

void ResetTexture(BuiltInTexture& texture) {
    if (g_icons.backend && texture.cpuHandle.ptr != 0) {
        g_icons.backend->FreeTextureDescriptor(texture.cpuHandle, texture.gpuHandle);
    }

    texture.texture.Reset();
    texture.textureId = {};
    texture.cpuHandle = {};
    texture.gpuHandle = {};
    texture.width = 0;
    texture.height = 0;
}

bool WaitForQueueIdle(ID3D12Device* device, ID3D12CommandQueue* commandQueue) {
    if (!device || !commandQueue) {
        return false;
    }

    ComPtr<ID3D12Fence> fence;
    if (FAILED(device->CreateFence(0, D3D12_FENCE_FLAG_NONE, IID_PPV_ARGS(&fence)))) {
        return false;
    }

    HANDLE eventHandle = CreateEventW(nullptr, FALSE, FALSE, nullptr);
    if (!eventHandle) {
        return false;
    }

    constexpr UINT64 kFenceValue = 1;
    const HRESULT signalHr = commandQueue->Signal(fence.Get(), kFenceValue);
    if (FAILED(signalHr)) {
        CloseHandle(eventHandle);
        return false;
    }

    if (fence->GetCompletedValue() < kFenceValue) {
        if (FAILED(fence->SetEventOnCompletion(kFenceValue, eventHandle))) {
            CloseHandle(eventHandle);
            return false;
        }
        WaitForSingleObject(eventHandle, INFINITE);
    }

    CloseHandle(eventHandle);
    return true;
}

bool LoadTextureFromFile(
    ImGuiBackendBridge& backend,
    ID3D12Device* device,
    ID3D12CommandQueue* commandQueue,
    const std::filesystem::path& filePath,
    BuiltInTexture& outTexture) {
    if (!device || !commandQueue || !std::filesystem::exists(filePath)) {
        return false;
    }

    int width = 0;
    int height = 0;
    int channels = 0;
    stbi_uc* pixels = stbi_load(filePath.string().c_str(), &width, &height, &channels, STBI_rgb_alpha);
    if (!pixels || width <= 0 || height <= 0) {
        if (pixels) {
            stbi_image_free(pixels);
        }
        return false;
    }

    const UINT srcRowPitch = static_cast<UINT>(width * 4);

    D3D12_RESOURCE_DESC textureDesc = {};
    textureDesc.Dimension = D3D12_RESOURCE_DIMENSION_TEXTURE2D;
    textureDesc.Alignment = 0;
    textureDesc.Width = static_cast<UINT64>(width);
    textureDesc.Height = static_cast<UINT>(height);
    textureDesc.DepthOrArraySize = 1;
    textureDesc.MipLevels = 1;
    textureDesc.Format = DXGI_FORMAT_R8G8B8A8_UNORM;
    textureDesc.SampleDesc.Count = 1;
    textureDesc.Layout = D3D12_TEXTURE_LAYOUT_UNKNOWN;

    D3D12_HEAP_PROPERTIES defaultHeap = {};
    defaultHeap.Type = D3D12_HEAP_TYPE_DEFAULT;

    ComPtr<ID3D12Resource> textureResource;
    if (FAILED(device->CreateCommittedResource(
            &defaultHeap,
            D3D12_HEAP_FLAG_NONE,
            &textureDesc,
            D3D12_RESOURCE_STATE_COPY_DEST,
            nullptr,
            IID_PPV_ARGS(&textureResource)))) {
        stbi_image_free(pixels);
        return false;
    }

    D3D12_PLACED_SUBRESOURCE_FOOTPRINT footprint = {};
    UINT numRows = 0;
    UINT64 rowSizeInBytes = 0;
    UINT64 uploadBufferSize = 0;
    device->GetCopyableFootprints(&textureDesc, 0, 1, 0, &footprint, &numRows, &rowSizeInBytes, &uploadBufferSize);

    D3D12_RESOURCE_DESC uploadDesc = {};
    uploadDesc.Dimension = D3D12_RESOURCE_DIMENSION_BUFFER;
    uploadDesc.Width = uploadBufferSize;
    uploadDesc.Height = 1;
    uploadDesc.DepthOrArraySize = 1;
    uploadDesc.MipLevels = 1;
    uploadDesc.SampleDesc.Count = 1;
    uploadDesc.Layout = D3D12_TEXTURE_LAYOUT_ROW_MAJOR;

    D3D12_HEAP_PROPERTIES uploadHeap = {};
    uploadHeap.Type = D3D12_HEAP_TYPE_UPLOAD;

    ComPtr<ID3D12Resource> uploadResource;
    if (FAILED(device->CreateCommittedResource(
            &uploadHeap,
            D3D12_HEAP_FLAG_NONE,
            &uploadDesc,
            D3D12_RESOURCE_STATE_GENERIC_READ,
            nullptr,
            IID_PPV_ARGS(&uploadResource)))) {
        stbi_image_free(pixels);
        return false;
    }

    std::uint8_t* mappedData = nullptr;
    if (FAILED(uploadResource->Map(0, nullptr, reinterpret_cast<void**>(&mappedData)))) {
        stbi_image_free(pixels);
        return false;
    }

    for (UINT row = 0; row < numRows; ++row) {
        std::memcpy(
            mappedData + footprint.Offset + static_cast<SIZE_T>(row) * footprint.Footprint.RowPitch,
            pixels + static_cast<size_t>(row) * srcRowPitch,
            srcRowPitch);
    }
    uploadResource->Unmap(0, nullptr);
    stbi_image_free(pixels);

    ComPtr<ID3D12CommandAllocator> commandAllocator;
    ComPtr<ID3D12GraphicsCommandList> commandList;
    if (FAILED(device->CreateCommandAllocator(D3D12_COMMAND_LIST_TYPE_DIRECT, IID_PPV_ARGS(&commandAllocator)))) {
        return false;
    }
    if (FAILED(device->CreateCommandList(
            0,
            D3D12_COMMAND_LIST_TYPE_DIRECT,
            commandAllocator.Get(),
            nullptr,
            IID_PPV_ARGS(&commandList)))) {
        return false;
    }

    D3D12_TEXTURE_COPY_LOCATION dst = {};
    dst.pResource = textureResource.Get();
    dst.Type = D3D12_TEXTURE_COPY_TYPE_SUBRESOURCE_INDEX;
    dst.SubresourceIndex = 0;

    D3D12_TEXTURE_COPY_LOCATION src = {};
    src.pResource = uploadResource.Get();
    src.Type = D3D12_TEXTURE_COPY_TYPE_PLACED_FOOTPRINT;
    src.PlacedFootprint = footprint;

    commandList->CopyTextureRegion(&dst, 0, 0, 0, &src, nullptr);

    D3D12_RESOURCE_BARRIER barrier = {};
    barrier.Type = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION;
    barrier.Transition.pResource = textureResource.Get();
    barrier.Transition.StateBefore = D3D12_RESOURCE_STATE_COPY_DEST;
    barrier.Transition.StateAfter = D3D12_RESOURCE_STATE_PIXEL_SHADER_RESOURCE;
    barrier.Transition.Subresource = D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES;
    commandList->ResourceBarrier(1, &barrier);

    if (FAILED(commandList->Close())) {
        return false;
    }

    ID3D12CommandList* commandLists[] = { commandList.Get() };
    commandQueue->ExecuteCommandLists(1, commandLists);

    if (!WaitForQueueIdle(device, commandQueue)) {
        return false;
    }

    backend.AllocateTextureDescriptor(&outTexture.cpuHandle, &outTexture.gpuHandle);

    D3D12_SHADER_RESOURCE_VIEW_DESC srvDesc = {};
    srvDesc.Shader4ComponentMapping = D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING;
    srvDesc.Format = DXGI_FORMAT_R8G8B8A8_UNORM;
    srvDesc.ViewDimension = D3D12_SRV_DIMENSION_TEXTURE2D;
    srvDesc.Texture2D.MipLevels = 1;
    device->CreateShaderResourceView(textureResource.Get(), &srvDesc, outTexture.cpuHandle);

    outTexture.texture = textureResource;
    outTexture.textureId = (ImTextureID)(static_cast<intptr_t>(outTexture.gpuHandle.ptr));
    outTexture.width = width;
    outTexture.height = height;
    return true;
}

ImVec2 ComputeFittedIconSize(const BuiltInTexture& texture, const ImVec2& min, const ImVec2& max) {
    const float availableWidth = max.x - min.x;
    const float availableHeight = max.y - min.y;
    if (availableWidth <= 0.0f || availableHeight <= 0.0f || texture.width <= 0 || texture.height <= 0) {
        return ImVec2(0.0f, 0.0f);
    }

    const float scale = (std::min)(
        availableWidth / static_cast<float>(texture.width),
        availableHeight / static_cast<float>(texture.height));
    return ImVec2(
        static_cast<float>(texture.width) * scale,
        static_cast<float>(texture.height) * scale);
}

void DrawTextureIcon(ImDrawList* drawList, const BuiltInTexture& texture, const ImVec2& min, const ImVec2& max) {
    if (!drawList || !texture.IsValid()) {
        return;
    }

    const ImVec2 size = ComputeFittedIconSize(texture, min, max);
    const float x = min.x + ((max.x - min.x) - size.x) * 0.5f;
    const float y = min.y + ((max.y - min.y) - size.y) * 0.5f;
    drawList->AddImage(texture.textureId, ImVec2(x, y), ImVec2(x + size.x, y + size.y));
}

void DrawBuiltInFolderFallback(ImDrawList* drawList, const ImVec2& min, const ImVec2& max) {
    if (!drawList) {
        return;
    }

    const float width = max.x - min.x;
    const float height = max.y - min.y;
    if (width <= 0.0f || height <= 0.0f) {
        return;
    }

    const float rounding = (std::max)(1.0f, (std::min)(width, height) * 0.18f);
    const ImU32 tabColor = ImGui::GetColorU32(BuiltInFolderIconTabColor());
    const ImU32 topColor = ImGui::GetColorU32(BuiltInFolderIconTopColor());
    const ImU32 bodyColor = ImGui::GetColorU32(BuiltInFolderIconBodyColor());

    const ImVec2 tabMin(min.x + width * 0.08f, min.y + height * 0.14f);
    const ImVec2 tabMax(min.x + width * 0.48f, min.y + height * 0.38f);
    const ImVec2 topMin(min.x + width * 0.24f, min.y + height * 0.22f);
    const ImVec2 topMax(min.x + width * 0.90f, min.y + height * 0.42f);
    const ImVec2 bodyMin(min.x + width * 0.06f, min.y + height * 0.32f);
    const ImVec2 bodyMax(min.x + width * 0.94f, min.y + height * 0.88f);

    drawList->AddRectFilled(tabMin, tabMax, tabColor, rounding);
    drawList->AddRectFilled(
        topMin,
        topMax,
        topColor,
        rounding,
        ImDrawFlags_RoundCornersTopLeft | ImDrawFlags_RoundCornersTopRight);
    drawList->AddRectFilled(bodyMin, bodyMax, bodyColor, rounding);
}

void DrawBuiltInFileIcon(ImDrawList* drawList, const ImVec2& min, const ImVec2& max) {
    if (!drawList) {
        return;
    }

    const ImU32 fillColor = ImGui::GetColorU32(AssetFileIconFillColor());
    const ImU32 lineColor = ImGui::GetColorU32(AssetFileIconLineColor());
    const ImVec2 foldA(max.x - 8.0f, min.y);
    const ImVec2 foldB(max.x, min.y + 8.0f);
    drawList->AddRectFilled(min, max, fillColor, 2.0f);
    drawList->AddRect(min, max, lineColor, 2.0f);
    drawList->AddTriangleFilled(foldA, ImVec2(max.x, min.y), foldB, ImGui::GetColorU32(AssetFileFoldColor()));
    drawList->AddLine(foldA, foldB, lineColor);
}

} // namespace

void InitializeBuiltInIcons(
    ImGuiBackendBridge& backend,
    ID3D12Device* device,
    ID3D12CommandQueue* commandQueue) {
    ShutdownBuiltInIcons();
    g_icons.backend = &backend;
    LoadTextureFromFile(backend, device, commandQueue, ResolveFolderIconPath(), g_icons.folder);
    LoadTextureFromFile(backend, device, commandQueue, ResolveGameObjectIconPath(), g_icons.gameObject);
}

void ShutdownBuiltInIcons() {
    ResetTexture(g_icons.folder);
    ResetTexture(g_icons.gameObject);
    g_icons.backend = nullptr;
}

void DrawAssetIcon(ImDrawList* drawList, const ImVec2& min, const ImVec2& max, AssetIconKind kind) {
    if (kind == AssetIconKind::Folder) {
        if (g_icons.folder.IsValid()) {
            DrawTextureIcon(drawList, g_icons.folder, min, max);
            return;
        }

        DrawBuiltInFolderFallback(drawList, min, max);
        return;
    }

    if (kind == AssetIconKind::GameObject) {
        if (g_icons.gameObject.IsValid()) {
            DrawTextureIcon(drawList, g_icons.gameObject, min, max);
            return;
        }

        DrawBuiltInFileIcon(drawList, min, max);
        return;
    }

    DrawBuiltInFileIcon(drawList, min, max);
}

} // namespace UI
} // namespace Editor
} // namespace XCEngine