#include "gpu_pipeline.hpp"
#include "gpu_sprite_batch.hpp"  // for GpuVertex layout

#include <SDL3/SDL_log.h>
#include <cstddef>  // offsetof

// ============================================================================
// MSL Shaders (Metal Shading Language, macOS)
// ============================================================================

// ---------------------------------------------------------------------------
// Sprite vertex shader
// Input: GpuVertex (pos=NDC float2, uv float2, col float4)
// Output: position, uv, col forwarded to fragment stage
// ---------------------------------------------------------------------------
static const char* kSpriteVertMSL = R"(
#include <metal_stdlib>
using namespace metal;

struct SpriteVIn {
    float2 pos [[attribute(0)]];
    float2 uv  [[attribute(1)]];
    float4 col [[attribute(2)]];
};
struct SpriteVOut {
    float4 pos [[position]];
    float2 uv;
    float4 col;
};

vertex SpriteVOut sprite_vs(SpriteVIn in [[stage_in]]) {
    SpriteVOut out;
    out.pos = float4(in.pos, 0.0, 1.0);
    out.uv  = in.uv;
    out.col = in.col;
    return out;
}
)";

// ---------------------------------------------------------------------------
// Sprite fragment shader
// Samples a texture and multiplies by vertex color (for tinting + alpha).
// ---------------------------------------------------------------------------
static const char* kSpriteFragMSL = R"(
#include <metal_stdlib>
using namespace metal;

struct SpriteVOut {
    float4 pos [[position]];
    float2 uv;
    float4 col;
};

fragment float4 sprite_fs(SpriteVOut in [[stage_in]],
                           texture2d<float> tex  [[texture(0)]],
                           sampler           samp [[sampler(0)]]) {
    float4 t = tex.sample(samp, in.uv);
    return float4(t.rgb * in.col.rgb, t.a * in.col.a);
}
)";

// ---------------------------------------------------------------------------
// PostFX vertex shader
// Generates a full-screen triangle from vertex_id (no vertex buffer needed).
// UV mapping: NDC(-1,-1)→UV(0,1)  NDC(-1,3)→UV(0,-1)  NDC(3,-1)→UV(2,1)
// ---------------------------------------------------------------------------
static const char* kPostFXVertMSL = 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;
}
)";

// ---------------------------------------------------------------------------
// PostFX fragment shader
// Samples the offscreen scene texture and applies a subtle vignette.
// ---------------------------------------------------------------------------
static const char* kPostFXFragMSL = R"(
#include <metal_stdlib>
using namespace metal;

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

fragment float4 postfx_fs(PostVOut         in    [[stage_in]],
                           texture2d<float> scene [[texture(0)]],
                           sampler          samp  [[sampler(0)]]) {
    float4 color = scene.sample(samp, in.uv);

    // Subtle vignette: darkens edges proportionally to distance from centre
    float2 d = in.uv - float2(0.5, 0.5);
    float vignette = 1.0 - dot(d, d) * 1.5;
    color.rgb *= clamp(vignette, 0.0, 1.0);

    return color;
}
)";

// ============================================================================
// GpuPipeline implementation
// ============================================================================

