#include "../../engine/Shared.h"
#include "../Global.h"
#include "../Instance.h"
#include "../Log.h"
#include "../Mesh.h"
#include "Rendering.h"

#include "SDL3/SDL_events.h"
#include "backends/imgui_impl_sdl3.h"
#include "bgfx/defines.h"
#include "bx/bx.h"
#include "bx/constants.h"
#include "bx/filepath.h"
#include "bx/math.h"
#include "bx/timer.h"
#include <bgfx/bgfx.h>
#include <bimg/bimg.h>
#include <bx/file.h>
#include <cstdio>
#include <thread>

#include "imgui-helper.h"
#include "imgui.h"

using namespace std::chrono_literals;

namespace Game
{
    namespace
    {
        const bgfx::Memory* loadFile(const char* path, bool appendZero = false, int32_t retryCount = 1)
        {
            FILE* file;
            for (int32_t i = 0; i < retryCount; ++i)
            {
                file = fopen(path, "rb");
                if (file == nullptr && i < retryCount - 1)
                {
                    std::this_thread::sleep_for(100ms);
                    break;
                }

                fseek(file, 0, SEEK_END);
                long fileSize = ftell(file);
                fseek(file, 0, SEEK_SET);

                void* rawMem = AllocateScratch(fileSize + 1);
                const bgfx::Memory* mem = bgfx::makeRef(rawMem, fileSize + 1);
                fread(mem->data, 1, fileSize, file);
                if (appendZero)
                {
                    mem->data[mem->size - 1] = '\0';
                }
                fclose(file);

                return mem;
            }

            return nullptr;
        }

        bgfx::ShaderHandle loadShader(const char* FILENAME)
        {
            const char* shaderPath = "???";

            switch (bgfx::getRendererType())
            {
            case bgfx::RendererType::Agc:
            case bgfx::RendererType::Nvn:
            case bgfx::RendererType::Count:
            case bgfx::RendererType::Noop:
                break;
            case bgfx::RendererType::Direct3D11:
            case bgfx::RendererType::Direct3D12:
                shaderPath = "game/compiled-shaders/dx11/";
                break;
            case bgfx::RendererType::Gnm:
                shaderPath = "game/compiled-shaders/pssl/";
                break;
            case bgfx::RendererType::Metal:
                shaderPath = "game/compiled-shaders/metal/";
                break;
            case bgfx::RendererType::OpenGL:
                shaderPath = "game/compiled-shaders/glsl/";
                break;
            case bgfx::RendererType::OpenGLES:
                shaderPath = "game/compiled-shaders/essl/";
                break;
            case bgfx::RendererType::Vulkan:
                shaderPath = "game/compiled-shaders/spirv/";
                break;
            }

            char buffer[512]{0};
            bx::strCopy(buffer, sizeof(buffer), shaderPath);
            bx::strCat(buffer, sizeof(buffer), FILENAME);
            bx::strCat(buffer, sizeof(buffer), ".bin");

            LOG("Loading shader at %s", buffer);
            const bgfx::Memory* mem = loadFile(buffer, true, 3);

            if (mem == nullptr)
            {
                LOG_WARN("Failed to load shader %s", FILENAME);
                return {};
            }
            return bgfx::createShader(mem);
        }

