#include "StrokeRenderingContext.h"
#include "../Assets.h"
#include "../GPU.h"
#include "../PrimitivesBuilder/PipelineBuilder.h"
#include "../PrimitivesBuilder/VertexInputStateBuilder.h"
#include "../GPUBuffer.h"


GraphicsPipeline* StrokeRenderingContext::pipelineStamp = nullptr;
GraphicsPipeline* StrokeRenderingContext::pipelineBlend = nullptr;
GraphicsPipeline* StrokeRenderingContext::pipelineEraseBlend = nullptr;
Shader* StrokeRenderingContext::vertShader = nullptr;
Shader* StrokeRenderingContext::fragShader = nullptr;
SDL_GPUSampler* StrokeRenderingContext::sampler = nullptr;

StrokeRenderingContext::StrokeRenderingContext(Brush* brush, int width, int height, float* initialColor, float* initialSize)
{
    assert(brush); // brush should not be null
    this->brush = brush;
    this->boundingBox = { 0.0f, 0.0f, -1.0f, -1.0f }; // no normal stamp should have exactly -1.0 size, hopefully
    this->width = width;
    this->height = height;
    for (int i = 0; i < 3; i++) {
        this->color[i] = initialColor[i];
    }
    this->color[3] = brush->Flow;
    this->opacity = brush->Opacity;
    brush->Size = *initialSize;
    if (!pipelineStamp || !pipelineBlend) {
        this->createPipeline();
    }
    this->offsetX = 0;
    this->offsetY = 0;
}

void StrokeRenderingContext::createTexture()
{
    SDL_GPUTextureCreateInfo createInfo = {};
    createInfo.width = width;
    createInfo.height = height;
    createInfo.format = SDL_GPU_TEXTUREFORMAT_R8G8B8A8_UNORM;
    createInfo.usage = SDL_GPU_TEXTUREUSAGE_COLOR_TARGET | SDL_GPU_TEXTUREUSAGE_SAMPLER;
    createInfo.layer_count_or_depth = 1;
    createInfo.num_levels = 1;
    createInfo.sample_count = SDL_GPU_SAMPLECOUNT_1;
    strokeTexture = new Texture(createInfo);
    // clear
    strokeTexture->Clear({ color[0], color[1], color[2], 0 });
}

void StrokeRenderingContext::createPipeline()
{
    Shader* vertShader = new Shader(Assets::Path("shaders/StrokeRenderingContext.vert.spirv").c_str(), "main", 0, 0, 0, 1);
    Shader* fragShader = new Shader(Assets::Path("shaders/StrokeRenderingContext.frag.spirv").c_str(), "main", 1, 0, 0, 0);

    auto vertexInputState = VertexInputStateBuilder()
        .AddAttribute(0, SDL_GPU_VERTEXELEMENTFORMAT_FLOAT4) // rect
        .AddAttribute(1, SDL_GPU_VERTEXELEMENTFORMAT_FLOAT4) // uv rect
        .AddAttribute(2, SDL_GPU_VERTEXELEMENTFORMAT_FLOAT) // opacity
        .AddAttribute(3, SDL_GPU_VERTEXELEMENTFORMAT_FLOAT) // rotation
        .Instanced()
        .BuildDeferred();

    auto genericPipeline = PipelineBuilder()
        .Shaders(vertShader, fragShader)
        .Format(SDL_GPU_TEXTUREFORMAT_R8G8B8A8_UNORM)
        .Vertex(vertexInputState.Use())
        .Primitive(SDL_GPU_PRIMITIVETYPE_TRIANGLELIST);

    pipelineStamp = genericPipeline
        .BlendAlphaKeepColor()
        .Build();

    pipelineBlend = genericPipeline
        .BlendAlpha()
        .Build();

    pipelineEraseBlend = genericPipeline
        .BlendErase()
        .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;
    samplerInfo.mip_lod_bias = 0.0f;
    samplerInfo.max_anisotropy = 1.0f;
    samplerInfo.compare_op = SDL_GPU_COMPAREOP_NEVER;
    samplerInfo.min_lod = 0.0f;
    samplerInfo.max_lod = 1000.0f;
    sampler = SDL_CreateGPUSampler(GPU::GetDevice(), &samplerInfo);

}

