internal
Byte_Range get_section_range(const App_State &app, Section_Id sec)
{
  switch (sec) {
    default: return { 0, 0 };
    case Sec_TFile_Header:   return Byte_Range { 0, app.tfile_data.root_file_header_size };
    case Sec_TFile_Object:   return app.tfile_data.rng_root_file_obj;
    case Sec_TFile_Info:     return app.tfile_data.rng_root_file_info;
    case Sec_TFile_FreeList: return app.tfile_data.rng_root_file_free;
    case Sec_TKey_List:      return app.rndata.rng_tkeys_list;
    case Sec_RNTuple_Anchor: return app.rndata.rng_anchor;
    case Sec_RNTuple_Header: return app.rndata.rng_header;
    case Sec_RNTuple_Footer: return app.rndata.rng_footer;
  }
}

internal
void accum_dt_ms(Delta_Time_Accum &accum, f32 dt)
{
	if (accum.count < accum.max) {
		assert(accum.start == 0);
		accum.base[accum.count++] = dt;
	} else {
		accum.base[accum.start++] = dt;
		if (accum.start == accum.max)
			accum.start = 0;
	}
}

internal
f32 calc_avg_dt_ms(const Delta_Time_Accum &accum)
{
	f32 res = 0;
	for (u16 idx = 0; idx < accum.count; ++idx)
		res += accum.base[idx];
	if (accum.count) res /= accum.count;
	return res;
}

internal
String8 to_pretty_size(Arena *arena, u64 bytes)
{
	if (bytes >= GiB(1)) return push_str8f(arena, "%.1f GiB", (f32)bytes / GiB(1));
	if (bytes >= MiB(1)) return push_str8f(arena, "%.1f MiB", (f32)bytes / MiB(1));
	if (bytes >= KiB(1)) return push_str8f(arena, "%.1f KiB", (f32)bytes / KiB(1));
	return push_str8f(arena, "%zu B", bytes);
}

