#include "core/rendering/sdl3gpu/sdl3gpu_shader.hpp"

#include <cstdint>
#include <cstring>
#include <iostream>

#ifndef __APPLE__
#include "core/rendering/sdl3gpu/spv/postfx_frag_spv.h"
#include "core/rendering/sdl3gpu/spv/postfx_vert_spv.h"
#endif

namespace {

#ifdef __APPLE__
    // NOLINTBEGIN(readability-identifier-naming)
    const char* POSTFX_VERT_MSL_ = R"(
#include <metal_stdlib>
using namespace metal;

struct PostVOut {
    float4 pos [[position]];
    float2 uv;
};

vertex PostVOut postfx_vs(uint vid [[vertex_id]]) {
    const float2 positions[3] = { {-1.0, -1.0}, {3.0, -1.0}, {-1.0, 3.0} };
    const float2 uvs[3]       = { { 0.0,  1.0}, {2.0,  1.0}, { 0.0,-1.0} };
    PostVOut out;
    out.pos = float4(positions[vid], 0.0, 1.0);
    out.uv  = uvs[vid];
    return out;
}
)";

    const char* POSTFX_FRAG_MSL_ = R"(
#include <metal_stdlib>
using namespace metal;

struct PostVOut {
    float4 pos [[position]];
    float2 uv;
};

struct PostFXUniforms {
    float vignette_strength;
    float chroma_strength;
    float scanline_strength;
    float screen_height;
    float mask_strength;
    float gamma_strength;
    float curvature;
    float bleeding;
    float pixel_scale;
    float time;
    float oversample;
    float flicker;
};

static float3 rgb_to_ycc(float3 rgb) {
    return float3(
         0.299f*rgb.r + 0.587f*rgb.g + 0.114f*rgb.b,
        -0.169f*rgb.r - 0.331f*rgb.g + 0.500f*rgb.b + 0.5f,
         0.500f*rgb.r - 0.419f*rgb.g - 0.081f*rgb.b + 0.5f
    );
}
static float3 ycc_to_rgb(float3 ycc) {
    float y  = ycc.x;
    float cb = ycc.y - 0.5f;
    float cr = ycc.z - 0.5f;
    return clamp(float3(
        y + 1.402f*cr,
        y - 0.344f*cb - 0.714f*cr,
        y + 1.772f*cb
    ), 0.0f, 1.0f);
}

