#include "atlas.h"

#include <cstring>

#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;
  atlas->cpu_dirty = false;

  success = 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 = {};
}