internal
Section find_section(App_State &app, u64 off)
{
  const RNTuple_Data &rdata = app.rndata;
  const TFile_Data &tdata = app.tfile_data;
  u64 rblob_sz = rdata.rblob_header_size; // @Incomplete
  i64 hilite_cluster = app.viewer.highlight_cluster ? app.viewer.highlighted_cluster : -1;
  b8 hilite = false;

  // TFile start
  if (off <= tdata.root_file_header_size) return { Sec_TFile_Header, { 0, tdata.root_file_header_size }, 0, hilite };
  if (tdata.rng_root_file_obj.start <= off && off < tdata.rng_root_file_obj.end())   return { Sec_TFile_Object, tdata.rng_root_file_obj, 0, hilite };
  if (tdata.rng_root_file_info.start <= off && off < tdata.rng_root_file_info.end()) return { Sec_TFile_Info, tdata.rng_root_file_info, 0, hilite };
  if (tdata.rng_root_file_free.start <= off && off < tdata.rng_root_file_free.end()) return { Sec_TFile_FreeList, tdata.rng_root_file_free, 0, hilite };

  /// Handle pages
  {
    // fast case: `off` is in the same page info as previous `off`.
    if (app.last_pinfo->range.start < off && off < app.last_pinfo->range.end()) {
      hilite = hilite_cluster >= 0 && app.last_pinfo->cluster_id == (u64)hilite_cluster;
      return { Sec_Page, app.last_pinfo->range, app.last_pinfo->checksum_size(), hilite };
    }

    // still fast case: `off is in the next page info as the previous.
    if (app.last_pinfo->next) // don't check if it's checksum, since it's the first byte of the page
      app.last_pinfo = app.last_pinfo->next;
    if (app.last_pinfo && app.last_pinfo->range.start <= off && off < app.last_pinfo->range.end()) {
      hilite = hilite_cluster >= 0 && app.last_pinfo->cluster_id == (u64)hilite_cluster;
      return { Sec_Page, app.last_pinfo->range, app.last_pinfo->checksum_size(), hilite };
    }
  }

  if (rdata.rng_anchor_key.start <= off && off < rdata.rng_anchor.end())
    return { Sec_RNTuple_Anchor, rdata.rng_anchor, 8, hilite };

  if (rdata.rng_header.start - rblob_sz <= off && off < rdata.rng_header.end()) 
    return { Sec_RNTuple_Header, rdata.rng_header, 8, hilite };

  if (rdata.rng_footer.start - rblob_sz <= off && off < rdata.rng_footer.end())
    return { Sec_RNTuple_Footer, rdata.rng_footer, 8, hilite };

  if (rdata.rng_tkeys_list.start <= off && off < rdata.rng_tkeys_list.end()) 
    return { Sec_TKey_List, rdata.rng_tkeys_list, 0, hilite };

  // @Speed
  for (u64 cg_idx = 0; cg_idx < rdata.n_cluster_groups; ++cg_idx) {
    Cluster_Group_Info &cg_info = rdata.cluster_groups[cg_idx];
    if (cg_info.rng_page_list.start - rblob_sz <= off && off < cg_info.rng_page_list.end())
      return { Sec_Page_List, cg_info.rng_page_list, 8, hilite };
  }

  // Slow page group lookup, ideally only done once per render when last_pinfo is invalid.
  for (Page_Info_Chunk *chunk = rdata.page_chunks; chunk; chunk = chunk->next) {
    // If we're at the start of a chunk, return a fake Sec_Page used to highlight the RBlob header bytes. 
    if (chunk->range.start - rblob_sz <= off && off < chunk->range.start)
      return { Sec_Page, { chunk->range.start, 0 }, 0, hilite };

    if (chunk->range.start <= off && off < chunk->range.end()) {
      for (u64 group_idx = chunk->first_group; group_idx < rdata.n_page_groups; ++group_idx) {
        const Page_Info_Group &group = rdata.page_groups[group_idx];
        if (off < group.range.start || off >= group.range.end())
          continue;

        for (Page_Info_Node *pinfo = group.first; pinfo; pinfo = pinfo->next) {
          if (pinfo->range.start <= off && off < pinfo->range.end()) {
            app.last_pinfo = pinfo;
            hilite = hilite_cluster >= 0 && pinfo->cluster_id == (u64)hilite_cluster;
            return { Sec_Page, pinfo->range, pinfo->checksum_size(), hilite };
          }
        }
      }

      fprintf(stderr, "Offset 0x%lX is in chunk 0x%lX - 0x%lX, but found in no page_info range!\n", 
              off, chunk->range.start, chunk->range.end());
      assert(false);
    }
  }

  return {};
}

internal
ImColor imcol(const f32 col[3])
{
  return ImColor(col[0], col[1], col[2]);
}

internal
u32 mem_edit_bg_color_fn(const u8 *, u64 off, void *user_data)
{
  App_State *app = reinterpret_cast<App_State *>(user_data);
  off += app->viewer.base_display_addr;

  Section section = find_section(*app, off);
  
  if (section.highlighted) return imcol(app->viewer.col_highlight);
  if (section.id == Sec_Page && off == section.range.start) return imcol(app->viewer.col_page_start);
  if (off < section.range.start) return imcol(app->viewer.col_key);
  if (section.range.end() - section.post_size <= off && off < section.range.end()) return imcol(app->viewer.col_checksum);
  return imcol(app->viewer.col_section[section.id]);
}

internal
void mem_edit_interact_fn(const u8 *, u64 off, void *user_data)
{
}

internal
MemoryEditor make_memory_editor(App_State &app, i32 n_cols)
{
  MemoryEditor mem_edit;
  mem_edit.Cols = n_cols ? n_cols : 32;
  mem_edit.OptShowDataPreview = true;
  // Do nothing on write.
  // Note that we don't use ReadOnly = true because that disables selecting bytes
  mem_edit.WriteFn = [] (ImU8*, size_t, ImU8) {};
  mem_edit.BgColorFn = mem_edit_bg_color_fn;
  mem_edit.BgColorFnUserData = &app;
  mem_edit.InteractFn = mem_edit_interact_fn;
  mem_edit.InteractFnUserData = &app;
  return mem_edit;
}