        bgfx::TextureHandle loadTexture(const bx::FilePath& _filePath,
                                        uint64_t _flags,
                                        uint8_t _skip,
                                        bgfx::TextureInfo* _info,
                                        bimg::Orientation::Enum* _orientation)
        {
            BX_UNUSED(_skip);
            bgfx::TextureHandle handle = BGFX_INVALID_HANDLE;
            bx::Error err;

            const bgfx::Memory* data = loadFile(_filePath.getCPtr());
            if (data == nullptr)
            {
                LOG_WARN("Failed to find image %s", _filePath.getCPtr());
                return handle;
            }
            bimg::ImageContainer imageContainer;
            if (!bimg::imageParse(imageContainer, data->data, data->size, &err))
            {
                LOG_WARN("Failed to load image %s", _filePath.getCPtr());
                return handle;
            }
            if (NULL != _orientation)
            {
                *_orientation = imageContainer.m_orientation;
            }

            const bgfx::Memory* mem =
                bgfx::makeRef(data->data + imageContainer.m_offset, data->size - imageContainer.m_offset);

            if (NULL != _info)
            {
                bgfx::calcTextureSize(*_info,
                                      uint16_t(imageContainer.m_width),
                                      uint16_t(imageContainer.m_height),
                                      uint16_t(imageContainer.m_depth),
                                      imageContainer.m_cubeMap,
                                      1 < imageContainer.m_numMips,
                                      imageContainer.m_numLayers,
                                      bgfx::TextureFormat::Enum(imageContainer.m_format));
            }

            if (imageContainer.m_cubeMap)
            {
                handle = bgfx::createTextureCube(uint16_t(imageContainer.m_width),
                                                 1 < imageContainer.m_numMips,
                                                 imageContainer.m_numLayers,
                                                 bgfx::TextureFormat::Enum(imageContainer.m_format),
                                                 _flags,
                                                 mem);
            }
            else if (1 < imageContainer.m_depth)
            {
                handle = bgfx::createTexture3D(uint16_t(imageContainer.m_width),
                                               uint16_t(imageContainer.m_height),
                                               uint16_t(imageContainer.m_depth),
                                               1 < imageContainer.m_numMips,
                                               bgfx::TextureFormat::Enum(imageContainer.m_format),
                                               _flags,
                                               mem);
            }
            else if (bgfx::isTextureValid(0,
                                          false,
                                          imageContainer.m_numLayers,
                                          bgfx::TextureFormat::Enum(imageContainer.m_format),
                                          _flags))
            {
                handle = bgfx::createTexture2D(uint16_t(imageContainer.m_width),
                                               uint16_t(imageContainer.m_height),
                                               1 < imageContainer.m_numMips,
                                               imageContainer.m_numLayers,
                                               bgfx::TextureFormat::Enum(imageContainer.m_format),
                                               _flags,
                                               mem);
            }

            if (bgfx::isValid(handle))
            {
                const bx::StringView name(_filePath);
                bgfx::setName(handle, name.getPtr(), name.getLength());
            }

            return handle;
        }

