#include "Primitives.h"
#include "../Assets.h"
#include "../GPU.h"
#include "../GPUBuffer.h"
#include "../PrimitivesBuilder/PipelineBuilder.h"
#include "../Shader.h"
#include "../Texture.h"
#include "Engine/Scale.h"
#include <SDL3/SDL_gpu.h>
#include <algorithm>
#include <array>
#include <cmath>
#include <vector>

namespace UI {


// TODO: clean up luna-5.6-not-reading-the-c++-style-guide-skill-slop
static Point translation;
static Rect clip;
static bool clipEnabled = false;
static GraphicsPipeline* pipeline = nullptr;
static SDL_GPUSampler* sampler = nullptr;
static Texture* fallbackTexture = nullptr;
static GPUBuffer* instanceBuffer = nullptr;
static std::vector<GPUBuffer*> obsoleteInstanceBuffers;
static size_t instanceCapacity = 0;
static std::vector<float> instances;

struct ClipState {
    Rect clip;
    bool enabled;
};

struct Batch {
    SDL_GPUTexture* texture;
    size_t instanceStart;
    size_t instanceCount;
    ClipState clipState;
};

static std::vector<Batch> batches;
static std::array<std::vector<SDL_GPUTransferBuffer*>, 3> transferBuffers;
static size_t transferFrame = 0;

constexpr size_t FloatsPerInstance = 16;

// pipeline + sampler + 1x1 white fallback texture
void EnsureResources() {
    if (pipeline) {
        return;
    }

    Shader vertShader(Assets::Path("shaders/UI.vert.spirv"), "main", 0, 0, 1, 1);
    Shader fragShader(Assets::Path("shaders/UI.frag.spirv"), "main", 1, 0, 0, 0);

    pipeline = PipelineBuilder()
        .FormatSwapchain()
        .Shaders(&vertShader, &fragShader)
        .BlendPremultiplied()
        .Primitive(SDL_GPU_PRIMITIVETYPE_TRIANGLESTRIP)
        .Build();

    SDL_GPUSamplerCreateInfo samplerInfo = {};
    samplerInfo.min_filter = SDL_GPU_FILTER_LINEAR;
    samplerInfo.mag_filter = SDL_GPU_FILTER_LINEAR;
    samplerInfo.mipmap_mode = SDL_GPU_SAMPLERMIPMAPMODE_LINEAR;
    samplerInfo.address_mode_u = SDL_GPU_SAMPLERADDRESSMODE_CLAMP_TO_EDGE;
    samplerInfo.address_mode_v = SDL_GPU_SAMPLERADDRESSMODE_CLAMP_TO_EDGE;
    samplerInfo.address_mode_w = SDL_GPU_SAMPLERADDRESSMODE_CLAMP_TO_EDGE;
    sampler = SDL_CreateGPUSampler(GPU::GetDevice(), &samplerInfo);

    SDL_GPUTextureCreateInfo textureInfo = {};
    textureInfo.type = SDL_GPU_TEXTURETYPE_2D;
    textureInfo.format = SDL_GPU_TEXTUREFORMAT_R8G8B8A8_UNORM;
    textureInfo.width = 1;
    textureInfo.height = 1;
    textureInfo.layer_count_or_depth = 1;
    textureInfo.num_levels = 1;
    textureInfo.usage = SDL_GPU_TEXTUREUSAGE_SAMPLER;
    fallbackTexture = new Texture(textureInfo);
    unsigned char whitePixel[] = { 255, 255, 255, 255 };
    fallbackTexture->Upload(whitePixel);
}

// if we need to increase capacity create a new buffer and obsolete the old one
void UpdateInstanceCapacity(size_t requiredInstances) {
    if (requiredInstances <= instanceCapacity) {
        return;
    }

    size_t newCapacity = instanceCapacity == 0 ? 1 : instanceCapacity;
    while (newCapacity < requiredInstances) {
        newCapacity *= 2;
    }

    GPUBuffer* newBuffer = new GPUBuffer(
        SDL_GPU_BUFFERUSAGE_GRAPHICS_STORAGE_READ,
        newCapacity * FloatsPerInstance * sizeof(float)
    );
    if (instanceBuffer) {
        obsoleteInstanceBuffers.push_back(instanceBuffer);
    }
    instanceBuffer = newBuffer;
    instanceCapacity = newCapacity;
}

void ReleaseObsoleteBuffers() {
    for (GPUBuffer* buffer : obsoleteInstanceBuffers) {
        delete buffer;
    }
    obsoleteInstanceBuffers.clear();
}

void AppendInstance(
    UI::Rect bounds,
    UI::Color color,
    UI::Color borderColor,
    float borderWidth,
    float roundness,
    float mode,
    SDL_GPUTexture* texture
) {
    EnsureResources();
    const bool startsNewBatch = batches.empty()
        || batches.back().texture != texture
        || batches.back().clipState.enabled != clipEnabled
        || (clipEnabled && batches.back().clipState.clip != clip);
    if (startsNewBatch) {
        batches.push_back({
            texture,
            instances.size() / FloatsPerInstance,
            0,
            {clip, clipEnabled}
        });
    }

    // logical units -> physical pixels, snapped onto the pixel grid so adjacent
    // elements share exact edges instead of sampling blurry half-pixels
    const float scale = GetScale();
    const float left = std::round((bounds.X + translation.x) * scale);
    const float top = std::round((bounds.Y + translation.y) * scale);
    const float right = std::round((bounds.X + translation.x + bounds.Width) * scale);
    const float bottom = std::round((bounds.Y + translation.y + bounds.Height) * scale);

    instances.insert(instances.end(), {
        left, top, right - left, bottom - top,
        color.r, color.g, color.b, color.a,
        borderColor.r, borderColor.g, borderColor.b, borderColor.a,
        borderWidth * scale, roundness * scale, mode, 0.0f
    });
    batches.back().instanceCount++;
}

void UploadInstances() {
    if (instances.empty()) {
        return;
    }

    EnsureResources();
    int width = 0;
    int height = 0;
    SDL_GetWindowSizeInPixels(GPU::GetWindow(), &width, &height);
    const Uint32 uploadSize = static_cast<Uint32>(instances.size() * sizeof(float));
    UpdateInstanceCapacity(instances.size() / FloatsPerInstance);

    SDL_GPUTransferBufferCreateInfo transferInfo = {};
    transferInfo.usage = SDL_GPU_TRANSFERBUFFERUSAGE_UPLOAD;
    transferInfo.size = uploadSize;
    SDL_GPUTransferBuffer* transfer = SDL_CreateGPUTransferBuffer(GPU::GetDevice(), &transferInfo);
    if (!transfer) {
        SDL_Log("Failed to create UI transfer buffer: %s", SDL_GetError());
        return;
    }

    void* mapped = SDL_MapGPUTransferBuffer(GPU::GetDevice(), transfer, false);
    if (!mapped) {
        SDL_Log("Failed to map UI transfer buffer: %s", SDL_GetError());
        SDL_ReleaseGPUTransferBuffer(GPU::GetDevice(), transfer);
        return;
    }
    memcpy(mapped, instances.data(), uploadSize);
    SDL_UnmapGPUTransferBuffer(GPU::GetDevice(), transfer);
    transferBuffers[transferFrame].push_back(transfer);

    GPU::BeginCopying();
    SDL_GPUTransferBufferLocation source = { .transfer_buffer = transfer, .offset = 0 };
    SDL_GPUBufferRegion destination = {
        .buffer = instanceBuffer->Get(),
        .offset = 0,
        .size = uploadSize
    };
    SDL_UploadToGPUBuffer(GPU::GetCopyPass(), &source, &destination, false);
    GPU::EndCopying();

    SDL_GPUColorTargetInfo colorTargetInfo = {};
    colorTargetInfo.texture = GPU::GetSwapchain();
    colorTargetInfo.load_op = SDL_GPU_LOADOP_LOAD;
    colorTargetInfo.store_op = SDL_GPU_STOREOP_STORE;
    colorTargetInfo.mip_level = 0;
    colorTargetInfo.layer_or_depth_plane = 0;
    colorTargetInfo.cycle = false;

    GPU::BeginRendering({ &colorTargetInfo }, nullptr);

    SDL_GPUViewport viewport = {};
    viewport.w = static_cast<float>(width);
    viewport.h = static_cast<float>(height);
    SDL_SetGPUViewport(GPU::GetRenderPass(), &viewport);
    SDL_GPUBuffer* buffer = instanceBuffer->Get();
    SDL_BindGPUVertexStorageBuffers(GPU::GetRenderPass(), 0, &buffer, 1);

    pipeline->Bind();
    for (const Batch& batch : batches) {
        clip = batch.clipState.clip;
        clipEnabled = batch.clipState.enabled;
        Primitives::ApplyRenderState();

        SDL_GPUTextureSamplerBinding samplerBinding = {
            .texture = batch.texture ? batch.texture : fallbackTexture->Get(),
            .sampler = sampler
        };
        SDL_BindGPUFragmentSamplers(GPU::GetRenderPass(), 0, &samplerBinding, 1);
        float pushConstants[4] = {
            static_cast<float>(width),
            static_cast<float>(height),
            static_cast<float>(batch.instanceStart),
            0.0f
        };
        SDL_PushGPUVertexUniformData(GPU::GetCommandBuffer(), 0, pushConstants, sizeof(pushConstants));
        SDL_DrawGPUPrimitives(
            GPU::GetRenderPass(),
            4,
            static_cast<Uint32>(batch.instanceCount),
            0,
            0
        );
    }

    instances.clear();
    batches.clear();
}

static std::vector<ClipState> clipStack;

void Primitives::BeginFrame() {
    translation = {};
    clip = {};
    clipEnabled = false;
    clipStack.clear();

    ReleaseObsoleteBuffers();
    transferFrame = (transferFrame + 1) % transferBuffers.size();
    for (SDL_GPUTransferBuffer* transfer : transferBuffers[transferFrame]) {
        SDL_ReleaseGPUTransferBuffer(GPU::GetDevice(), transfer);
    }
    transferBuffers[transferFrame].clear();
}

void Primitives::EndFrame() {
    UploadInstances();
    GPU::EndRendering();
}


void Primitives::SubmitRectangle(Rect bounds, Color color, Color borderColor, float borderWidth, float roundness) {
    EnsureResources();
    AppendInstance(bounds, color, borderColor, borderWidth, roundness, 0.0f, fallbackTexture->Get());
}

void Primitives::SubmitImage(Texture* texture, Rect bounds, Color color) {
    if (!texture) {
        return;
    }
    AppendInstance(bounds, color, Color(), 0.0f, 0.0f, 1.0f, texture->Get());
}

void Primitives::SubmitText(Texture* texture, Rect bounds, Color color) {
    if (!texture) {
        return;
    }
    AppendInstance(bounds, color, Color(), 0.0f, 0.0f, 2.0f, texture->Get());
}


void Primitives::SetTranslation(Point newTranslation) {
    translation = newTranslation;
}

Point Primitives::GetTranslation() {
    return translation;
}

void Primitives::SetClip(Rect newClip) {
    clip = newClip;
    clipEnabled = true;
}

void Primitives::ClearClip() {
    clipEnabled = false;
}

void Primitives::PushClip(Rect newClip) {
    clipStack.push_back({clip, clipEnabled});
    clip = clipEnabled ? clip.Intersect(newClip) : newClip;
    clipEnabled = true;
}

void Primitives::PopClip() {
    if (clipStack.empty()) {
        clipEnabled = false;
        return;
    }

    const ClipState previous = clipStack.back();
    clipStack.pop_back();
    clip = previous.clip;
    clipEnabled = previous.enabled;
}

void Primitives::ApplyRenderState() {
    int width = 0;
    int height = 0;
    SDL_GetWindowSizeInPixels(GPU::GetWindow(), &width, &height);

    SDL_Rect scissor{0, 0, width, height};
    if (clipEnabled) {
        const float scale = GetScale();
        const int left = static_cast<int>(std::round(clip.X * scale));
        const int top = static_cast<int>(std::round(clip.Y * scale));
        const int right = static_cast<int>(std::round((clip.X + clip.Width) * scale));
        const int bottom = static_cast<int>(std::round((clip.Y + clip.Height) * scale));
        scissor.x = left;
        scissor.y = top;
        scissor.w = std::max(0, right - left);
        scissor.h = std::max(0, bottom - top);
    }

    SDL_SetGPUScissor(GPU::GetRenderPass(), &scissor);
}

void Primitives::ClearResources() {
    if (pipeline) {
        delete pipeline;
        pipeline = nullptr;
    }
}

}