#include "sprite_batch.h"

#include <cstdlib>
#include <cstring>

namespace fn::gfx {

namespace {

struct Vertex {
    float x, y;
    float u, v;
    float r, g, b, a;
};

const char* VS =
#if defined(__EMSCRIPTEN__)
    "#version 300 es\n"
#else
    "#version 330 core\n"
#endif
    "uniform mat4 u_view_proj;\n"
    "layout(location = 0) in vec2 a_pos;\n"
    "layout(location = 1) in vec2 a_uv;\n"
    "layout(location = 2) in vec4 a_color;\n"
    "out vec2 v_uv;\n"
    "out vec4 v_color;\n"
    "void main() {\n"
    "  v_uv = a_uv;\n"
    "  v_color = a_color;\n"
    "  gl_Position = u_view_proj * vec4(a_pos, 0.0, 1.0);\n"
    "}\n";

const char* FS =
#if defined(__EMSCRIPTEN__)
    "#version 300 es\n"
    "precision mediump float;\n"
#else
    "#version 330 core\n"
#endif
    "in vec2 v_uv;\n"
    "in vec4 v_color;\n"
    "out vec4 frag;\n"
    "uniform sampler2D u_tex;\n"
    "void main() {\n"
    "  frag = texture(u_tex, v_uv) * v_color;\n"
    "}\n";

void flush(SpriteBatch& b) {
    if (b.cpu_count_quads == 0) return;
    int verts = b.cpu_count_quads * 6;
    sg_range data{ b.cpu_buffer, (size_t)verts * sizeof(Vertex) };
    sg_update_buffer(b.vbo, &data);

    sg_apply_pipeline(b.pipeline);

    sg_bindings bind{};
    bind.vertex_buffers[0] = b.vbo;
    bind.views[0] = b.current_view;
    bind.samplers[0] = b.sampler;
    sg_apply_bindings(&bind);

    sg_range proj_range{ b.proj, sizeof(b.proj) };
    sg_apply_uniforms(0, &proj_range);

    sg_draw(0, verts, 1);
    b.cpu_count_quads = 0;
}

}  // namespace

SpriteBatch sprite_batch_create(int cpu_capacity_quads) {
    SpriteBatch b{};
    b.cpu_capacity_quads = cpu_capacity_quads > 0 ? cpu_capacity_quads : 4096;
    b.cpu_buffer = std::malloc((size_t)b.cpu_capacity_quads * 6 * sizeof(Vertex));
    if (!b.cpu_buffer) return b;

    sg_buffer_desc bd{};
    bd.size = (size_t)b.cpu_capacity_quads * 6 * sizeof(Vertex);
    bd.usage.vertex_buffer = true;
    bd.usage.stream_update = true;
    b.vbo = sg_make_buffer(&bd);

    sg_shader_desc sd{};
    sd.vertex_func.source = VS;
    sd.fragment_func.source = FS;
    sd.attrs[0].glsl_name = "a_pos";
    sd.attrs[1].glsl_name = "a_uv";
    sd.attrs[2].glsl_name = "a_color";
    sd.uniform_blocks[0].stage = SG_SHADERSTAGE_VERTEX;
    sd.uniform_blocks[0].size = 16 * sizeof(float);
    sd.uniform_blocks[0].layout = SG_UNIFORMLAYOUT_NATIVE;
    sd.uniform_blocks[0].glsl_uniforms[0].type = SG_UNIFORMTYPE_MAT4;
    sd.uniform_blocks[0].glsl_uniforms[0].glsl_name = "u_view_proj";
    sd.views[0].texture.stage = SG_SHADERSTAGE_FRAGMENT;
    sd.views[0].texture.image_type = SG_IMAGETYPE_2D;
    sd.views[0].texture.sample_type = SG_IMAGESAMPLETYPE_FLOAT;
    sd.samplers[0].stage = SG_SHADERSTAGE_FRAGMENT;
    sd.samplers[0].sampler_type = SG_SAMPLERTYPE_FILTERING;
    sd.texture_sampler_pairs[0].stage = SG_SHADERSTAGE_FRAGMENT;
    sd.texture_sampler_pairs[0].view_slot = 0;
    sd.texture_sampler_pairs[0].sampler_slot = 0;
    sd.texture_sampler_pairs[0].glsl_name = "u_tex";
    sg_shader shd = sg_make_shader(&sd);

    sg_pipeline_desc pd{};
    pd.shader = shd;
    pd.layout.attrs[0].format = SG_VERTEXFORMAT_FLOAT2;
    pd.layout.attrs[1].format = SG_VERTEXFORMAT_FLOAT2;
    pd.layout.attrs[2].format = SG_VERTEXFORMAT_FLOAT4;
    pd.colors[0].blend.enabled = true;
    pd.colors[0].blend.src_factor_rgb = SG_BLENDFACTOR_SRC_ALPHA;
    pd.colors[0].blend.dst_factor_rgb = SG_BLENDFACTOR_ONE_MINUS_SRC_ALPHA;
    pd.colors[0].blend.src_factor_alpha = SG_BLENDFACTOR_ONE;
    pd.colors[0].blend.dst_factor_alpha = SG_BLENDFACTOR_ONE_MINUS_SRC_ALPHA;
    pd.primitive_type = SG_PRIMITIVETYPE_TRIANGLES;
    b.pipeline = sg_make_pipeline(&pd);

    sg_sampler_desc smd{};
    smd.min_filter = SG_FILTER_LINEAR;
    smd.mag_filter = SG_FILTER_LINEAR;
    smd.wrap_u = SG_WRAP_CLAMP_TO_EDGE;
    smd.wrap_v = SG_WRAP_CLAMP_TO_EDGE;
    b.sampler = sg_make_sampler(&smd);

    b.ok = true;
    return b;
}

void sprite_batch_destroy(SpriteBatch& b) {
    if (b.cpu_buffer) { std::free(b.cpu_buffer); b.cpu_buffer = nullptr; }
    if (b.pipeline.id) sg_destroy_pipeline(b.pipeline);
    if (b.vbo.id)      sg_destroy_buffer(b.vbo);
    if (b.sampler.id)  sg_destroy_sampler(b.sampler);
    b = {};
}

void sprite_batch_begin(SpriteBatch& b, const float view_proj_col_major[16]) {
    std::memcpy(b.proj, view_proj_col_major, sizeof(b.proj));
    b.cpu_count_quads = 0;
    b.current_view = {};
    b.in_pass = true;
}

void sprite_batch_draw(SpriteBatch& b, sg_view view, int img_w, int img_h,
                       fn::math2d::Rect src, fn::math2d::Rect dst,
                       fn::math2d::Color tint) {
    if (!b.in_pass || !b.ok) return;
    if (b.current_view.id != view.id) {
        flush(b);
        b.current_view = view;
    }
    if (b.cpu_count_quads >= b.cpu_capacity_quads) flush(b);

    float u0 = src.x / (float)img_w;
    float v0 = src.y / (float)img_h;
    float u1 = (src.x + src.w) / (float)img_w;
    float v1 = (src.y + src.h) / (float)img_h;

    Vertex* base = (Vertex*)b.cpu_buffer + b.cpu_count_quads * 6;
    Vertex tl{ dst.x,         dst.y,         u0, v0, tint.r, tint.g, tint.b, tint.a };
    Vertex tr{ dst.x + dst.w, dst.y,         u1, v0, tint.r, tint.g, tint.b, tint.a };
    Vertex bl{ dst.x,         dst.y + dst.h, u0, v1, tint.r, tint.g, tint.b, tint.a };
    Vertex br{ dst.x + dst.w, dst.y + dst.h, u1, v1, tint.r, tint.g, tint.b, tint.a };
    base[0] = tl; base[1] = tr; base[2] = br;
    base[3] = tl; base[4] = br; base[5] = bl;
    b.cpu_count_quads++;
}

void sprite_batch_end(SpriteBatch& b) {
    if (!b.in_pass) return;
    flush(b);
    b.in_pass = false;
}

}  // namespace fn::gfx