        void DitherGen(DitherData& data, int32_t recursion)
        {
            data.Points[0] = {0.0f, 0.0f};
            data.Points[1] = {0.5f, 0.5f};
            data.Points[2] = {0.5f, 0.0f};
            data.Points[3] = {0.0f, 0.5f};
            data.PointCount = 4;
            for (int32_t i = 0; i < data.BrightnessBucketCount; ++i)
                data.BrightnessBuckets[i] = 0;

            for (int32_t recursionLevel = 0; recursionLevel < recursion; ++recursionLevel)
            {
                int32_t startCount = data.PointCount;
                float offset = bx::pow(0.5f, recursionLevel + 1);
                for (int32_t i = 0; i < 4; ++i)
                {
                    for (int32_t j = 0; j < startCount; ++j)
                    {
                        data.Points[data.PointCount] = data.Points[j] + data.Points[i] * offset;
                        data.PointCount++;
                    }
                }
            }

            uint64_t dotsPerSide = bx::round(bx::pow(2, recursion));
            data.DitherTexDepth = dotsPerSide * dotsPerSide;
            data.DitherTexWH = 16 * dotsPerSide;
            uint64_t texPixelCount = data.DitherTexWH * data.DitherTexWH * data.DitherTexDepth;

            if (BX_COUNTOF(DitherData::DitherTex) < texPixelCount)
            {
                LOG_ERROR("Too many pixels: %llu", texPixelCount);
                return;
            }

            float invRes = 1.0f / data.DitherTexWH;
            for (int32_t z = 0; z < data.DitherTexDepth; ++z)
            {
                int32_t dotCount = z + 1;
                float dotArea = 0.5f / dotCount;
                float dotRadius = bx::sqrt(dotArea / bx::kPi);

                int zOffset = z * data.DitherTexWH * data.DitherTexWH;
                for (int32_t y = 0; y < data.DitherTexWH; ++y)
                {
                    int32_t yOffset = y * data.DitherTexWH;
                    for (int32_t x = 0; x < data.DitherTexWH; ++x)
                    {
                        Vec2 point{(x + 0.5f) * invRes, (y + 0.5f) * invRes};
                        float dist = bx::kFloatInfinity;
                        for (int32_t i = 0; i < dotCount; ++i)
                        {
                            Vec2 vec = point - data.Points[i];
                            Vec2 wrappedVec{bx::mod(vec.x + 0.5f, 1.0f) - 0.5f, bx::mod(vec.y + 0.5f, 1.0f) - 0.5f};
                            float curDist = wrappedVec.Magnitude();
                            dist = bx::min(dist, curDist);
                        }

                        dist = dist / (dotRadius * 2.4f);
                        float val = bx::clamp(1.0f - dist, 0.0f, 1.0f);
                        data.DitherTex[x + yOffset + zOffset] = Vec4{val, val, val, 1.0f};
                        // data.DitherTex[x + yOffset + zOffset] = Vec4{1.0, 0.0f, 0.0f, 1.0f};
                        int32_t bucket = bx::clamp(uint32_t(val * DitherData::BrightnessBucketCount),
                                                   0,
                                                   DitherData::BrightnessBucketCount - 1);
                        data.BrightnessBuckets[bucket] += 1;
                    }
                }
            }

            // Brightness ramp
            int32_t sum = 0;
            for (int32_t i = 0; i < data.BrightnessBucketCount; ++i)
            {
                sum += data.BrightnessBuckets[data.BrightnessBucketCount - 1 - i];
                data.BrightnessRamp[i + 1] = sum / (float)texPixelCount;
            }

            // Upload textures
            if (isValid(data.PreviewTex))
            {
                bgfx::destroy(data.PreviewTex);
                data.PreviewTex = BGFX_INVALID_HANDLE;
            }
            if (isValid(data.FinalTex))
            {
                bgfx::destroy(data.FinalTex);
                data.FinalTex = BGFX_INVALID_HANDLE;
            }
            if (isValid(data.RampTex))
            {
                bgfx::destroy(data.RampTex);
                data.RampTex = BGFX_INVALID_HANDLE;
            }
            const bgfx::Memory* memPreview = bgfx::makeRef(data.DitherTex, texPixelCount * sizeof(Vec4));
            const bgfx::Memory* memFinal = bgfx::makeRef(data.DitherTex, texPixelCount * sizeof(Vec4));
            const bgfx::Memory* memRamp = bgfx::makeRef(data.BrightnessRamp, sizeof(data.BrightnessRamp));
            data.PreviewTex = bgfx::createTexture2D(data.DitherTexWH,
                                                    data.DitherTexWH * data.DitherTexDepth,
                                                    false,
                                                    false,
                                                    bgfx::TextureFormat::RGBA32F,
                                                    0,
                                                    memPreview);
            data.FinalTex = bgfx::createTexture3D(data.DitherTexWH,
                                                  data.DitherTexWH,
                                                  data.DitherTexDepth,
                                                  false,
                                                  bgfx::TextureFormat::RGBA32F,
                                                  0,
                                                  memFinal);
            data.RampTex = bgfx::createTexture2D(
                BX_COUNTOF(data.BrightnessRamp), 1, false, 1, bgfx::TextureFormat::R32F, 0, memRamp);

            if (!isValid(data.DitherSampler))
            {
                data.DitherSampler = bgfx::createUniform("s_ditherSampler", bgfx::UniformType::Sampler);
            }
            if (!isValid(data.RampSampler))
            {
                data.RampSampler = bgfx::createUniform("s_rampSampler", bgfx::UniformType::Sampler);
            }
        }

        GameRendering* Instance = nullptr;
    } // namespace

    GameRendering& GameRendering::Get()
    {
        assert(Instance != nullptr);
        return *Instance;
    }