fragment float4 postfx_fs(PostVOut                 in    [[stage_in]],
                           texture2d<float>         scene [[texture(0)]],
                           sampler                  samp  [[sampler(0)]],
                           constant PostFXUniforms& u     [[buffer(0)]]) {
    float2 uv = in.uv;

    if (u.curvature > 0.0f) {
        float2 c = uv - 0.5f;
        float rsq = dot(c, c);
        float2 dist = float2(0.05f, 0.1f) * u.curvature;
        float2 barrelScale = 1.0f - 0.23f * dist;
        c += c * (dist * rsq);
        c *= barrelScale;
        if (abs(c.x) >= 0.5f || abs(c.y) >= 0.5f) {
            return float4(0.0f, 0.0f, 0.0f, 1.0f);
        }
        uv = c + 0.5f;
    }

    float3 base = scene.sample(samp, uv).rgb;

    float3 colour;
    if (u.bleeding > 0.0f) {
        float tw    = float(scene.get_width());
        float step  = u.oversample / tw;
        float3 ycc    = rgb_to_ycc(base);
        float3 ycc_l2 = rgb_to_ycc(scene.sample(samp, uv - float2(2.0f*step, 0.0f)).rgb);
        float3 ycc_l1 = rgb_to_ycc(scene.sample(samp, uv - float2(1.0f*step, 0.0f)).rgb);
        float3 ycc_r1 = rgb_to_ycc(scene.sample(samp, uv + float2(1.0f*step, 0.0f)).rgb);
        float3 ycc_r2 = rgb_to_ycc(scene.sample(samp, uv + float2(2.0f*step, 0.0f)).rgb);
        ycc.yz = (ycc_l2.yz + ycc_l1.yz*2.0f + ycc.yz*2.0f + ycc_r1.yz*2.0f + ycc_r2.yz) / 8.0f;
        colour = mix(base, ycc_to_rgb(ycc), u.bleeding);
    } else {
        colour = base;
    }

    float ca = u.chroma_strength * 0.005f * (1.0f + 0.15f * sin(u.time * 7.3f));
    colour.r = scene.sample(samp, uv + float2(ca, 0.0f)).r;
    colour.b = scene.sample(samp, uv - float2(ca, 0.0f)).b;

    if (u.gamma_strength > 0.0f) {
        float3 lin = pow(colour, float3(2.4f));
        colour = mix(colour, lin, u.gamma_strength);
    }

    const float SCAN_DARK_RATIO = 0.333f;
    const float SCAN_DARK_FLOOR = 0.42f;
    if (u.scanline_strength > 0.0f) {
        float ps          = max(1.0f, round(u.pixel_scale));
        float frac_in_row = fract(uv.y * u.screen_height);
        float row_pos     = floor(frac_in_row * ps);
        float bright_rows = (ps < 2.0f) ? ps : ((ps < 3.0f) ? 1.0f : floor(ps * (1.0f - SCAN_DARK_RATIO)));
        float is_dark     = step(bright_rows, row_pos);
        float scan        = mix(1.0f, SCAN_DARK_FLOOR, is_dark);
        colour           *= mix(1.0f, scan, u.scanline_strength);
    }

    if (u.gamma_strength > 0.0f) {
        float3 enc = pow(colour, float3(1.0f/2.2f));
        colour = mix(colour, enc, u.gamma_strength);
    }

    float2 d = uv - 0.5f;
    float vignette = 1.0f - dot(d, d) * u.vignette_strength;
    colour *= clamp(vignette, 0.0f, 1.0f);

    if (u.mask_strength > 0.0f) {
        float whichMask = fract(in.pos.x * 0.3333333f);
        float3 mask = float3(0.80f);
        if (whichMask < 0.3333333f)       mask.x = 1.0f;
        else if (whichMask < 0.6666667f)  mask.y = 1.0f;
        else                              mask.z = 1.0f;
        colour = mix(colour, colour * mask, u.mask_strength);
    }

    if (u.flicker > 0.0f) {
        float flicker_wave = sin(u.time * 100.0f) * 0.5f + 0.5f;
        colour *= 1.0f - u.flicker * 0.04f * flicker_wave;
    }

    return float4(colour, 1.0f);
}
)";
    // NOLINTEND(readability-identifier-naming)
#endif

#ifdef __APPLE__
    auto createShaderMSL(SDL_GPUDevice* device,
        const char* msl_source,
        const char* entrypoint,
        SDL_GPUShaderStage stage,
        std::uint32_t num_samplers,
        std::uint32_t num_uniform_buffers) -> SDL_GPUShader* {
        SDL_GPUShaderCreateInfo info{};
        info.code = reinterpret_cast<const Uint8*>(msl_source);
        info.code_size = std::strlen(msl_source) + 1;
        info.entrypoint = entrypoint;
        info.format = SDL_GPU_SHADERFORMAT_MSL;
        info.stage = stage;
        info.num_samplers = num_samplers;
        info.num_uniform_buffers = num_uniform_buffers;
        return SDL_CreateGPUShader(device, &info);
    }
#else
    auto createShaderSPIRV(SDL_GPUDevice* device,
        const std::uint8_t* spv_code,
        std::size_t spv_size,
        const char* entrypoint,
        SDL_GPUShaderStage stage,
        std::uint32_t num_samplers,
        std::uint32_t num_uniform_buffers) -> SDL_GPUShader* {
        SDL_GPUShaderCreateInfo info{};
        info.code = spv_code;
        info.code_size = spv_size;
        info.entrypoint = entrypoint;
        info.format = SDL_GPU_SHADERFORMAT_SPIRV;
        info.stage = stage;
        info.num_samplers = num_samplers;
        info.num_uniform_buffers = num_uniform_buffers;
        return SDL_CreateGPUShader(device, &info);
    }
#endif

}  // namespace

namespace Rendering {

    SDL3GPUShader::~SDL3GPUShader() { cleanup(); }

    auto SDL3GPUShader::init(SDL_Window* window, int game_width, int game_height) -> bool {
        window_ = window;
        game_width_ = game_width;
        game_height_ = game_height;
        uniforms_.screen_height = static_cast<float>(game_height);

        if (!createDevice()) { return false; }
        if (!createShaders()) {
            cleanup();
            return false;
        }
        if (!createPipeline()) {
            cleanup();
            return false;
        }
        if (!createResources()) {
            cleanup();
            return false;
        }
        // El pipeline CrtPi es crea sempre — l'usuari pot alternar en calent
        // via setActiveShader() sense tornar a init'ar la GPU.
        crtpi_pipeline_ = std::make_unique<CrtPiPipeline>();
        if (!crtpi_pipeline_->init(device_, window_)) {
            cleanup();
            return false;
        }

        is_ready_ = true;
        return true;
    }