bool GpuPipeline::init(SDL_GPUDevice* device,
                       SDL_GPUTextureFormat target_format,
                       SDL_GPUTextureFormat offscreen_format) {
    SDL_GPUShaderFormat supported = SDL_GetGPUShaderFormats(device);
    if (!(supported & SDL_GPU_SHADERFORMAT_MSL)) {
        SDL_Log("GpuPipeline: device does not support MSL shaders (format mask=%u)", supported);
        return false;
    }

    // ----------------------------------------------------------------
    // Sprite pipeline
    // ----------------------------------------------------------------
    SDL_GPUShader* sprite_vert = createShader(device, kSpriteVertMSL, "sprite_vs",
                                              SDL_GPU_SHADERSTAGE_VERTEX, 0, 0);
    SDL_GPUShader* sprite_frag = createShader(device, kSpriteFragMSL, "sprite_fs",
                                              SDL_GPU_SHADERSTAGE_FRAGMENT, 1, 0);
    if (!sprite_vert || !sprite_frag) {
        SDL_Log("GpuPipeline: failed to create sprite shaders");
        if (sprite_vert) SDL_ReleaseGPUShader(device, sprite_vert);
        if (sprite_frag) SDL_ReleaseGPUShader(device, sprite_frag);
        return false;
    }

    // Vertex input: GpuVertex layout
    SDL_GPUVertexBufferDescription vb_desc = {};
    vb_desc.slot              = 0;
    vb_desc.pitch             = sizeof(GpuVertex);
    vb_desc.input_rate        = SDL_GPU_VERTEXINPUTRATE_VERTEX;
    vb_desc.instance_step_rate = 0;

    SDL_GPUVertexAttribute attrs[3] = {};
    attrs[0].location    = 0;
    attrs[0].buffer_slot = 0;
    attrs[0].format      = SDL_GPU_VERTEXELEMENTFORMAT_FLOAT2;
    attrs[0].offset      = static_cast<Uint32>(offsetof(GpuVertex, x));

    attrs[1].location    = 1;
    attrs[1].buffer_slot = 0;
    attrs[1].format      = SDL_GPU_VERTEXELEMENTFORMAT_FLOAT2;
    attrs[1].offset      = static_cast<Uint32>(offsetof(GpuVertex, u));

    attrs[2].location    = 2;
    attrs[2].buffer_slot = 0;
    attrs[2].format      = SDL_GPU_VERTEXELEMENTFORMAT_FLOAT4;
    attrs[2].offset      = static_cast<Uint32>(offsetof(GpuVertex, r));

    SDL_GPUVertexInputState vertex_input = {};
    vertex_input.vertex_buffer_descriptions = &vb_desc;
    vertex_input.num_vertex_buffers         = 1;
    vertex_input.vertex_attributes          = attrs;
    vertex_input.num_vertex_attributes      = 3;

    // Alpha blend state (SRC_ALPHA, ONE_MINUS_SRC_ALPHA)
    SDL_GPUColorTargetBlendState blend = {};
    blend.enable_blend             = true;
    blend.src_color_blendfactor    = SDL_GPU_BLENDFACTOR_SRC_ALPHA;
    blend.dst_color_blendfactor    = SDL_GPU_BLENDFACTOR_ONE_MINUS_SRC_ALPHA;
    blend.color_blend_op           = SDL_GPU_BLENDOP_ADD;
    blend.src_alpha_blendfactor    = SDL_GPU_BLENDFACTOR_ONE;
    blend.dst_alpha_blendfactor    = SDL_GPU_BLENDFACTOR_ONE_MINUS_SRC_ALPHA;
    blend.alpha_blend_op           = SDL_GPU_BLENDOP_ADD;
    blend.enable_color_write_mask  = false;  // write all channels

    SDL_GPUColorTargetDescription color_target_desc = {};
    color_target_desc.format      = offscreen_format;
    color_target_desc.blend_state = blend;

    SDL_GPUGraphicsPipelineCreateInfo sprite_pipe_info = {};
    sprite_pipe_info.vertex_shader      = sprite_vert;
    sprite_pipe_info.fragment_shader    = sprite_frag;
    sprite_pipe_info.vertex_input_state = vertex_input;
    sprite_pipe_info.primitive_type     = SDL_GPU_PRIMITIVETYPE_TRIANGLELIST;
    sprite_pipe_info.target_info.num_color_targets          = 1;
    sprite_pipe_info.target_info.color_target_descriptions  = &color_target_desc;

    sprite_pipeline_ = SDL_CreateGPUGraphicsPipeline(device, &sprite_pipe_info);

    SDL_ReleaseGPUShader(device, sprite_vert);
    SDL_ReleaseGPUShader(device, sprite_frag);

    if (!sprite_pipeline_) {
        SDL_Log("GpuPipeline: sprite pipeline creation failed: %s", SDL_GetError());
        return false;
    }

    // ----------------------------------------------------------------
    // UI overlay pipeline (same as sprite but renders to swapchain format)
    // Reuse sprite shaders with different target format.
    // We create a second version of the sprite pipeline for swapchain.
    // ----------------------------------------------------------------
    // (postfx pipeline targets swapchain; UI overlay also targets swapchain
    //  but needs its own pipeline with swapchain format.)
    // For simplicity, the sprite pipeline is used for the offscreen pass only.
    // The UI overlay is composited via a separate postfx-like pass below.

    // ----------------------------------------------------------------
    // PostFX pipeline
    // ----------------------------------------------------------------
    SDL_GPUShader* postfx_vert = createShader(device, kPostFXVertMSL, "postfx_vs",
                                              SDL_GPU_SHADERSTAGE_VERTEX, 0, 0);
    SDL_GPUShader* postfx_frag = createShader(device, kPostFXFragMSL, "postfx_fs",
                                              SDL_GPU_SHADERSTAGE_FRAGMENT, 1, 0);
    if (!postfx_vert || !postfx_frag) {
        SDL_Log("GpuPipeline: failed to create postfx shaders");
        if (postfx_vert) SDL_ReleaseGPUShader(device, postfx_vert);
        if (postfx_frag) SDL_ReleaseGPUShader(device, postfx_frag);
        return false;
    }

    // PostFX: no vertex input (uses vertex_id), no blend (replace output)
    SDL_GPUColorTargetBlendState no_blend = {};
    no_blend.enable_blend = false;
    no_blend.enable_color_write_mask = false;

    SDL_GPUColorTargetDescription postfx_target_desc = {};
    postfx_target_desc.format      = target_format;
    postfx_target_desc.blend_state = no_blend;

    SDL_GPUVertexInputState no_input = {};

    SDL_GPUGraphicsPipelineCreateInfo postfx_pipe_info = {};
    postfx_pipe_info.vertex_shader      = postfx_vert;
    postfx_pipe_info.fragment_shader    = postfx_frag;
    postfx_pipe_info.vertex_input_state = no_input;
    postfx_pipe_info.primitive_type     = SDL_GPU_PRIMITIVETYPE_TRIANGLELIST;
    postfx_pipe_info.target_info.num_color_targets         = 1;
    postfx_pipe_info.target_info.color_target_descriptions = &postfx_target_desc;

    postfx_pipeline_ = SDL_CreateGPUGraphicsPipeline(device, &postfx_pipe_info);

    SDL_ReleaseGPUShader(device, postfx_vert);
    SDL_ReleaseGPUShader(device, postfx_frag);

    if (!postfx_pipeline_) {
        SDL_Log("GpuPipeline: postfx pipeline creation failed: %s", SDL_GetError());
        return false;
    }

    SDL_Log("GpuPipeline: sprite and postfx pipelines created successfully");
    return true;
}