    void GameRendering::Setup()
    {
        LOG("--- RENDERING STARTUP ---");
        if (Instance != nullptr) LOG_WARN("old rendering wasn't destroyed!");
        Instance = this;
        SharedData& shared = GetShared();

        bgfx::Init init;
        init.type = bgfx::RendererType::Direct3D12;
        init.debug = true;
        init.platformData.nwh = shared.Window.Handle;
        init.platformData.ndt = nullptr;
        init.platformData.type = bgfx::NativeWindowHandleType::Default;
        init.resolution.width = shared.Window.WindowWidth;
        init.resolution.height = shared.Window.WindowHeight;
        init.resolution.reset = ResetFlags;

        LastWidth = shared.Window.WindowWidth;
        LastHeight = shared.Window.WindowHeight;

        LOG("%i by %i", init.resolution.width, init.resolution.height);

        if (!bgfx::init(init))
        {
            LOG_ERROR("BGFX setup failed!");
        }
        else
        {
            LOG("BGFX setup succeded!");
        }
        // bgfx::setDebug(BGFX_DEBUG_TEXT);
        bgfx::setViewClear(MainViewID, BGFX_CLEAR_COLOR | BGFX_CLEAR_DEPTH, 0x3399FFff, 1.0f, 0);
        bgfx::setViewRect(MainViewID, 0, 0, shared.Window.WindowWidth, shared.Window.WindowHeight);

        DefaultSampler = bgfx::createUniform("s_texColor", bgfx::UniformType::Sampler);
        Textures[0].Handle = loadTexture(
            bx::FilePath{"models/body.dds"}, BGFX_TEXTURE_NONE | BGFX_SAMPLER_NONE, 0, &Textures[0].Info, nullptr);
        Textures[0].SamplerHandle = DefaultSampler;
        LoadModels(Models, ModelCount);

        Materials[0] =
            Material::LoadFromShader("vert", "frag", MainViewID, Textures[0].Handle, Textures[0].SamplerHandle);
        imguiCreate();

        if (!ImGui_ImplSDL3_InitForOther(shared.Window.SDLWindow))
        {
            LOG_ERROR("Failed to set up imgui implementation!");
            return;
        }

        // ImGui::GetIO().BackendFlags |= ImGuiBackendFlags_RendererHasViewports;
        ImGui::GetIO().ConfigFlags |= ImGuiConfigFlags_DockingEnable;
        // ImGui::GetIO().ConfigFlags |= ImGuiConfigFlags_ViewportsEnable;
        // auto& platIO = ImGui::GetPlatformIO();
        // platIO.Renderer_CreateWindow = TODO;
        // platIO.Platform_DestroyWindow = TODO;
        // platIO.Platform_SetWindowSize = TODO;
        // platIO.Platform_RenderWindow = TODO;

        GameInstance& inst = GetInstance();
        if (!inst.IsInitialized)
        {
            inst.Time.StartTime = bx::getHPCounter();
        }

        if (inst.DebugData.ImguiIniSize > 0)
        {
            ImGui::LoadIniSettingsFromMemory(inst.DebugData.ImguiIni, inst.DebugData.ImguiIniSize);
        }
        else
        {
            ImGui::LoadIniSettingsFromDisk("imgui.ini");
        }
        DitherGen(DitherTextures, DitherRecursion);
    }