    auto SDL3GPUShader::createDevice() -> bool {
#ifdef __APPLE__
        const SDL_GPUShaderFormat FORMAT = SDL_GPU_SHADERFORMAT_MSL;
#else
        const SDL_GPUShaderFormat FORMAT = SDL_GPU_SHADERFORMAT_SPIRV;
#endif
        const char* driver = preferred_driver_.empty() ? nullptr : preferred_driver_.c_str();
        device_ = SDL_CreateGPUDevice(FORMAT, false, driver);
        if (device_ == nullptr) {
            std::cerr << "SDL3GPUShader: SDL_CreateGPUDevice ha fallat: " << SDL_GetError() << '\n';
            return false;
        }
        if (!SDL_ClaimWindowForGPUDevice(device_, window_)) {
            std::cerr << "SDL3GPUShader: SDL_ClaimWindowForGPUDevice ha fallat: " << SDL_GetError() << '\n';
            return false;
        }
        SDL_SetGPUSwapchainParameters(device_, window_, SDL_GPU_SWAPCHAINCOMPOSITION_SDR, bestPresentMode(vsync_));

        const char* name = SDL_GetGPUDeviceDriver(device_);
        driver_name_ = (name != nullptr) ? name : "";
        return true;
    }

    auto SDL3GPUShader::createShaders() -> bool {
#ifdef __APPLE__
        vertex_shader_ = createShaderMSL(device_, POSTFX_VERT_MSL_, "postfx_vs", SDL_GPU_SHADERSTAGE_VERTEX, 0, 0);
        fragment_shader_ = createShaderMSL(device_, POSTFX_FRAG_MSL_, "postfx_fs", SDL_GPU_SHADERSTAGE_FRAGMENT, 1, 1);
#else
        vertex_shader_ = createShaderSPIRV(device_, kpostfx_vert_spv, kpostfx_vert_spv_size, "main", SDL_GPU_SHADERSTAGE_VERTEX, 0, 0);
        fragment_shader_ = createShaderSPIRV(device_, kpostfx_frag_spv, kpostfx_frag_spv_size, "main", SDL_GPU_SHADERSTAGE_FRAGMENT, 1, 1);
#endif
        if (vertex_shader_ == nullptr || fragment_shader_ == nullptr) {
            std::cerr << "SDL3GPUShader: SDL_CreateGPUShader ha fallat: " << SDL_GetError() << '\n';
            return false;
        }
        return true;
    }

    auto SDL3GPUShader::createPipeline() -> bool {
        const SDL_GPUTextureFormat SWAPCHAIN_FORMAT = SDL_GetGPUSwapchainTextureFormat(device_, window_);

        SDL_GPUColorTargetBlendState no_blend{};
        SDL_GPUColorTargetDescription color_target{};
        color_target.format = SWAPCHAIN_FORMAT;
        color_target.blend_state = no_blend;

        SDL_GPUGraphicsPipelineCreateInfo info{};
        info.vertex_shader = vertex_shader_;
        info.fragment_shader = fragment_shader_;
        info.primitive_type = SDL_GPU_PRIMITIVETYPE_TRIANGLELIST;
        info.target_info.num_color_targets = 1;
        info.target_info.color_target_descriptions = &color_target;

        pipeline_ = SDL_CreateGPUGraphicsPipeline(device_, &info);
        if (pipeline_ == nullptr) {
            std::cerr << "SDL3GPUShader: SDL_CreateGPUGraphicsPipeline ha fallat: " << SDL_GetError() << '\n';
            return false;
        }
        return true;
    }

