SPR_Ctx SPR = Zi;

////////////////////////////////////////////////////////////
//~ Bootstrap

void SPR_Bootstrap(void)
{
  // FIXME: Initialize nil/unready sprite texture here

  OnAsyncTick(SPR_TickAsync);
}

////////////////////////////////////////////////////////////
//~ Helpers

SPR_SheetKey SPR_SheetKeyFromResource(ResourceKey resource)
{
  SPR_SheetKey result = Zi;
  result.r = resource;
  return result;
}

SPR_SpanKey SPR_SpanKeyFromName(String name)
{
  SPR_SpanKey result = Zi;
  result.v = HashString(name);
  return result;
}

b32 SPR_IsSheetKeyNil(SPR_SheetKey key)
{
  return key.r.v == 0;
}

b32 SPR_IsSpanKeyNil(SPR_SpanKey key)
{
  return key.v == 0;
}

String SPR_NameFromRayKind(SPR_RayKind kind)
{
  PERSIST Readonly String names[SPR_RayKind_COUNT] = {
    #define X(kind_name, layer_name, ...) [SPR_RayKind_##kind_name] = CompLit(#layer_name),
      SPR_RayKindXList(X)
    #undef X
  };
  String result = Zi;
  if (kind >= 0 && kind < countof(names))
  {
    result = names[kind];
  }
  return result;
}

SPR_LayerKind SPR_LayerKindFromName(String name)
{
  SPR_LayerKind result = SPR_LayerKind_Visual;
  if (StringBeginsWith(name, Lit(".")) || StringBeginsWith(name, Lit("#")))
  {
    result = SPR_LayerKind_Hidden;
  }
  return result;
}

////////////////////////////////////////////////////////////
//~ Lookup