    void GameRendering::Update()
    {
        SharedData& shared = GetShared();

        // Handle events
        for (uint16_t i = 0; i < shared.Window.SDLEventCount; ++i)
        {
            ImGui_ImplSDL3_ProcessEvent(&shared.Window.SDLEvents[i]);
        }
        shared.Window.SDLEventCount = 0;

        // Resize if necessary
        if (shared.Window.WindowWidth != LastWidth || shared.Window.WindowHeight != LastHeight)
        {
            bgfx::reset(shared.Window.WindowWidth, shared.Window.WindowHeight, ResetFlags);
            bgfx::setViewRect(MainViewID, 0, 0, shared.Window.WindowWidth, shared.Window.WindowHeight);

            LastWidth = shared.Window.WindowWidth;
            LastHeight = shared.Window.WindowHeight;
        }

        // Reload shaders if necessary
        FileChangeNotification* shaderChange = nullptr;
        if (shared.Dev.ChangedShaderCount > 0)
        {
            shared.Dev.ChangedShaderCount = 0;

            // TODO: when to destroy shader?
            // TODO: only reload changed shaders
            bgfx::ShaderHandle vertexShader = loadShader("vert");
            bgfx::ShaderHandle fragmentShader = loadShader("frag");
            if (isValid(vertexShader) && isValid(fragmentShader))
            {
                bgfx::ProgramHandle newProgram = bgfx::createProgram(vertexShader, fragmentShader, true);
                if (isValid(newProgram))
                {
                    Materials[0].Shader = newProgram;
                    LastShaderLoadTime = GetInstance().Time.Now;
                }
                else
                {
                    LOG_WARN("Failed to load shader!");
                    LastShaderLoadTime = -1.0f;
                }
            }
        }

        // Start Rendering
        imguiBeginFrame(20);
        ImGui_ImplSDL3_NewFrame();
        ImGui::DockSpaceOverViewport(0, 0, ImGuiDockNodeFlags_PassthruCentralNode);

        auto& level = GetInstance().GameLevel;
        auto& debug = GetInstance().DebugData;
        if (UIVisible == UIVisibilityState::Debug)
        {
            if (ImGui::Begin("Rendering"))
            {
                if (LastShaderLoadTime >= 0.0f)
                {
                    ImGui::TextColored({0.2f, 0.9f, 0.2f, 1.0f},
                                       "Shader loaded %.0f seconds ago",
                                       GetInstance().Time.Now - LastShaderLoadTime);
                }
                else
                {
                    ImGui::TextColored({0.9f, 0.2f, 0.2f, 1.0f}, "Shader load Failiure!");
                }
                if (ImGui::Button("Reload Meshes"))
                {
                    LoadModels(Models, ModelCount);
                }
                ImGui::SameLine();
                if (ImGui::Button("Reload Level"))
                {
                    auto& lvl = GetInstance().GameLevel;
                    lvl = {};
                    lvl.Setup(shared.Game);
                }
                if (ImGui::Button("Dithergen"))
                {
                    DitherGen(DitherTextures, DitherRecursion);
                }
                ImGui::SameLine();
                ImGui::SliderInt("Recursion", &DitherRecursion, 1, 4);
                ImGui::Text(
                    "%ux%ux%u", DitherTextures.DitherTexWH, DitherTextures.DitherTexWH, DitherTextures.DitherTexDepth);
                if (!isValid(DitherTextures.PreviewTex))
                {
                    ImGui::Text("Invalid Texture");
                }
                else
                {
                    ImGui::Image(DitherTextures.PreviewTex.idx,
                                 {(float)DitherTextures.DitherTexWH,
                                  (float)DitherTextures.DitherTexWH * DitherTextures.DitherTexDepth});
                }
                if (isValid(DitherTextures.RampTex))
                {
                    ImGui::Image(DitherTextures.RampTex.idx, {BX_COUNTOF(DitherTextures.BrightnessRamp), 8});
                }
                Vec3 quadPos = level.UIQuads.Get({0}).EData.Transform.GetPosition();
                ImGui::Text("%f %f %f", quadPos.x, quadPos.y, quadPos.z);
                ImGui::Text("Shader log:");
                ImGui::TextWrapped("%s", GetShared().Dev.ShaderLog);
            }
            ImGui::End();
            if (ImGui::Begin("Puzzles"))
            {
                ImGui::Text("List");
                for (int32_t i = 0; i < BX_COUNTOF(level.Puzzles); ++i)
                {
                    auto& puzzleData = level.Puzzles[i].Data;
                    bool isSelected = debug.SelectedDebugLevel == i;
                    ImGui::PushID("selectable");
                    if (ImGui::Selectable(puzzleData.PuzzleName, isSelected))
                    {
                        debug.SelectedDebugLevel = isSelected ? UINT16_MAX : i;
                    }
                    ImGui::PopID();
                    if (isSelected)
                    {
                        ImGui::PushID("edit field");
                        ImGui::InputText("", puzzleData.PuzzleName, sizeof(Puzzle::PuzzleData::PuzzleName));
                        ImGui::PopID();

                        if (!puzzleData.RenderDebugUI())
                        {
                            debug.SelectedDebugLevel = UINT16_MAX;
                        }
                    }
                }
            }
            ImGui::End();
        }

        GetInstance().GameLevel.Update();
        GetInstance().GameLevel.Render(MainViewID, Models, Materials);

        // bgfx::dbgTextPrintf(1, 1, 0x0F, "Time: %.1fs", GetInstance().Time.Now);
        // for (int32_t i = 0; i < (int32_t)PerfCounterType::COUNT; ++i)
        // {
        //     bgfx::dbgTextPrintf(
        //         1, 2 + i, 0x0F, "%s Max: %.3fs", PerfCounterNames[i], shared.Window.PerfCounters[i].GetMax());
        // }

        // Finish Frame
        imguiEndFrame();

        if (ImGui::GetIO().ConfigFlags & ImGuiConfigFlags_ViewportsEnable)
        {
            ImGui::UpdatePlatformWindows();
            ImGui::RenderPlatformWindowsDefault();
        }

        START_PERF();
        bgfx::frame();
        END_PERF(shared.Window.PerfCounters, PerfCounterType::Submit, shared.Window.FrameCounter);
    }

