#include "Texture.h"
#include "GPU.h"
#include <SDL3/SDL_gpu.h>
#include <stdexcept>
#include <SDL3_image/SDL_image.h>
#include "Assets.h"
#define QOI_IMPLEMENTATION
#include "qoi.h"

Texture::Texture(SDL_GPUTextureCreateInfo createInfo)
{
    this->createInfo = createInfo;
    this->texture = SDL_CreateGPUTexture(GPU::GetDevice(), &createInfo);
    if (!texture) {
        throw std::runtime_error("Texture creation failed: " + std::string(SDL_GetError()));
    }
}

Texture::Texture(std::string path)
{
    SDL_GPUCommandBuffer* cmd = SDL_AcquireGPUCommandBuffer(GPU::GetDevice());
    SDL_GPUCopyPass* copy_pass = SDL_BeginGPUCopyPass(cmd);
    createInfo = {};
    createInfo.format = SDL_GPU_TEXTUREFORMAT_R8G8B8A8_UNORM;
    createInfo.num_levels = 1;
    createInfo.layer_count_or_depth = 1;
    createInfo.usage = SDL_GPU_TEXTUREUSAGE_SAMPLER;
    texture = IMG_LoadGPUTexture(GPU::GetDevice(), copy_pass, Assets::Path(path).c_str(), (int*)&createInfo.width, (int*)&createInfo.height);
    if (!texture) {
        SDL_CancelGPUCommandBuffer(cmd);
        throw std::runtime_error("Texture creation failed: " + std::string(SDL_GetError()));
    }
    SDL_EndGPUCopyPass(copy_pass);
    SDL_SubmitGPUCommandBuffer(cmd);
}

Texture::~Texture()
{
    SDL_ReleaseGPUTexture(GPU::GetDevice(), texture);
    texture = nullptr;
}

SDL_GPUTexture* Texture::Get()
{
    return texture;
}

SDL_GPUTextureCreateInfo Texture::GetCreateInfo()
{
    return createInfo;
}

int Texture::GetWidth()
{
    return createInfo.width;
}

int Texture::GetHeight()
{
    return createInfo.height;
}

SDL_GPUTextureFormat Texture::GetFormat()
{
    return createInfo.format;
}

void Texture::Clear(SDL_FColor color)
{
    GPU::AquirePushCommandBuffer();
    SDL_GPUColorTargetInfo targetInfo = {};
    targetInfo.texture = texture;
    targetInfo.clear_color = color;
    targetInfo.load_op = SDL_GPU_LOADOP_CLEAR;
    targetInfo.store_op = SDL_GPU_STOREOP_STORE;
    targetInfo.mip_level = 0;
    targetInfo.layer_or_depth_plane = 0;
    targetInfo.cycle = false;
    GPU::BeginRendering({ &targetInfo }, nullptr);
    GPU::EndRendering();
    GPU::SubmitCommandsPopCommandBuffer();
}

void Texture::CopyTo(Texture* target, SDL_Point sourceOffset, SDL_Point destinationOffset, SDL_Point size)
{
    GPU::AquirePushCommandBuffer();
    SDL_GPUTextureLocation sourceLocation = {};
    sourceLocation.texture = this->Get();
    sourceLocation.x = sourceOffset.x;
    sourceLocation.y = sourceOffset.y;
    SDL_GPUTextureLocation destinationLocation = {};
    destinationLocation.texture = target->Get();
    destinationLocation.x = destinationOffset.x;
    destinationLocation.y = destinationOffset.y;
    GPU::BeginCopying();
    SDL_CopyGPUTextureToTexture(GPU::GetCopyPass(), &sourceLocation, &destinationLocation, size.x, size.y, 1, false);
    GPU::EndCopying();
    GPU::SubmitCommandsPopCommandBuffer();
}