internal
void make_viewer(App_State &app, u16 n_cols)
{
  Viewer viewer {};
  viewer.mem_edit = make_memory_editor(app, (i32)n_cols);

#define COL(c, r, g, b) viewer.c[0] = r/255.0, viewer.c[1] = g/255.0, viewer.c[2] = b/255.0
  COL(col_key,    0, 100, 50);
  COL(col_page_start, 200, 0, 200);
  COL(col_checksum, 134, 65, 25);
  COL(col_highlight, 190, 190, 190);
#define COL_S(c, r, g, b) viewer.col_section[c][0] = r/255.0, viewer.col_section[c][1] = g/255.0, viewer.col_section[c][2] = b/255.0
  COL_S(Sec_RNTuple_Anchor, 150, 150, 0);
  COL_S(Sec_RNTuple_Header, 150, 0, 50);
  COL_S(Sec_RNTuple_Footer, 50, 0, 150);
  COL_S(Sec_TFile_Header,  90, 90, 90);
  COL_S(Sec_TFile_Object,  120, 120, 120);
  COL_S(Sec_TFile_Info, 95, 76, 76);
  COL_S(Sec_TFile_FreeList, 60, 60, 90);
  COL_S(Sec_Page,   125, 0, 125);
  COL_S(Sec_Page_List, 60, 110, 120);
  COL_S(Sec_TKey_List, 100, 140, 100);
#undef COL
#undef COL_S

  app.viewer = viewer;
}

internal
void viewer_jump_to(Viewer &viewer, u64 addr)
{
  viewer.base_display_addr = addr;
  viewer.mem_edit.GotoAddr = 0;
}

internal
void viewer_jump_to_page(App_State &app, u64 page_idx)
{
  assert(app.rndata.n_pages > 0);
  page_idx = (page_idx + app.rndata.n_pages) % app.rndata.n_pages;
  
  // @Speed
  Page_Info_Node *page = app.rndata.pages;
  for (u64 i = 0; i < page_idx; ++i) {
    page = page->next;    
    assert(page);
  }

  app.viewer.latest_page_gone_to = page_idx;
  viewer_jump_to(app.viewer, page->range.start);
}

internal
void viewer_jump_to_page_list(App_State &app, u64 page_list_idx)
{
  assert(app.rndata.n_cluster_groups > 0);
  page_list_idx = (page_list_idx + app.rndata.n_cluster_groups) % app.rndata.n_cluster_groups;
  
  Cluster_Group_Info &cg_info = app.rndata.cluster_groups[page_list_idx];

  app.viewer.latest_page_list_gone_to = page_list_idx;
  viewer_jump_to(app.viewer, cg_info.rng_page_list.start);
}

internal
void viewer_jump_to_cluster(App_State &app, u64 cluster_idx)
{
  assert(app.rndata.n_clusters > 0);
  cluster_idx = (cluster_idx + app.rndata.n_clusters) % app.rndata.n_clusters;

  // @Speed: this is slow! Consider an acceleration structure, or maybe we can reuse
  // Page_Info_Groups + binary search? (depends on whether cluster_idx are sorted)
  Page_Info_Node *page = app.rndata.pages;
  for (u64 i = 0; i < app.rndata.n_pages; ++i) {
    if (page->cluster_id == cluster_idx)
      break;
    page = page->next;    
    assert(page);
  }

  app.viewer.highlighted_cluster = cluster_idx;
  viewer_jump_to(app.viewer, page->range.start);
}