SPR_Sprite SPR_SpriteFromSheet(SPR_SheetKey sheet_key, SPR_SpanKey span_key, i64 frame_seq)
{
  //////////////////////////////
  //- Fetch sheet

  SPR_SheetEntry *sheet = 0;
  {
    SPR_SheetBin *sheet_bin = &SPR.sheet_bins[sheet_key.r.v % countof(SPR.sheet_bins)];
    Lock sheet_bin_lock = LockS(&sheet_bin->mutex);
    {
      for (sheet = sheet_bin->first; sheet; sheet = sheet->next_in_bin)
      {
        if (sheet->key.r.v == sheet_key.r.v)
        {
          break;
        }
      }
    }
    Unlock(&sheet_bin_lock);
  }

  //////////////////////////////
  //- Decode sheet

  if (!sheet)
  {
    SPR_SheetBin *sheet_bin = &SPR.sheet_bins[sheet_key.r.v % countof(SPR.sheet_bins)];
    Lock sheet_bin_lock = LockE(&sheet_bin->mutex);
    {
      for (sheet = sheet_bin->first; sheet; sheet = sheet->next_in_bin)
      {
        if (sheet->key.r.v == sheet_key.r.v)
        {
          break;
        }
      }
      if (!sheet)
      {
        Arena *perm = PermArena();
        sheet = PushStruct(perm, SPR_SheetEntry);
        SllStackPushN(sheet_bin->first, sheet, next_in_bin);
        sheet->key = sheet_key;

        String sheet_data = DataFromResource(sheet->key.r);
        String sheet_name = NameFromResource(sheet->key.r);
        LogInfoF("Decoding sprite sheet %F \"%F\" (%F bytes)", FmtHandle(sheet->key.r.v), FmtString(sheet_name), FmtUint(sheet_data.len));
        sheet->meta = ASE_DecodeMeta(perm, sheet_data);

        //- Init slices
        sheet->slices_count = MaxI64(sheet->meta.frames_count, 1);
        sheet->slices = PushStructs(perm, SPR_SliceEntry, sheet->slices_count);
        for (i64 slice_idx = 0; slice_idx < sheet->slices_count; ++slice_idx)
        {
          SPR_SliceEntry *slice = &sheet->slices[slice_idx];
          slice->canvas_rect = Rng2Empty;
          Atomic64Set(&slice->atlas_copy_completion_target, I64Max);
          for (SPR_RayKind ray_kind = 0; ray_kind < SPR_RayKind_COUNT; ++ray_kind)
          {
            slice->rays[ray_kind].dir.x = 1;
          }
        }

        //- Compute slice rect
        for (ASE_Layer *ase_layer = sheet->meta.first_layer; ase_layer; ase_layer = ase_layer->next)
        {
          SPR_LayerKind kind = SPR_LayerKindFromName(ase_layer->name);
          if (kind == SPR_LayerKind_Visual)
          {
            for (i64 slice_idx = 0; slice_idx < sheet->meta.frames_count; ++slice_idx)
            {
              ASE_Cel *cel = &ase_layer->cels[slice_idx];
              SPR_SliceEntry *slice = &sheet->slices[slice_idx];
              slice->canvas_rect = UnionRng2(slice->canvas_rect, cel->bounds);
            }
          }
        }

        //- Push async rasterization commands
        if (sheet->meta.frames_count > 0)
        {
          Lock cmds_lock = LockE(&SPR.submit.mutex);
          {
            i64 submit_count = 0;
            for (i64 slice_idx = 0; slice_idx < sheet->meta.frames_count; ++slice_idx)
            {
              SPR_SliceEntry *slice = &sheet->slices[slice_idx];
              if (!IsRng2Empty(slice->canvas_rect))
              {
                SPR_CmdNode *cmd_node = SPR.submit.first_free;
                if (cmd_node)
                {
                  SllStackPop(SPR.submit.first_free);
                  ZeroStruct(cmd_node);
                }
                else
                {
                  cmd_node = PushStruct(perm, SPR_CmdNode);
                }
                cmd_node->cmd.sheet = sheet;
                cmd_node->cmd.slice_idx = slice_idx;
                SllQueuePush(SPR.submit.first, SPR.submit.last, cmd_node);
                submit_count += 1;
              }
            }
            if (submit_count > 0)
            {
              SPR.submit.count += submit_count;
              SignalAsyncTick();
            }
          }
          Unlock(&cmds_lock);
        }

        //- Compute rays
        {
          TempArena scratch = BeginScratchNoConflict();
          for (SPR_RayKind ray_kind = 0; ray_kind < SPR_RayKind_COUNT; ++ray_kind)
          {
            String ray_name = SPR_NameFromRayKind(ray_kind);
            b32 match = 0;
            for (ASE_Layer *ase_layer = sheet->meta.last_layer; ase_layer && !match; ase_layer = ase_layer->prev)
            {
              if (MatchString(ray_name, ase_layer->name))
              {
                match = 1;
                for (i64 slice_idx = 0; slice_idx < sheet->meta.frames_count; ++slice_idx)
                {
                  ASE_Cel *ase_cel = &ase_layer->cels[slice_idx];
                  SPR_SliceEntry *slice = &sheet->slices[slice_idx];
                  ASE_Image image = ASE_DecompressImageFromCel(scratch.arena, ase_cel);
                  if (!IsRng2Empty(image.bounds))
                  {
                    u32 ray_pix = image.pixels[0];
                    u32 alpha = (ray_pix >> 24) & 0xFF;
                    if (alpha > 0)
                    {
                      // TODO: Different quantization so that 128 equals 0, instead of approximately 0
                      f32 dir_x = (((f32)((ray_pix >> 0) & 0xFF) / 255.0) * 2.0) - 1;
                      f32 dir_y = (((f32)((ray_pix >> 8) & 0xFF) / 255.0) * 2.0) - 1;
                      Vec2 dir = NormVec2(VEC2(dir_x, dir_y));
                      slice->rays[ray_kind].pos = AddVec2(SubVec2(ase_cel->bounds.p0, slice->canvas_rect.p0), VEC2(0.5, 0.5));
                      slice->rays[ray_kind].dir = dir;
                    }
                  }
                }
              }
            }
          }
          EndScratch(scratch);
        }

        //- Init spans
        {
          sheet->span_bins_count = MaxI64(1, NextPow2U64(sheet->meta.spans_count * 4));
          sheet->span_bins = PushStructs(perm, SPR_SpanBin, sheet->span_bins_count);
          for (ASE_Span *ase_span = sheet->meta.first_span; ase_span; ase_span = ase_span->next)
          {
            SPR_SpanKey new_span_key = { .v = HashString(ase_span->name) };
            SPR_SpanBin *span_bin = &sheet->span_bins[new_span_key.v % sheet->span_bins_count];
            SPR_SpanEntry *span = 0;
            for (span = span_bin->first; span; span = span->next_in_bin)
            {
              if (span->key.v == new_span_key.v)
              {
                break;
              }
            }
            if (!span)
            {
              span = PushStruct(perm, SPR_SpanEntry);
              SllQueuePush(sheet->first_span, sheet->last_span, span);
              SllStackPushN(span_bin->first, span, next_in_bin);
              span->key = new_span_key;
              span->from = ase_span->from;
              span->to = ase_span->to;
            }
          }

          // Insert nil span
          {
            SPR_SpanBin *span_bin = &sheet->span_bins[0];
            SPR_SpanEntry *span = PushStruct(perm, SPR_SpanEntry);
            SllQueuePush(sheet->first_span, sheet->last_span, span);
            SllStackPushN(span_bin->first, span, next_in_bin);
            span->key = SPR_NilSpanKey;
            span->from = 0;
            span->to = sheet->slices_count;
          }
        }

        // TODO: More rigorous validation
        sheet->ok = sheet->meta.ok;
      }
    }
    Unlock(&sheet_bin_lock);
  }

  //////////////////////////////
  //- Fetch span

  // FIXME: Ensure slices_count always > 0
  i64 span_start = 0;
  i64 span_end = sheet->slices_count;
  b32 span_matched = 0;
  {
    SPR_SpanEntry *span = 0;
    if (sheet->ok)
    {
      SPR_SpanBin *span_bin = &sheet->span_bins[span_key.v % sheet->span_bins_count];
      for (span = span_bin->first; span; span = span->next_in_bin)
      {
        if (span->key.v == span_key.v)
        {
          span_matched = 1;
          break;
        }
      }
    }
    if (!span)
    {
      span = sheet->first_span;
    }
    if (span)
    {
      span_start = span->from;
      span_end = span->to;
    }
  }

  //////////////////////////////
  //- Fetch slice

  SPR_SliceEntry *slice = 0;
  {
    // FIXME: Ensure span->end is never <= span->start
    i64 slice_idx = span_start + (frame_seq % (span_end - span_start));
    slice = &sheet->slices[slice_idx];
  }

  //////////////////////////////
  //- Compute result

  SPR_Sprite result = Zi;
  {
    b32 slice_ready = 0;
    if (sheet->ok)
    {
      i64 completion = G_CompletionValueFromQueue(G_QueueKind_AsyncCopy);
      if (completion >= Atomic64Fetch(&slice->atlas_copy_completion_target))
      {
        slice_ready = 1;
      }
    }

    // Fill tex info
    if (slice_ready)
    {
      result.tex_dims = Vec2FromVec(slice->atlas->dims);
      result.tex = slice->atlas->tex;
      result.tex_rect = RNG2(slice->atlas_pos, AddVec2(slice->atlas_pos, DimsFromRng2(slice->canvas_rect)));
    }
    else
    {
      result.tex_dims = SPR.unready_tex_dims;
      result.tex = SPR.unready_tex;
      result.tex_rect = RNG2(VEC2(0, 0), SPR.unready_tex_dims);
    }

    // Fill rays
    StaticAssert(countof(result.rays) == countof(slice->rays));
    CopyStructs(result.rays, slice->rays, countof(result.rays));

    result.matched = span_matched;
    result.ready = slice_ready;
  }

  return result;
}

