#include "atlas.h"

#include <cstring>

#include "cglm/types.h"
#include "mimalloc.h"
#include "utils.h"

static void wayc_atlas_gpu_usage_desc(struct sg_image_usage* usage) {
  wayc_notnull(usage);

  usage->immutable = false;
  usage->dynamic_update = true;
}

static void wayc_atlas_gpu_desc(struct sg_image_desc* desc, u32 width,
                                u32 height, enum sg_pixel_format format) {
  wayc_notnull(desc);

  struct sg_image_usage usage = {};
  wayc_atlas_gpu_usage_desc(&usage);

  desc->type = SG_IMAGETYPE_2D;
  desc->width = width;
  desc->height = height;
  desc->num_mipmaps = 1;
  desc->pixel_format = format;
  desc->usage = usage;
}

static inline usize wayc_atlas_size(u32 width, u32 height,
                                    enum sg_pixel_format format) {
  return sg_query_surface_pitch(format, width, height, 1);
}

enum atlas_error_e wayc_atlas_init(struct atlas_s* atlas, u32 width, u32 height,
                                   enum sg_pixel_format format) {
  wayc_notnull(atlas);
  memset(atlas, 0, sizeof(*atlas));
  bool success = false;

  usize atlas_size = wayc_atlas_size(width, height, format);
  u8* cpu_atlas = (u8*)mi_malloc(atlas_size);
  wayc_defer_cond(mi_free(cpu_atlas), success, true);
  memset(cpu_atlas, 0, atlas_size);

  struct sg_image_desc desc = {};
  wayc_atlas_gpu_desc(&desc, width, height, format);

  struct sg_image gpu_atlas = sg_make_image(&desc);
  if (sg_query_image_state(gpu_atlas) != SG_RESOURCESTATE_VALID)
    return ATLAS_ERROR_GPU_ATLAS;

  atlas->width = width;
  atlas->height = height;
  atlas->format = format;
  atlas->cpu_atlas = cpu_atlas;
  atlas->gpu_atlas = gpu_atlas;

  success = true;
  return ATLAS_ERROR_NONE;
}

enum atlas_error_e wayc_atlas_insert(struct atlas_s* atlas, u32 width,
                                     u32 height, const u8* data, vec2* uv0,
                                     vec2* uv1) {
  wayc_notnull(atlas);
  wayc_notnull(data);
  wayc_notnull(uv0);
  wayc_notnull(uv1);

  struct atlas_packer_s* packer = &atlas->packer;

  u32 x = WAYC_X(packer->cursor);
  u32 y = WAYC_Y(packer->cursor);

  // wrap
  if (x + width > atlas->width) {
    x = 0;
    y += packer->row_height;
    packer->row_height = 0;
  }

  if (y + height > atlas->height) return ATLAS_ERROR_ATLAS_FULL;

  const int atlas_pitch =
      sg_query_surface_pitch(atlas->format, atlas->width, 1, 1);
  const int bpp = sg_query_surface_pitch(atlas->format, 1, 1, 1);
  const int src_pitch = width * bpp;

  for (u32 row = 0; row < height; row++) {
    const u8* src_row = data + row * src_pitch;
    u8* dst_row = atlas->cpu_atlas + (y + row) * atlas_pitch;
    u8* dst_px = dst_row + x * bpp;

    memcpy(dst_px, src_row, src_pitch);
  }

  const f32 W = (f32)atlas->width;
  const f32 H = (f32)atlas->height;

  WAYC_X(*uv0) = (f32)x / W;
  WAYC_Y(*uv0) = (f32)y / H;
  WAYC_X(*uv1) = (f32)(x + width) / W;
  WAYC_Y(*uv1) = (f32)(y + height) / H;

  WAYC_X(packer->cursor) = x + width;
  WAYC_Y(packer->cursor) = y;
  packer->row_height = wayc_max(packer->row_height, height);
  atlas->cpu_dirty = true;

  return ATLAS_ERROR_NONE;
}

enum atlas_error_e wayc_atlas_flush(struct atlas_s* atlas) {
  wayc_notnull(atlas);
  if (!atlas->cpu_dirty) return ATLAS_ERROR_NONE;

  usize size = wayc_atlas_size(atlas->width, atlas->height, atlas->format);

  struct sg_range range = {};
  range.ptr = atlas->cpu_atlas;
  range.size = size;

  struct sg_image_data data = {};
  data.mip_levels[0] = range;

  sg_update_image(atlas->gpu_atlas, &data);
  if (sg_query_image_state(atlas->gpu_atlas) != SG_RESOURCESTATE_VALID)
    return ATLAS_ERROR_UPLOAD_FAILED;

  atlas->cpu_dirty = false;
  return ATLAS_ERROR_NONE;
}

void wayc_atlas_deinit(struct atlas_s* atlas) {
  wayc_notnull(atlas);

  enum sg_resource_state state = sg_query_image_state(atlas->gpu_atlas);
  if (atlas->cpu_atlas == nullptr || state != SG_RESOURCESTATE_VALID) return;

  sg_destroy_image(atlas->gpu_atlas);
  mi_free(atlas->cpu_atlas);

  atlas->width = 0;
  atlas->height = 0;
  atlas->cpu_atlas = nullptr;
  atlas->gpu_atlas = {};
}