internal
void update_and_render(Arena *arena, App_State &app, f32 delta_time_ms)
{
  accum_dt_ms(app.delta_time_accum, delta_time_ms);
  
  ImGui::SetNextWindowPos({ 0, 0 });
  ImGui::SetNextWindowSize({ (f32)app.win_data.width, (f32)app.win_data.height });

  Temp scratch = scratch_begin(&arena, 1);
  defer { scratch_end(scratch); };

  const auto main_win_flags = ImGuiWindowFlags_NoResize | ImGuiWindowFlags_NoTitleBar | ImGuiWindowFlags_NoDecoration;
  if (ImGui::Begin("main", nullptr, main_win_flags)) {
        
    String8 ntpl_desc = rntuple_description(scratch.arena, app.rndata);
    ImGui::Text("RNTuple '%s' (%s) from file \"%s\"", app.ntpl_name.c(), ntpl_desc.c(), app.inspected_file.name.c());

    // Draw stats
    {
      ImGui::SameLine();

      f32 avg_dt = calc_avg_dt_ms(app.delta_time_accum);

      String8 mem_used = to_pretty_size(scratch.arena, arena->mem_used);
      String8 mem_peak = to_pretty_size(scratch.arena, arena->mem_peak_used);
      String8 stat_txt = push_str8f(scratch.arena, "mem used: %s (peak: %s) | avg dt: %.1f", 
                                    mem_used.c(), mem_peak.c(), avg_dt);

      f32 pos_x = (ImGui::GetCursorPosX() + ImGui::GetColumnWidth() - ImGui::CalcTextSize(stat_txt.c()).x 
        - ImGui::GetScrollX() - 2 * ImGui::GetStyle().ItemSpacing.x);
      if (pos_x > ImGui::GetCursorPosX())
        ImGui::SetCursorPosX(pos_x);

      ImGui::Text("%s", stat_txt.c());
    }

    ImGui::Separator();

    // Draw main content
    {
      const i64 step_i64 = 1;

      ImGui::BeginTable("Hex View", 2, ImGuiTableFlags_Resizable);

      ImGui::TableNextColumn();

      assert(app.viewer.base_display_addr < app.inspected_file.size);
      void *content = app.inspected_file.mem + app.viewer.base_display_addr;
      u64 content_size = app.inspected_file.size - app.viewer.base_display_addr;
      app.last_pinfo = &invalid_pinfo;
      app.viewer.mem_edit.DrawContents(content, content_size, app.viewer.base_display_addr);

      ImGui::TableNextColumn();
      ImGuiColorEditFlags flags =  ImGuiColorEditFlags_NoInputs|ImGuiColorEditFlags_NoLabel;

      // Unique sections: just display a button that jumps to the start of it and show their size
      for (u32 i = 0; i < Sec_COUNT; ++i) {
        Byte_Range range = get_section_range(app, static_cast<Section_Id>(i));
        if (!range.len) continue;

        String8 sec_name = section_names[i];
        String8 col_label = push_str8f(scratch.arena, "_%s", sec_name.c());
        ImGui::ColorEdit3(col_label.c(), app.viewer.col_section[i], flags);
        ImGui::SameLine(); 
        if (ImGui::Button(sec_name.c())) 
          viewer_jump_to(app.viewer, range.start);
        ImGui::SameLine(); 
        ImGui::Text("%s", to_pretty_size(scratch.arena, range.len).c());
      }

      // Repeated sections: allow jumping to the N-th
      ImGui::ColorEdit3("_Page Start", app.viewer.col_page_start, flags);
      ImGui::SameLine(); 
      if (ImGui::Button("Page Start"))
        viewer_jump_to_page(app, 0);

      ImGui::ColorEdit3("_Page", app.viewer.col_section[Sec_Page], flags);
      ImGui::SameLine();
      if (ImGui::Button("Page"))
        viewer_jump_to_page(app, app.viewer.latest_page_gone_to);
      ImGui::SameLine();
      {
        const i64 step_fast_i64 = app.rndata.n_pages / 100;
        i64 page_to_go_to = app.viewer.latest_page_gone_to;
        ImGui::PushItemWidth(100.f);
        if (ImGui::InputScalar("##page_viewed", ImGuiDataType_S64, &page_to_go_to, &step_i64, &step_fast_i64, "%u"))
          viewer_jump_to_page(app, page_to_go_to);
        ImGui::PopItemWidth();
      }
      ImGui::SameLine();
      ImGui::Text("%s", to_pretty_size(scratch.arena, app.rndata.tot_page_size).c());

      ImGui::ColorEdit3("_Checksum", app.viewer.col_checksum, flags);
      ImGui::SameLine(); 
      if (ImGui::Button("Checksum")) {} // TODO jump to next checksum

      ImGui::ColorEdit3("_Page List", app.viewer.col_section[Sec_Page_List], flags);
      ImGui::SameLine(); 
      if (ImGui::Button("Page List"))      viewer_jump_to_page_list(app, app.viewer.latest_page_list_gone_to);
      ImGui::SameLine();
      {
        i64 page_list_to_go_to = app.viewer.latest_page_list_gone_to;
        ImGui::PushItemWidth(80.f);
        if (ImGui::InputScalar("##page_list_viewed", ImGuiDataType_S64, &page_list_to_go_to, &step_i64, nullptr, "%u"))
          viewer_jump_to_page_list(app, page_list_to_go_to);
        ImGui::PopItemWidth();
      }
      ImGui::SameLine(); 
      ImGui::Text("%s", to_pretty_size(scratch.arena, app.rndata.tot_page_list_size).c());

      // -------------------------------
      ImGui::Separator();
      String8 root_version_str = push_str8f(scratch.arena, "%u.%u.%u", 
                                            app.tfile_data.root_version_major,
                                            app.tfile_data.root_version_minor,
                                            app.tfile_data.root_version_patch);
      ImGui::Text("ROOT version: %s", root_version_str.c());
      ImGui::Text("TFile compression: %u", app.tfile_data.compression);
      ImGui::Text("Num pages: %lu", app.rndata.n_pages);
      ImGui::Text("Num elements: %lu", app.rndata.n_elems);

      {
        const i64 step_fast_i64 = app.rndata.n_clusters / 100;
        i64 cluster_to_highlight = app.viewer.highlighted_cluster;
        ImGui::PushItemWidth(100.f);
        if (ImGui::InputScalar("##highlighted_cluster", ImGuiDataType_S64, &cluster_to_highlight, &step_i64, &step_fast_i64, "%u")) {
          app.viewer.highlight_cluster = true;
          viewer_jump_to_cluster(app, cluster_to_highlight);
        }
        ImGui::PopItemWidth();
        ImGui::SameLine();
        ImGui::Checkbox("Highlight cluster", &app.viewer.highlight_cluster);
      }

      ImGui::EndTable();
    }

    ImGui::End();
  }
}