void* Texture::Download() {
    int bufferSize = GetWidth() * GetHeight() * 4;
    SDL_GPUTransferBufferCreateInfo transferBufferCreateInfo = {};
    transferBufferCreateInfo.size = bufferSize;
    transferBufferCreateInfo.usage = SDL_GPU_TRANSFERBUFFERUSAGE_DOWNLOAD;    
    
    SDL_GPUTransferBuffer* transferBuffer = SDL_CreateGPUTransferBuffer(GPU::GetDevice(), &transferBufferCreateInfo);
    if (!transferBuffer) {
        throw std::runtime_error("Texture download failed: " + std::string(SDL_GetError()));
    }

    SDL_GPUTextureRegion sourceRegion = {};
    sourceRegion.texture = texture;
    sourceRegion.x = 0;
    sourceRegion.y = 0;
    sourceRegion.w = GetWidth();
    sourceRegion.h = GetHeight();
    sourceRegion.mip_level = 0;
    sourceRegion.layer = 0;

    SDL_GPUTextureTransferInfo destinationRegion = {};
    destinationRegion.transfer_buffer = transferBuffer;
    destinationRegion.offset = 0;
    destinationRegion.pixels_per_row = GetWidth();
    destinationRegion.rows_per_layer = GetHeight();

    GPU::AquirePushCommandBuffer();
    GPU::BeginCopying();

    SDL_DownloadFromGPUTexture(GPU::GetCopyPass(), &sourceRegion, &destinationRegion);
    GPU::EndCopying();
    GPU::SubmitCommandsPopCommandBufferAwaitFence();   
    
    void* data = SDL_MapGPUTransferBuffer(GPU::GetDevice(), transferBuffer, false);
    if (!data) {
        throw std::runtime_error("Texture download failed: " + std::string(SDL_GetError()));
    }

    void* copy = malloc(bufferSize);
    memcpy(copy, data, bufferSize);

    SDL_UnmapGPUTransferBuffer(GPU::GetDevice(), transferBuffer);
    SDL_ReleaseGPUTransferBuffer(GPU::GetDevice(), transferBuffer);
    
    return copy;
}

void* Texture::DownloadAsQoI(size_t& outSize) {
    void* rawData = Download();
    
    qoi_desc desc = {
        .width = (uint32_t)GetWidth(),
        .height = (uint32_t)GetHeight(),
        .channels = 4,
        .colorspace = QOI_SRGB
    };
    
    int encodedSize = 0;
    void* encodedData = qoi_encode(rawData, &desc, &encodedSize);
    outSize = encodedSize;
    
    free(rawData);
    
    return encodedData;
}

void Texture::Upload(void* data) {
    int bufferSize = GetWidth() * GetHeight() * 4;
    SDL_GPUTransferBufferCreateInfo transferBufferCreateInfo = {};
    transferBufferCreateInfo.size = bufferSize;
    transferBufferCreateInfo.usage = SDL_GPU_TRANSFERBUFFERUSAGE_UPLOAD;
    SDL_GPUTransferBuffer* transfer = SDL_CreateGPUTransferBuffer(GPU::GetDevice(), &transferBufferCreateInfo);
    if (!transfer) {
        throw std::runtime_error("Texture upload failed: " + std::string(SDL_GetError()));
    }
    void* mapped = SDL_MapGPUTransferBuffer(GPU::GetDevice(), transfer, false);
    if (!mapped) {
        SDL_ReleaseGPUTransferBuffer(GPU::GetDevice(), transfer);
        throw std::runtime_error("Texture upload failed: " + std::string(SDL_GetError()));
    }
    memcpy(mapped, data, bufferSize);
    SDL_UnmapGPUTransferBuffer(GPU::GetDevice(), transfer);

    GPU::AquirePushCommandBuffer();
    GPU::BeginCopying();

    SDL_GPUTextureTransferInfo src = {};
    src.transfer_buffer = transfer;
    src.offset = 0;
    src.pixels_per_row = GetWidth();
    src.rows_per_layer = GetHeight();

    SDL_GPUTextureRegion dst = {};
    dst.texture = texture;
    dst.mip_level = 0;
    dst.layer = 0;
    dst.x = 0;
    dst.y = 0;
    dst.z = 0;
    dst.w = (Uint32)GetWidth();
    dst.h = (Uint32)GetHeight();
    dst.d = 1;

    SDL_UploadToGPUTexture(GPU::GetCopyPass(), &src, &dst, false);
    GPU::EndCopying();
    GPU::SubmitCommandsPopCommandBufferAwaitFence();

    SDL_ReleaseGPUTransferBuffer(GPU::GetDevice(), transfer);
}

void Texture::UploadAsQoI(const void* data, size_t size) {
    qoi_desc desc;
    void* decoded = qoi_decode(data, (int)size, &desc, 4);
    if (!decoded) {
        throw std::runtime_error("QOI decode failed");
    }
    Upload(decoded);
    free(decoded);
}