////////////////////////////////////////////////////////////
//~ Async

void SPR_TickAsync(WaveLaneCtx *lane, AsyncFrameLaneCtx *base_async_lane_frame)
{
  Arena *perm = PermArena();
  SPR_AsyncCtx *async = &SPR.async;
  Arena *frame_arena = base_async_lane_frame->arena;

  // TODO: Distribute rasterization accross wave
  if (lane->idx == 0)
  {
    i64 cmds_count = 0;
    SPR_Cmd *cmds = 0;
    {
      Lock lock = LockE(&SPR.submit.mutex);
      if (SPR.submit.count > 0)
      {
        cmds_count = SPR.submit.count;
        cmds = PushStructsNoZero(frame_arena, SPR_Cmd, cmds_count);
        i64 cmd_idx = 0;
        for (SPR_CmdNode *n = SPR.submit.first; n; n = n->next)
        {
          cmds[cmd_idx] = n->cmd;
          ++cmd_idx;
        }
        // Reset submission queue
        SPR.submit.first_free = SPR.submit.first;
        SPR.submit.first = 0;
        SPR.submit.last = 0;
        SPR.submit.count = 0;
      }
      Unlock(&lock);
    }

    if (cmds_count > 0)
    {
      G_CommandListHandle cl = G_PrepareCommandList(G_QueueKind_AsyncCopy);
      for (i64 cmd_idx = 0; cmd_idx < cmds_count; ++cmd_idx)
      {
        TempArena scratch = BeginScratchNoConflict();
        SPR_Cmd *cmd = &cmds[cmd_idx];
        SPR_SheetEntry *sheet = cmd->sheet;
        ASE_Meta meta = sheet->meta;

        // String encoded = DataFromResource(sheet->key.r);
        // String name = NameFromResource(sheet->key.r);
        // LogInfoF("Rasterizing sprite sheet %F \"%F\" (%F bytes)", FmtHandle(sheet->key.r), FmtString(name), FmtUint(encoded.len));

        SPR_SliceEntry *slice = &sheet->slices[cmd->slice_idx];
        Vec2 slice_dims = DimsFromRng2(slice->canvas_rect);

        //////////////////////////////
        //- Composite

        ASE_Image composite = ASE_PushBlankImage(frame_arena, slice->canvas_rect);
        for (ASE_Layer *ase_layer = sheet->meta.first_layer; ase_layer; ase_layer = ase_layer->next)
        {
          SPR_LayerKind kind = SPR_LayerKindFromName(ase_layer->name);
          if (kind == SPR_LayerKind_Visual)
          {
            ASE_Cel *cel = &ase_layer->cels[cmd->slice_idx];
            {
              TempArena temp = BeginTempArena(scratch.arena);
              {
                // TODO: Reuse decompressed images for linked cels
                ASE_Image image = ASE_DecompressImageFromCel(temp.arena, cel);
                ASE_BlendImage(image, composite);
              }
              EndTempArena(temp);
            }
          }
        }

        //////////////////////////////
        //- Write atlas

        // TODO: Use a more efficient atlas packing algorithm for less wasted space
        SPR_Atlas *atlas = SPR.first_atlas;
        b32 can_use_atlas = 0;
        Vec2I32 atlas_pos = Zi;
        while (can_use_atlas == 0)
        {
          // Create atlas
          if (!atlas)
          {
            atlas = PushStruct(perm, SPR_Atlas);
            i32 atlas_size = MaxI32(1024, NextPow2U64(MaxI32(slice_dims.x, slice_dims.y)));
            atlas->dims = VEC2I32(atlas_size, atlas_size);
            {
              G_ArenaHandle gpu_perm = G_PermArena();
              atlas->tex_res = G_PushTexture2D(
                gpu_perm, cl,
                G_Format_R8G8B8A8_Unorm_Srgb,
                atlas->dims,
                .name = Lit("Sprite atlas")
              );
              atlas->tex = G_PushTexture2DRef(gpu_perm, atlas->tex_res);
            }
            SllStackPush(SPR.first_atlas, atlas);
            ++SPR.atlases_count;
          }
          if (atlas->cur_pos.x + slice_dims.x > atlas->dims.x)
          {
            atlas->cur_pos.x = 0;
            atlas->cur_pos.y += atlas->cur_row_height;
            atlas->cur_row_height = 0;
          }
          if (atlas->cur_pos.x + slice_dims.x < atlas->dims.x && atlas->cur_pos.y + slice_dims.y < atlas->dims.y)
          {
            atlas_pos = atlas->cur_pos;
            atlas->cur_row_height = MaxI32(atlas->cur_row_height, slice_dims.y);
            atlas->cur_pos.x += slice_dims.x;
            can_use_atlas = 1;
          }
          else
          {
            // New atlas required
            atlas = 0;
          }
        }

        // Fill slice_entry atlas info
        {
          slice->atlas = atlas;
          slice->atlas_pos = Vec2FromVec(atlas_pos);
        }

        // Copy to atlas
        G_CopyCpuToTexture(
          cl,
          atlas->tex_res, VEC3I32(atlas_pos.x, atlas_pos.y, 0),
          composite.pixels, VEC3I32(slice_dims.x, slice_dims.y, 1),
          RNG3I32(
            VEC3I32(0, 0, 0),
            VEC3I32(slice_dims.x, slice_dims.y, 1)
          )
        );
        EndScratch(scratch);
      }
      i64 completion_target = G_CommitCommandList(cl);

      //////////////////////////////
      //- Update completion targets

      for (i64 cmd_idx = 0; cmd_idx < cmds_count; ++cmd_idx)
      {
        SPR_Cmd *cmd = &cmds[cmd_idx];
        SPR_SliceEntry *slice = &cmd->sheet->slices[cmd->slice_idx];
        Atomic64Set(&slice->atlas_copy_completion_target, completion_target);
      }
    }
  }
}