internal
Term_Viewer make_term_viewer(u16 n_cols)
{
  Term_Viewer viewer {};
  viewer.max_cols = n_cols ? n_cols : 32;
  
  viewer.col_key        = ACol_Green;
  viewer.col_page_start = ACol_Bright_Magenta;
  viewer.col_checksum   = ACol_Yellow;
  viewer.col_highlight  = ACol_Bright_White;
  viewer.col_section[Sec_RNTuple_Anchor] = ACol_Bright_Yellow;
  viewer.col_section[Sec_RNTuple_Header] = ACol_Cyan;
  viewer.col_section[Sec_RNTuple_Footer] = ACol_Blue;
  viewer.col_section[Sec_TFile_Header]   = ACol_White;
  viewer.col_section[Sec_TFile_Object]   = ACol_Grey;
  viewer.col_section[Sec_TFile_Info]     = ACol_Red;
  viewer.col_section[Sec_TFile_FreeList] = ACol_Bright_Yellow;
  viewer.col_section[Sec_Page]           = ACol_Magenta;
  viewer.col_section[Sec_Page_List]      = ACol_Bright_Cyan;
  viewer.col_section[Sec_TKey_List]      = ACol_Bright_Green;
  return viewer;
}

internal
Ansi_Color viewer_to_ansi_color(const Viewer &viewer, const Term_Viewer &tviewer, u32 col)
{
  if (col == imcol(viewer.col_key)) return tviewer.col_key;
  if (col == imcol(viewer.col_page_start)) return tviewer.col_page_start;
  if (col == imcol(viewer.col_checksum)) return tviewer.col_checksum;
  if (col == imcol(viewer.col_highlight)) return tviewer.col_highlight;
  for (u64 section = 0; section < Sec_COUNT; ++section)
    if (col == imcol(viewer.col_section[section]))
      return tviewer.col_section[section];
  return ACol_None;
}