void StrokeRenderingContext::SubmitInputs(std::list<PointerEvent> events)
{
    inputEvents.insert(inputEvents.end(), events.begin(), events.end());
    immutableInputEvents.insert(immutableInputEvents.end(), events.begin(), events.end());
}

Texture* StrokeRenderingContext::RenderQueued(SDL_FRect* outputRectBoundingBox)
{
    if (this->minimized) {
        return strokeTexture;
    }
    if (!this->strokeTexture) {
        this->createTexture();
    }
    std::list<BrushStamp> stamps = brush->CalculateStampList(inputEvents);
    std::vector<GPUStampInstance> instances;
    while (!stamps.empty()) {
        BrushStamp currentStamp = stamps.front();
        stamps.pop_front();
        SDL_FRect stampRect = {
            (float)(currentStamp.X - (currentStamp.Size / 2.0)),
            (float)(currentStamp.Y - (currentStamp.Size / 2.0)),
            (float)currentStamp.Size,
            (float)currentStamp.Size
        };
        // expand bounding box
        if (this->boundingBox.w == -1.0f && this->boundingBox.h == -1.0f) {
            this->boundingBox = stampRect;
        }

        float minX = SDL_min(this->boundingBox.x, stampRect.x - 1);
        float minY = SDL_min(this->boundingBox.y, stampRect.y - 1);
        float maxX = SDL_max(this->boundingBox.x + this->boundingBox.w, stampRect.x + stampRect.w + 2);
        float maxY = SDL_max(this->boundingBox.y + this->boundingBox.h, stampRect.y + stampRect.h + 2);
        // limit to texture bounds
        this->boundingBox.x = SDL_max(minX, 0.0f);
        this->boundingBox.y = SDL_max(minY, 0.0f);
        this->boundingBox.w = SDL_min(maxX, (float)width) - this->boundingBox.x;
        this->boundingBox.h = SDL_min(maxY, (float)height) - this->boundingBox.y;

        // build 
        instances.push_back({
            .x = stampRect.x,
            .y = stampRect.y,
            .w = stampRect.w,
            .h = stampRect.h,
            .u0 = 0.0f,
            .v0 = 0.0f,
            .u1 = 1.0f,
            .v1 = 1.0f,
            .pressure = (float)color[3],
            .rotation = (float)currentStamp.Rotation
        });
    }

    if (!instances.empty() and boundingBox.w != -1.0f and boundingBox.h != -1.0f) {
        upload(instances, this->strokeTexture, brush->Stamp, pipelineStamp);
    }

    // copy bounding box to output
    *outputRectBoundingBox = this->boundingBox;
    return strokeTexture;
}

void StrokeRenderingContext::RenderResultOnto(Texture* target, bool erase)
{
    std::vector<GPUStampInstance> instances;
    if (boundingBox.w == -1.0f && boundingBox.h == -1.0f) {
        instances.push_back({
            .x = 0.0f,
            .y = 0.0f,
            .w = (float)target->GetWidth(),
            .h = (float)target->GetHeight(),
            .u0 = 0.0f,
            .v0 = 0.0f,
            .u1 = 1.0f,
            .v1 = 1.0f,
            .pressure = opacity,
            .rotation = 0.0f
        });
    }
    else {
        instances.push_back({
            .x = boundingBox.x,
            .y = boundingBox.y,
            .w = boundingBox.w,
            .h = boundingBox.h,
            .u0 = minimized ? 0.0f : boundingBox.x / (float)width,
            .v0 = minimized ? 0.0f : boundingBox.y / (float)height,
            .u1 = minimized ? 1.0f : (boundingBox.x + boundingBox.w) / (float)width,
            .v1 = minimized ? 1.0f : (boundingBox.y + boundingBox.h) / (float)height,
            .pressure = opacity,
            .rotation = 0.0f
        });
    }
    upload(instances, target, this->strokeTexture, erase ? pipelineEraseBlend : pipelineBlend); // lol
}