void GpuPipeline::destroy(SDL_GPUDevice* device) {
    if (sprite_pipeline_) { SDL_ReleaseGPUGraphicsPipeline(device, sprite_pipeline_); sprite_pipeline_ = nullptr; }
    if (postfx_pipeline_) { SDL_ReleaseGPUGraphicsPipeline(device, postfx_pipeline_); postfx_pipeline_ = nullptr; }
}

SDL_GPUShader* GpuPipeline::createShader(SDL_GPUDevice* device,
                                          const char* msl_source,
                                          const char* entrypoint,
                                          SDL_GPUShaderStage stage,
                                          Uint32 num_samplers,
                                          Uint32 num_uniform_buffers) {
    SDL_GPUShaderCreateInfo info = {};
    info.code                = reinterpret_cast<const Uint8*>(msl_source);
    info.code_size           = static_cast<size_t>(strlen(msl_source) + 1);
    info.entrypoint          = entrypoint;
    info.format              = SDL_GPU_SHADERFORMAT_MSL;
    info.stage               = stage;
    info.num_samplers        = num_samplers;
    info.num_storage_textures = 0;
    info.num_storage_buffers  = 0;
    info.num_uniform_buffers  = num_uniform_buffers;

    SDL_GPUShader* shader = SDL_CreateGPUShader(device, &info);
    if (!shader) {
        SDL_Log("GpuPipeline: shader '%s' failed: %s", entrypoint, SDL_GetError());
    }
    return shader;
}