internal
String8 render_legend_to_string(Arena *arena, const Term_Viewer &viewer)
{
  Temp scratch = scratch_begin(&arena, 1);
  defer { scratch_end(scratch); };
  
  struct String8_Node {
    String8_Node *next;
    String8 str;
  } *head = nullptr, *tail = nullptr;
  
  u64 tot_len = 0;
  for (u64 section = 0; section < Sec_COUNT; ++section) {
    Ansi_Color color = viewer.col_section[section];
    String8 color_str = ansi_color_table[color];
    String8 sec_name = section_names[section];
    String8 s = push_str8f(scratch.arena, "%s%s  ", color_str.c(), sec_name.c());
    String8_Node *node = arena_push<String8_Node>(scratch.arena);
    node->str = s;
    if (!tail) {
      head = node;
    } else {
      tail->next = node;
    }
    tail = node;
    tot_len += s.size;
  }

  String8 legend { arena_push_array_nozero<u8>(arena, tot_len + 1), tot_len + 1 };
  legend.str[tot_len] = 0;
  u64 cur_size = 0;

  for (String8_Node *node = head; node; node = node->next) {
    memcpy(legend.str + cur_size, node->str.str, node->str.size);
    cur_size += node->str.size;
  }

  return legend;
} 

internal
String8 render_range_to_string(Arena *arena, App_State &app, u64 len, u64 n_cols)
{
  Term_Viewer viewer = make_term_viewer(n_cols);

  String8 legend = render_legend_to_string(arena, viewer);

  // NOTE: +3 because we need 2 chars for each byte + 1 whitespace.
  u64 buf_size = (ACOL_MAX_LEN + 3) * len;
  u64 n_newlines = len / viewer.max_cols;
  buf_size += n_newlines;
  buf_size += 2; // two newlines
  buf_size += legend.size;
  buf_size += 1; // trailing zero
  u8 *buf = arena_push_array_nozero<u8>(arena, buf_size);
  buf[len] = 0;

  u64 start = app.viewer.base_display_addr;

  // We need to properly initialize this before calling mem_edit_bg_color_fn!
  app.last_pinfo = &invalid_pinfo;

  const u8 *data = app.inspected_file.mem;
  u64 max_addr = app.inspected_file.size;
  assert(start <= max_addr);
  len = min(len, max_addr - start);
  
  Temp scratch = scratch_begin(&arena, 1);
  defer { scratch_end(scratch); };
  
  u8 *cur = buf;
  for (u64 i = 0; i < len; ++i) {
    u64 off = start + i;

    /// Select color
    u32 byte_col = mem_edit_bg_color_fn(data, off, &app);
    // piggyback off mem_edit_bg_color_fn instead of rewriting the color-choosing logic.
    // Kinda dumb code but we don't need speed here since we're printing this one-off.
    Ansi_Color acol = viewer_to_ansi_color(app.viewer, viewer, byte_col);
    String8 acol_str = ansi_color_table[acol];
    memcpy(cur, acol_str.str, acol_str.size);
    cur += acol_str.size;

    /// Write the human-readable byte
    String8 byte_str = push_str8f(scratch.arena, "%02X ", data[off]);
    memcpy(cur, byte_str.str, byte_str.size);
    cur += byte_str.size;

    if ((i + 1) % viewer.max_cols == 0)
      *cur++ = '\n';

    assert((u64)(cur - buf) < (u64)buf_size);
  }

  *cur++ = '\n';
  *cur++ = '\n';
  memcpy(cur, legend.str, legend.size);
  cur += legend.size;

  *cur = 0;

  return { buf, (u64)(cur - buf) };
}