void StrokeRenderingContext::upload(std::vector<GPUStampInstance> instances, Texture* target, Texture* source, GraphicsPipeline* pipeline)
{
    SDL_GPUBufferCreateInfo vertexBufferInfo = {};
    vertexBufferInfo.size = instances.size() * sizeof(GPUStampInstance);
    vertexBufferInfo.usage = SDL_GPU_BUFFERUSAGE_VERTEX;

    GPUBuffer* vertexBuffer = new GPUBuffer(&vertexBufferInfo);//SDL_CreateGPUBuffer(GPU::GetDevice(), &vertexBufferInfo);

    GPU::AquirePushCommandBuffer();

    vertexBuffer->Transfer(instances.data());

    SDL_GPUColorTargetInfo colorTargetInfo = {};
    colorTargetInfo.texture = target->Get();
    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);
    pipeline->Bind();
    SDL_GPUBufferBinding vertexBufferBinding = { .buffer = vertexBuffer->Get(), .offset = 0 };
    SDL_BindGPUVertexBuffers(GPU::GetRenderPass(), 0, &vertexBufferBinding, 1);
    SDL_GPUTextureSamplerBinding textureSamplerBinding = { .texture = source->Get(), .sampler = sampler };
    SDL_BindGPUFragmentSamplers(GPU::GetRenderPass(), 0, &textureSamplerBinding, 1);
    float pushConstants[5] = { (float)width, (float)height, (float)offsetX, (float)offsetY, target == strokeTexture ? 1.0f : 0.0f };
    SDL_PushGPUVertexUniformData(GPU::GetCommandBuffer(), 0, pushConstants, sizeof(pushConstants));
    SDL_DrawGPUPrimitives(GPU::GetRenderPass(), 6, instances.size(), 0, 0);
    GPU::EndRendering();

    GPU::SubmitCommandsPopCommandBuffer();

    instances.clear();
    delete vertexBuffer;
}

bool StrokeRenderingContext::IsComplete()
{
    return inputEvents.empty() && !brush->IsDrawing && !immutableInputEvents.empty() && boundingBox.w > 0 && boundingBox.h > 0; // a srctx without any events is never complete
}

void StrokeRenderingContext::Minimize()
{
    if (minimized) return;
    SDL_GPUTextureCreateInfo createInfo = {};
    createInfo.width = boundingBox.w;
    createInfo.height = boundingBox.h;
    createInfo.format = SDL_GPU_TEXTUREFORMAT_R8G8B8A8_UNORM;
    createInfo.usage = SDL_GPU_TEXTUREUSAGE_COLOR_TARGET | SDL_GPU_TEXTUREUSAGE_SAMPLER;
    createInfo.layer_count_or_depth = 1;
    createInfo.num_levels = 1;
    createInfo.sample_count = SDL_GPU_SAMPLECOUNT_1;
    auto tempStrokeTexture = new Texture(createInfo);

    strokeTexture->CopyTo(tempStrokeTexture, { (int)boundingBox.x, (int)boundingBox.y }, { 0, 0 }, { (int)boundingBox.w, (int)boundingBox.h });

    if (offsetX == 0 && offsetY == 0) {
        offsetX = boundingBox.x;
        offsetY = boundingBox.y;
    }
    boundingBox.x = 0.0f;
    boundingBox.y = 0.0f;
    boundingBox.w = (float)tempStrokeTexture->GetWidth();
    boundingBox.h = (float)tempStrokeTexture->GetHeight();

    delete strokeTexture;
    strokeTexture = tempStrokeTexture;
    minimized = true;
}

std::list<PointerEvent> StrokeRenderingContext::GetCompleteInputs()
{
    return immutableInputEvents;
}