    void GameRendering::Shutdown()
    {
        LOG("--- RENDERING_SHUTDOWN ---");
        ImGui::SaveIniSettingsToDisk("imgui.ini");
        auto& debug = GetInstance().DebugData;
        const char* iniData = ImGui::SaveIniSettingsToMemory(reinterpret_cast<uint64_t*>(&debug.ImguiIniSize));
        assert(debug.ImguiIniSize <= BX_COUNTOF(InstanceDebugData::ImguiIni));
        bx::memCopy(debug.ImguiIni, iniData, bx::min(debug.ImguiIniSize, BX_COUNTOF(InstanceDebugData::ImguiIni)));
        ImGui_ImplSDL3_Shutdown();
        imguiDestroy();
        bgfx::shutdown();
        Instance = nullptr;
    }

    Material Material::LoadFromShader(
        const char* vertPath, const char* fragPath, uint16_t view, bgfx::TextureHandle tex, bgfx::UniformHandle sampler)
    {
        BX_ASSERT(vertPath != nullptr && fragPath != nullptr, "Invalid shader path!");
        bgfx::ShaderHandle vertexShader = loadShader(vertPath);
        bgfx::ShaderHandle fragmentShader = loadShader(fragPath);

        Material mat;
        mat.Shader = bgfx::createProgram(vertexShader, fragmentShader, true);
        mat.State = 0 | BGFX_STATE_WRITE_RGB | BGFX_STATE_WRITE_A | BGFX_STATE_WRITE_Z | BGFX_STATE_DEPTH_TEST_LESS |
                    BGFX_STATE_CULL_CCW | BGFX_STATE_MSAA;
        mat.Uniforms[Material::UTime] = bgfx::createUniform("u_time", bgfx::UniformType::Vec4);
        mat.Uniforms[Material::UDotColor] = bgfx::createUniform("u_testColor", bgfx::UniformType::Vec4);
        mat.Uniforms[Material::UTexInfo] = bgfx::createUniform("u_texInfo", bgfx::UniformType::Vec4);
        mat.Textures[0].Handle = tex;
        mat.Textures[0].SamplerHandle = sampler;
        mat.ViewID = view;
        return mat;
    }

    uint16_t GameRendering::GetModelHandleFromPath(const char* path)
    {
        uint32_t AssetHandle = CrcPath(path);
        for (int32_t i = 0; i < ModelCount; ++i)
        {
            if (Models[i].AssetHandle == AssetHandle)
            {
                return i;
            }
        }
        return 0;
    }

} // namespace Game