    auto SDL3GPUShader::createResources() -> bool {
        SDL_GPUTextureCreateInfo tex_info{};
        tex_info.type = SDL_GPU_TEXTURETYPE_2D;
        tex_info.format = SDL_GPU_TEXTUREFORMAT_B8G8R8A8_UNORM;
        tex_info.usage = SDL_GPU_TEXTUREUSAGE_SAMPLER;
        tex_info.width = static_cast<Uint32>(game_width_);
        tex_info.height = static_cast<Uint32>(game_height_);
        tex_info.layer_count_or_depth = 1;
        tex_info.num_levels = 1;
        scene_texture_ = SDL_CreateGPUTexture(device_, &tex_info);

        SDL_GPUTransferBufferCreateInfo tb_info{};
        tb_info.usage = SDL_GPU_TRANSFERBUFFERUSAGE_UPLOAD;
        tb_info.size = static_cast<Uint32>(game_width_) * static_cast<Uint32>(game_height_) * 4U;
        upload_buffer_ = SDL_CreateGPUTransferBuffer(device_, &tb_info);

        SDL_GPUSamplerCreateInfo samp_info{};
        samp_info.min_filter = SDL_GPU_FILTER_NEAREST;
        samp_info.mag_filter = SDL_GPU_FILTER_NEAREST;
        samp_info.mipmap_mode = SDL_GPU_SAMPLERMIPMAPMODE_NEAREST;
        samp_info.address_mode_u = SDL_GPU_SAMPLERADDRESSMODE_CLAMP_TO_EDGE;
        samp_info.address_mode_v = SDL_GPU_SAMPLERADDRESSMODE_CLAMP_TO_EDGE;
        samp_info.address_mode_w = SDL_GPU_SAMPLERADDRESSMODE_CLAMP_TO_EDGE;
        sampler_ = SDL_CreateGPUSampler(device_, &samp_info);

        if (scene_texture_ == nullptr || upload_buffer_ == nullptr || sampler_ == nullptr) {
            std::cerr << "SDL3GPUShader: creació de recursos GPU ha fallat: " << SDL_GetError() << '\n';
            return false;
        }
        return true;
    }

    void SDL3GPUShader::uploadPixels(const std::uint32_t* pixels, int width, int height) {
        if (!is_ready_ || upload_buffer_ == nullptr) { return; }
        void* mapped = SDL_MapGPUTransferBuffer(device_, upload_buffer_, false);
        if (mapped == nullptr) { return; }
        std::memcpy(mapped, pixels, static_cast<std::size_t>(width) * static_cast<std::size_t>(height) * 4U);
        SDL_UnmapGPUTransferBuffer(device_, upload_buffer_);
    }

    void SDL3GPUShader::render() {
        if (!is_ready_) { return; }

        SDL_GPUCommandBuffer* cmd = SDL_AcquireGPUCommandBuffer(device_);
        if (cmd == nullptr) { return; }

        // Copy pass: transfer buffer -> scene texture.
        SDL_GPUCopyPass* copy = SDL_BeginGPUCopyPass(cmd);
        if (copy != nullptr) {
            SDL_GPUTextureTransferInfo src{};
            src.transfer_buffer = upload_buffer_;
            src.pixels_per_row = static_cast<Uint32>(game_width_);
            src.rows_per_layer = static_cast<Uint32>(game_height_);

            SDL_GPUTextureRegion dst{};
            dst.texture = scene_texture_;
            dst.w = static_cast<Uint32>(game_width_);
            dst.h = static_cast<Uint32>(game_height_);
            dst.d = 1;

            SDL_UploadToGPUTexture(copy, &src, &dst, false);
            SDL_EndGPUCopyPass(copy);
        }

        // Agafa la textura del swapchain (pot fallar si la finestra està minimitzada).
        SDL_GPUTexture* swapchain = nullptr;
        Uint32 sw = 0;
        Uint32 sh = 0;
        if (!SDL_AcquireGPUSwapchainTexture(cmd, window_, &swapchain, &sw, &sh) || swapchain == nullptr) {
            SDL_SubmitGPUCommandBuffer(cmd);
            return;
        }

        // Camí CrtPi: delega tot el render pass al pipeline dedicat.
        if (active_shader_ == ShaderType::CRTPI && crtpi_pipeline_ != nullptr) {
            crtpi_pipeline_->render(cmd, swapchain, scene_texture_, sampler_, game_width_, game_height_);
            SDL_SubmitGPUCommandBuffer(cmd);
            return;
        }

        uniforms_.time = static_cast<float>(SDL_GetTicks()) / 1000.0F;

        SDL_GPUColorTargetInfo color_target{};
        color_target.texture = swapchain;
        color_target.load_op = SDL_GPU_LOADOP_CLEAR;
        color_target.store_op = SDL_GPU_STOREOP_STORE;
        color_target.clear_color = {.r = 0.0F, .g = 0.0F, .b = 0.0F, .a = 1.0F};

        SDL_GPURenderPass* pass = SDL_BeginGPURenderPass(cmd, &color_target, 1, nullptr);
        if (pass != nullptr) {
            SDL_BindGPUGraphicsPipeline(pass, pipeline_);

            SDL_GPUTextureSamplerBinding binding{};
            binding.texture = scene_texture_;
            binding.sampler = sampler_;
            SDL_BindGPUFragmentSamplers(pass, 0, &binding, 1);

            SDL_PushGPUFragmentUniformData(cmd, 0, &uniforms_, sizeof(PostFXUniforms));
            SDL_DrawGPUPrimitives(pass, 3, 1, 0, 0);
            SDL_EndGPURenderPass(pass);
        }

        SDL_SubmitGPUCommandBuffer(cmd);
    }

