#include "MetaCoreRender/MetaCoreImGuiHelper.h"

#include <vector>
#include <unordered_map>
#include <fstream>
#include <iostream>

#include <imgui.h>

#include <filament/Camera.h>
#include <filament/Fence.h>
#include <filament/IndexBuffer.h>
#include <filament/Material.h>
#include <filament/MaterialInstance.h>
#include <filament/RenderableManager.h>
#include <filament/Scene.h>
#include <filament/Texture.h>
#include <filament/TextureSampler.h>
#include <filament/TransformManager.h>
#include <filament/VertexBuffer.h>

#include <utils/EntityManager.h>

using namespace filament::math;
using namespace filament;
using namespace utils;

using MinFilter = TextureSampler::MinFilter;
using MagFilter = TextureSampler::MagFilter;

namespace MetaCore {

MetaCoreImGuiHelper::MetaCoreImGuiHelper(Engine* engine, filament::View* view, const std::filesystem::path& fontPath,
        ImGuiContext *imGuiContext)
        : mEngine(engine), mView(view), mScene(engine->createScene()),
        mImGuiContext(imGuiContext ? imGuiContext : ImGui::CreateContext()),
        mOwnsImGuiContext(imGuiContext == nullptr) {
    
    ImGui::SetCurrentContext(mImGuiContext);
    ImGuiIO& io = ImGui::GetIO();
    
    // 加载我们编译好的材质
    std::ifstream file("uiBlit.filamat", std::ios::binary);
    if (!file.is_open()) {
        file.open("../uiBlit.filamat", std::ios::binary);
    }
    if (!file.is_open()) {
        file.open("build/uiBlit.filamat", std::ios::binary);
    }
    
    if (!file.is_open()) {
        std::cerr << "MetaCoreImGuiHelper: Failed to open uiBlit.filamat from any expected location!" << std::endl;
    }
    
    std::vector<char> buffer((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());
    
    if (buffer.empty()) {
        std::cerr << "MetaCoreImGuiHelper: Material buffer is empty! Skipping material creation to avoid crash." << std::endl;
        mMaterial2d = nullptr;
    } else {
        mMaterial2d = Material::Builder()
                .package(buffer.data(), buffer.size())
                .build(*engine);
    }

    if (!fontPath.empty() && std::filesystem::exists(fontPath)) {
        io.Fonts->AddFontFromFileTTF(fontPath.string().c_str(), 16.0f, nullptr, io.Fonts->GetGlyphRangesChineseSimplifiedCommon());
    }
    createAtlasTexture(engine);

    utils::EntityManager& em = utils::EntityManager::get();
    mCameraEntity = em.create();
    mCamera = mEngine->createCamera(mCameraEntity);

    view->setCamera(mCamera);
    view->setPostProcessingEnabled(false);
    view->setBlendMode(View::BlendMode::TRANSLUCENT);
    view->setShadowingEnabled(false);

    // 绑定场景
    view->setScene(mScene);

    mRenderable = em.create();
    mScene->addEntity(mRenderable);

    ImGui::StyleColorsDark();
}

void MetaCoreImGuiHelper::createAtlasTexture(Engine* engine) {
    if (mTexture) {
        engine->destroy(mTexture);
    }
    
    ImGuiIO& io = ImGui::GetIO();
    unsigned char* pixels;
    int width, height;
    io.Fonts->GetTexDataAsRGBA32(&pixels, &width, &height);
    size_t size = (size_t) (width * height * 4);
    
    Texture::PixelBufferDescriptor pb(
            pixels, size,
            Texture::Format::RGBA, Texture::Type::UBYTE);
            
    mTexture = Texture::Builder()
            .width((uint32_t) width)
            .height((uint32_t) height)
            .levels((uint8_t) 1)
            .format(Texture::InternalFormat::RGBA8)
            .sampler(Texture::Sampler::SAMPLER_2D)
            .build(*engine);
            
    mTexture->setImage(*engine, 0, std::move(pb));

    mSampler = TextureSampler(MinFilter::LINEAR, MagFilter::LINEAR);
    if (mMaterial2d) {
        mMaterial2d->setDefaultParameter("albedo", mTexture, mSampler);
    }
}

MetaCoreImGuiHelper::~MetaCoreImGuiHelper() {
    mEngine->destroy(mScene);
    mEngine->destroy(mRenderable);
    mEngine->destroyCameraComponent(mCameraEntity);

    for (auto& mi : mMaterial2dInstances) {
        mEngine->destroy(mi);
    }
    mEngine->destroy(mMaterial2d);
    mEngine->destroy(mTexture);
    
    for (auto& vb : mVertexBuffers) {
        mEngine->destroy(vb);
    }
    for (auto& ib : mIndexBuffers) {
        mEngine->destroy(ib);
    }

    for (auto itextures: mImGuiTextures) {
        mEngine->destroy(itextures);
    }

    EntityManager& em = utils::EntityManager::get();
    em.destroy(mRenderable);
    em.destroy(mCameraEntity);

    if (mImGuiContext == ImGui::GetCurrentContext()) {
        ImGui::SetCurrentContext(nullptr);
    }
    if (mOwnsImGuiContext) {
        ImGui::DestroyContext(mImGuiContext);
    }
}

void MetaCoreImGuiHelper::setDisplaySize(int width, int height, float scaleX, float scaleY,
        bool flipVertical) {
    ImGuiIO& io = ImGui::GetIO();
    io.DisplaySize = ImVec2(width, height);
    io.DisplayFramebufferScale.x = scaleX;
    io.DisplayFramebufferScale.y = scaleY;
    mFlipVertical = flipVertical;
    
    if (flipVertical) {
        mCamera->setProjection(Camera::Projection::ORTHO, 0.0, double(width),
                0.0, double(height), 0.0, 1.0);
    } else {
        mCamera->setProjection(Camera::Projection::ORTHO, 0.0, double(width),
                double(height), 0.0, 0.0, 1.0);
    }
}

void MetaCoreImGuiHelper::render(float timeStepInSeconds, Callback imguiCommands) {
    ImGui::SetCurrentContext(mImGuiContext);
    ImGuiIO& io = ImGui::GetIO();
    io.DeltaTime = timeStepInSeconds;
    
    ImGui::NewFrame();
    imguiCommands(mEngine, mView);
    ImGui::Render();
    
    processImGuiCommands(ImGui::GetDrawData(), ImGui::GetIO());
}

void MetaCoreImGuiHelper::processImGuiCommands(ImDrawData* commands, const ImGuiIO& io) {
    ImGui::SetCurrentContext(mImGuiContext);
    
    mHasSynced = false;
    auto& rcm = mEngine->getRenderableManager();

    int fbwidth = (int)(io.DisplaySize.x * io.DisplayFramebufferScale.x);
    int fbheight = (int)(io.DisplaySize.y * io.DisplayFramebufferScale.y);
    if (fbwidth == 0 || fbheight == 0)
        return;
    commands->ScaleClipRects(io.DisplayFramebufferScale);

    createBuffers(commands->CmdListsCount);

    size_t nPrims = 0;
    for (int cmdListIndex = 0; cmdListIndex < commands->CmdListsCount; cmdListIndex++) {
        const ImDrawList* cmds = commands->CmdLists[cmdListIndex];
        nPrims += cmds->CmdBuffer.size();
    }
    
    auto rbuilder = RenderableManager::Builder(nPrims);
    rbuilder.boundingBox({{ 0, 0, 0 }, { 10000, 10000, 10000 }}).culling(false);

    rcm.destroy(mRenderable);
    int bufferIndex = 0;
    int primIndex = 0;
    int material2dIndex = 0;
    
    for (int cmdListIndex = 0; cmdListIndex < commands->CmdListsCount; cmdListIndex++) {
        const ImDrawList* cmds = commands->CmdLists[cmdListIndex];
        populateVertexData(bufferIndex,
                cmds->VtxBuffer.Size * sizeof(ImDrawVert), cmds->VtxBuffer.Data,
                cmds->IdxBuffer.Size * sizeof(ImDrawIdx), cmds->IdxBuffer.Data);
                
        for (const auto& pcmd : cmds->CmdBuffer) {
            if (pcmd.UserCallback) {
                pcmd.UserCallback(cmds, &pcmd);
            } else {
                auto texture = (Texture const*)pcmd.GetTexID();
                MaterialInstance* materialInstance;
                
                if (material2dIndex == mMaterial2dInstances.size()) {
                    if (mMaterial2d) {
                        mMaterial2dInstances.push_back(mMaterial2d->createInstance());
                    } else {
                        mMaterial2dInstances.push_back(nullptr);
                    }
                }
                materialInstance = mMaterial2dInstances[material2dIndex++];

                if (materialInstance) {
                    auto scissorLeft = static_cast<uint32_t>(std::max(0.0f, pcmd.ClipRect.x));
                    auto scissorBottom = static_cast<uint32_t>(std::max(0.0f,
                            mFlipVertical ? pcmd.ClipRect.y : (fbheight - pcmd.ClipRect.w)));
                    auto scissorWidth = static_cast<uint32_t>(pcmd.ClipRect.z - pcmd.ClipRect.x);
                    auto scissorHeight = static_cast<uint32_t>(pcmd.ClipRect.w - pcmd.ClipRect.y);
                    
                    materialInstance->setScissor(scissorLeft, scissorBottom, scissorWidth, scissorHeight);
                    
                    if (texture) {
                        TextureSampler sampler(MinFilter::LINEAR, MagFilter::LINEAR);
                        materialInstance->setParameter("albedo", texture, sampler);
                    } else {
                        materialInstance->setParameter("albedo", mTexture, mSampler);
                    }
                    
                    rbuilder
                            .geometry(primIndex, RenderableManager::PrimitiveType::TRIANGLES,
                                    mVertexBuffers[bufferIndex], mIndexBuffers[bufferIndex],
                                    pcmd.IdxOffset, pcmd.ElemCount)
                            .blendOrder(primIndex, primIndex)
                            .material(primIndex, materialInstance);
                    primIndex++;
                }
            }
        }
        bufferIndex++;
    }
    if (commands->CmdListsCount > 0) {
        rbuilder.build(*mEngine, mRenderable);
    }
}

void MetaCoreImGuiHelper::createVertexBuffer(size_t bufferIndex, size_t capacity) {
    syncThreads();
    if (bufferIndex < mVertexBuffers.size() && mVertexBuffers[bufferIndex]) {
        mEngine->destroy(mVertexBuffers[bufferIndex]);
    }
    
    mVertexBuffers[bufferIndex] = VertexBuffer::Builder()
            .vertexCount(capacity)
            .bufferCount(1)
            .attribute(VertexAttribute::POSITION, 0, VertexBuffer::AttributeType::FLOAT2, 0,
                    sizeof(ImDrawVert))
            .attribute(VertexAttribute::UV0, 0, VertexBuffer::AttributeType::FLOAT2,
                    sizeof(filament::math::float2), sizeof(ImDrawVert))
            .attribute(VertexAttribute::COLOR, 0, VertexBuffer::AttributeType::UBYTE4,
                    2 * sizeof(filament::math::float2), sizeof(ImDrawVert))
            .normalized(VertexAttribute::COLOR)
            .build(*mEngine);
}

void MetaCoreImGuiHelper::createIndexBuffer(size_t bufferIndex, size_t capacity) {
    syncThreads();
    if (bufferIndex < mIndexBuffers.size() && mIndexBuffers[bufferIndex]) {
        mEngine->destroy(mIndexBuffers[bufferIndex]);
    }
    
    mIndexBuffers[bufferIndex] = IndexBuffer::Builder()
            .indexCount(capacity)
            .bufferType(IndexBuffer::IndexType::USHORT)
            .build(*mEngine);
}

void MetaCoreImGuiHelper::createBuffers(int numRequiredBuffers) {
    if (numRequiredBuffers > mVertexBuffers.size()) {
        size_t previousSize = mVertexBuffers.size();
        mVertexBuffers.resize(numRequiredBuffers, nullptr);
        for (size_t i = previousSize; i < mVertexBuffers.size(); i++) {
            createVertexBuffer(i, 1000);
        }
    }
    if (numRequiredBuffers > mIndexBuffers.size()) {
        size_t previousSize = mIndexBuffers.size();
        mIndexBuffers.resize(numRequiredBuffers, nullptr);
        for (size_t i = previousSize; i < mIndexBuffers.size(); i++) {
            createIndexBuffer(i, 5000);
        }
    }
}

void MetaCoreImGuiHelper::populateVertexData(size_t bufferIndex, size_t vbSizeInBytes, void* vbImguiData,
        size_t ibSizeInBytes, void* ibImguiData)
{
    size_t requiredVertCount = vbSizeInBytes / sizeof(ImDrawVert);
    size_t capacityVertCount = mVertexBuffers[bufferIndex]->getVertexCount();
    if (requiredVertCount > capacityVertCount) {
        createVertexBuffer(bufferIndex, requiredVertCount);
    }
    size_t nVbBytes = requiredVertCount * sizeof(ImDrawVert);
    void* vbFilamentData = malloc(nVbBytes);
    memcpy(vbFilamentData, vbImguiData, nVbBytes);
    mVertexBuffers[bufferIndex]->setBufferAt(*mEngine, 0,
            VertexBuffer::BufferDescriptor(vbFilamentData, nVbBytes,
                [](void* buffer, size_t size, void* user) {
                    free(buffer);
                }, /* user = */ nullptr));

    size_t requiredIndexCount = ibSizeInBytes / 2;
    size_t capacityIndexCount = mIndexBuffers[bufferIndex]->getIndexCount();
    if (requiredIndexCount > capacityIndexCount) {
        createIndexBuffer(bufferIndex, requiredIndexCount);
    }
    size_t nIbBytes = requiredIndexCount * 2;
    void* ibFilamentData = malloc(nIbBytes);
    memcpy(ibFilamentData, ibImguiData, nIbBytes);
    mIndexBuffers[bufferIndex]->setBuffer(*mEngine,
            IndexBuffer::BufferDescriptor(ibFilamentData, nIbBytes,
                [](void* buffer, size_t size, void* user) {
                    free(buffer);
                }, /* user = */ nullptr));
}

void MetaCoreImGuiHelper::syncThreads() {
#if UTILS_HAS_THREADING
    if (!mHasSynced) {
        Fence::waitAndDestroy(mEngine->createFence());
        mHasSynced = true;
    }
#endif
}

} // namespace MetaCore