struct VertexInput {
    @location(0) offset: vec2<f32>,
    @location(1) uv: vec2<f32>,
};

struct VertexOutput {
    @builtin(position) position: vec4<f32>,
    @location(0) tex_coords: vec2<f32>,
    @location(1) color: vec4<f32>,
    @location(2) @interpolate(flat) texture_slot: u32,
    @location(3) @interpolate(flat) texture_slot2: u32,
    @location(4) blend_factor: f32,
};

struct SpriteInstance {
    position: vec2<f32>,
    size: vec2<f32>,
    uv_min: vec2<f32>,
    uv_max: vec2<f32>,
    color: vec4<f32>,
    rotation_scale: vec4<f32>,
    depth: f32,
    texture_slot: u32,
    texture_slot2: u32,
    blend_factor: f32,
};

struct GlobalUniforms {
    projection: mat4x4<f32>,
    screen_size: vec2<f32>,
    atlas_slots_per_row: f32,
    atlas_slot_uv_size: f32,
};

@group(0) @binding(0)
var<uniform> globals: GlobalUniforms;

@group(0) @binding(1)
var<storage, read> instances: array<SpriteInstance>;

@group(1) @binding(0)
var sprite_atlas: texture_2d<f32>;

@group(1) @binding(1)
var sprite_sampler: sampler;

@vertex
fn vs_main(
    in: VertexInput,
    @builtin(instance_index) instance_index: u32,
) -> VertexOutput {
    let instance = instances[instance_index];

    let scaled_corner = in.offset * instance.size * instance.rotation_scale.zw;

    let cos_r = instance.rotation_scale.x;
    let sin_r = instance.rotation_scale.y;
    let rotated = vec2<f32>(
        scaled_corner.x * cos_r - scaled_corner.y * sin_r,
        scaled_corner.x * sin_r + scaled_corner.y * cos_r
    );

    let world_position = instance.position + rotated;

    let uv = mix(instance.uv_min, instance.uv_max, in.uv);

    var output: VertexOutput;
    output.position = globals.projection * vec4<f32>(world_position, instance.depth, 1.0);
    output.tex_coords = uv;
    output.color = instance.color;
    output.texture_slot = instance.texture_slot;
    output.texture_slot2 = instance.texture_slot2;
    output.blend_factor = instance.blend_factor;

    return output;
}

fn sample_atlas_slot(slot: u32, uv: vec2<f32>) -> vec4<f32> {
    let slots_per_row = u32(globals.atlas_slots_per_row);
    let row = f32(slot / slots_per_row);
    let col = f32(slot % slots_per_row);

    let slot_u = col * globals.atlas_slot_uv_size;
    let slot_v = row * globals.atlas_slot_uv_size;

    let atlas_uv = vec2<f32>(
        slot_u + uv.x * globals.atlas_slot_uv_size,
        slot_v + uv.y * globals.atlas_slot_uv_size
    );

    return textureSample(sprite_atlas, sprite_sampler, atlas_uv);
}

@fragment
fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
    let tex1 = sample_atlas_slot(in.texture_slot, in.tex_coords);
    let tex2 = sample_atlas_slot(in.texture_slot2, in.tex_coords);
    let blended = mix(tex1, tex2, in.blend_factor);
    return blended * in.color;
}