    void SDL3GPUShader::setCrtPiParams(const CrtPiParams& params) {
        if (crtpi_pipeline_ != nullptr) {
            crtpi_pipeline_->setParams(params);
        }
    }

    void SDL3GPUShader::setPostFXParams(const PostFXParams& params) {
        uniforms_.vignette_strength = params.vignette;
        uniforms_.scanline_strength = params.scanlines;
        uniforms_.chroma_strength = params.chroma;
        uniforms_.mask_strength = params.mask;
        uniforms_.gamma_strength = params.gamma;
        uniforms_.curvature = params.curvature;
        uniforms_.bleeding = params.bleeding;
        uniforms_.flicker = params.flicker;
    }

    void SDL3GPUShader::setVSync(bool vsync) {
        vsync_ = vsync;
        if (device_ != nullptr && window_ != nullptr) {
            SDL_SetGPUSwapchainParameters(device_, window_, SDL_GPU_SWAPCHAINCOMPOSITION_SDR, bestPresentMode(vsync_));
        }
    }

    auto SDL3GPUShader::bestPresentMode(bool vsync) const -> SDL_GPUPresentMode {
        if (vsync) { return SDL_GPU_PRESENTMODE_VSYNC; }
        if (device_ != nullptr && window_ != nullptr) {
            if (SDL_WindowSupportsGPUPresentMode(device_, window_, SDL_GPU_PRESENTMODE_IMMEDIATE)) {
                return SDL_GPU_PRESENTMODE_IMMEDIATE;
            }
            if (SDL_WindowSupportsGPUPresentMode(device_, window_, SDL_GPU_PRESENTMODE_MAILBOX)) {
                return SDL_GPU_PRESENTMODE_MAILBOX;
            }
        }
        return SDL_GPU_PRESENTMODE_VSYNC;
    }

    void SDL3GPUShader::releaseWindow() {
        if (device_ != nullptr && window_ != nullptr) {
            SDL_WaitForGPUIdle(device_);
            SDL_ReleaseWindowFromGPUDevice(device_, window_);
            window_ = nullptr;
        }
    }

    void SDL3GPUShader::cleanup() {
        is_ready_ = false;
        if (device_ != nullptr) {
            SDL_WaitForGPUIdle(device_);

            // El pipeline CrtPi ha d'alliberar-se mentre el device segueix viu.
            crtpi_pipeline_.reset();

            if (pipeline_ != nullptr) {
                SDL_ReleaseGPUGraphicsPipeline(device_, pipeline_);
                pipeline_ = nullptr;
            }
            if (vertex_shader_ != nullptr) {
                SDL_ReleaseGPUShader(device_, vertex_shader_);
                vertex_shader_ = nullptr;
            }
            if (fragment_shader_ != nullptr) {
                SDL_ReleaseGPUShader(device_, fragment_shader_);
                fragment_shader_ = nullptr;
            }
            if (scene_texture_ != nullptr) {
                SDL_ReleaseGPUTexture(device_, scene_texture_);
                scene_texture_ = nullptr;
            }
            if (upload_buffer_ != nullptr) {
                SDL_ReleaseGPUTransferBuffer(device_, upload_buffer_);
                upload_buffer_ = nullptr;
            }
            if (sampler_ != nullptr) {
                SDL_ReleaseGPUSampler(device_, sampler_);
                sampler_ = nullptr;
            }

            // Si no s'ha cridat releaseWindow() abans, el desreclam és implícit dins
            // DestroyGPUDevice; però això pot bloquejar-se en drivers Wayland si la
            // Wayland surface ja no existeix. El camí recomanat és cridar releaseWindow()
            // des de Screen::destroySdl() ABANS de SDL_DestroyWindow.
            SDL_DestroyGPUDevice(device_);
            device_ = nullptr;
        }
        window_ = nullptr;
    }

}  // namespace Rendering