// The capture viewer: a grid of every .mcap's preview, and one capture at full size in a window of
// its own. Scrolling never waits on the disk: previews are read and decoded on worker threads,
// those in view first and then the rows around them, and are uploaded a few per frame. Textures
// for previews far out of view are released once there are more than VIEWER_MAX_TEXTURES.

VIEWER_FOLDER           :: "D:/SnagitIndex/mcap";   // hardcoded for now
VIEWER_WORKER_COUNT     :: 4;
VIEWER_MAX_TEXTURES     :: 1200;   // previews on the GPU at once, about 200 KB each
VIEWER_PREFETCH_ROWS    :: 8;      // most rows decoded ahead of the view, above and below
VIEWER_UPLOAD_BUDGET_MS :: 4.0;    // time per frame spent uploading previews
VIEWER_FULL_WINDOW      :: "Capture";
VIEWER_MIN_ZOOM         :: 5.0;    // percent, for the full-size window
VIEWER_MAX_ZOOM         :: 800.0;

GPU_TEXTURE_MAX_SIDE :: 16384;     // what D3D11 guarantees; larger images are scaled down to fit

// A texture made by gpu_texture_create (graphics/DX11.jai, base/platform/sdl.jai).
Gpu_Texture :: struct {
  handle: *void;   // the backend's own: an ID3D11ShaderResourceView or an SDL_GPUTexture
  width, height: s32;
}

// Bumped when the renderer moves to another device (DX11's WARP-first startup, or device loss),
// which takes every texture with it.
g_texture_generation: u64;

gpu_texture_ref :: (texture: Gpu_Texture) -> ImGui.ImTextureRef {
  return .{null, cast(ImGui.ImTextureID) texture.handle};
}

Viewer_Capture :: struct {
  name: string;                 // the file name without .mcap
  path: string;
  texture: Gpu_Texture;         // its preview, once uploaded
  last_drawn_frame: u64;
  last_wanted_frame: u64;       // last in the rows viewer_request_previews asked for
  failed: bool;
  in_flight: bool;              // under viewer.mutex: queued for a worker, or being decoded
}

// How the full-size window sizes its image.
Viewer_Fit :: enum u8 {
  OFF;      // at the zoom the slider sets
  WINDOW;   // the whole image in the window
  WIDTH;    // the window's width, scrolling down
  HEIGHT;   // the window's height, scrolling across
}

// Pixels a worker decoded, waiting for the main thread to upload them.
Decoded :: struct {
  index: s32;
  pixels: *u8;                  // RGBA; from stb_image, or from the default allocator when scaled down
  scaled: bool;                 // pixels came from the default allocator
  width, height: s32;
  source_width, source_height: s32;   // before any scaling down
  failed: bool;
}

viewer: struct {
  captures: [..] Viewer_Capture;
  scan_error: string;
  thumb_size: float = 200;
  frame: u64;
  texture_generation: u64;
  with_texture: [..] s32;       // captures holding a preview texture

  // Shared with the workers: everything below up to the full view, under mutex.
  mutex: Mutex;
  wake: Semaphore;
  queue: [..] s32;              // previews wanted, nearest the view first
  queue_next: s64;
  done: [..] Decoded;
  full_wanted: s32 = -1;
  full_request: u64;            // bumped for every new full-size request
  full_taken: u64;              // the request a worker last took
  full_done: Decoded;
  full_done_request: u64;       // 0 while there is none
  quit: bool;
  workers: [VIEWER_WORKER_COUNT] Thread;
  workers_started: s64;

  // The full view; main thread only.
  full_open: bool;
  full_index: s32 = -1;         // the capture shown, or being loaded
  full_texture: Gpu_Texture;
  full_source_width, full_source_height: s32;
  full_loading: bool;
  full_failed: bool;
  full_fit := Viewer_Fit.WINDOW;
  full_zoom: float = 100;       // percent of the size it was captured at; set by the fit, unless OFF
}

// ---------------------------------------------------------------- startup and shutdown

viewer_startup :: () {
  push_allocator(context.default_allocator);
  viewer.captures.allocator     = context.default_allocator;
  viewer.with_texture.allocator = context.default_allocator;
  viewer.queue.allocator        = context.default_allocator;
  viewer.done.allocator         = context.default_allocator;

  init(*viewer.mutex, "viewer");
  init(*viewer.wake);

  start := current_time_monotonic();
  files := file_list(VIEWER_FOLDER,, temp);
  for files {
    name := path_filename(it);
    if !ends_with_nocase(name, tprint(".%", MCAP_EXTENSION))  continue;
    capture := array_add(*viewer.captures);
    capture.path = copy_string(it);
    capture.name = copy_string(path_strip_extension(name));
  }
  // Newest first, roughly: dated names (2026-...) sort after the older {GUID} ones.
  quick_sort(viewer.captures, (a: Viewer_Capture, b: Viewer_Capture) -> s64 { return compare(b.name, a.name); });
  if !viewer.captures.count  viewer.scan_error = sprint("No .% files in %.", MCAP_EXTENSION, VIEWER_FOLDER);
  log("[viewer] Found % captures in % in % ms.", viewer.captures.count, VIEWER_FOLDER, FF(to_float64_seconds(current_time_monotonic() - start) * 1000.0, 1));

  for * viewer.workers {
    // Reading a front takes 64 KB of temporary storage, and a full image several MB.
    if !thread_init(it, viewer_worker, temporary_storage_size = 1024 * 1024, name = tprint("viewer %", it_index)) {
      log_error("[viewer] Could not start worker %.", it_index);
      continue;
    }
    thread_start(it);
    viewer.workers_started += 1;
  }
}

// On the shutdown thread: stops the workers. The textures go with the device.
viewer_shutdown :: () {
  lock(*viewer.mutex);
  viewer.quit = true;
  unlock(*viewer.mutex);
  for 1..viewer.workers_started  signal(*viewer.wake);
  for 0..viewer.workers_started - 1  thread_is_done(*viewer.workers[it], -1);
}

// ---------------------------------------------------------------- the workers

viewer_worker :: (thread: *Thread) -> s64 {
  while true {
    wait_for(*viewer.wake, 250);
    while true {
      index: s32 = -1;
      full: bool;
      request: u64;
      path: string;

      lock(*viewer.mutex);
      if viewer.quit { unlock(*viewer.mutex); return 0; }
      if viewer.full_request != viewer.full_taken {
        full = true;
        request = viewer.full_request;
        viewer.full_taken = request;
        index = viewer.full_wanted;
      } else if viewer.queue_next < viewer.queue.count {
        index = viewer.queue[viewer.queue_next];
        viewer.queue_next += 1;
      }
      if index >= 0  path = viewer.captures[index].path;   // the array no longer grows after startup
      unlock(*viewer.mutex);
      if index < 0  break;   // nothing to do: wait to be woken

      result := ifx full then viewer_decode_full(path) else viewer_decode_preview(path);
      result.index = index;

      lock(*viewer.mutex);
      if full {
        if request == viewer.full_request {   // still the one wanted
          viewer_free_pixels(*viewer.full_done);
          viewer.full_done = result;
          viewer.full_done_request = request;
        } else {
          viewer_free_pixels(*result);
        }
      } else {
        array_add(*viewer.done, result);
      }
      unlock(*viewer.mutex);
    }
  }
  return 0;
}

// The preview at the front of the file: only the header, text and preview are read.
viewer_decode_preview :: (path: string) -> Decoded {
  result: Decoded;
  result.failed = true;
  push_allocator(temp);
  defer reset_temporary_storage();

  m, ok, error := mcap_read_front(path);
  if !ok { log_error("[viewer] % %", path, error); return result; }
  if !m.preview.png.count { log_error("[viewer] % has no preview.", path); return result; }
  return viewer_decode_png(path, m.preview.png);
}

// The rendered image, annotations and all, or the original when there is none.
viewer_decode_full :: (path: string) -> Decoded {
  result: Decoded;
  result.failed = true;
  push_allocator(temp);
  defer reset_temporary_storage();

  data, read_ok := read_entire_file(path, log_errors = false);
  if !read_ok { log_error("[viewer] Cannot read %.", path); return result; }
  m, ok, error := mcap_parse(cast([] u8) data);
  if !ok { log_error("[viewer] % %", path, error); return result; }
  png := mcap_rendered(m);
  if !png.count  png = m.original;
  return viewer_decode_png(path, png);
}

viewer_decode_png :: (path: string, png: [] u8) -> Decoded {
  result: Decoded;
  result.failed = true;
  width, height, channels: s32;
  pixels := stbi_load_from_memory(png.data, cast(s32) png.count, *width, *height, *channels, 4);
  if !pixels || width <= 0 || height <= 0 {
    log_error("[viewer] Cannot decode the image in %.", path);
    if pixels  stbi_image_free(pixels);
    return result;
  }
  result.pixels = pixels;
  result.width, result.height = width, height;
  result.source_width, result.source_height = width, height;
  result.failed = false;

  // Taller or wider than a texture can be (long scrolling captures): scaled down to fit.
  longest := max(width, height);
  if longest > GPU_TEXTURE_MAX_SIDE {
    scale := cast(float64) GPU_TEXTURE_MAX_SIDE / longest;
    out_width  := max(cast(s32) 1, cast(s32) (width * scale));
    out_height := max(cast(s32) 1, cast(s32) (height * scale));
    scaled := cast(*u8) alloc(cast(s64) out_width * out_height * 4,, context.default_allocator);
    if stbir_resize_uint8_srgb(pixels, width, height, 0, scaled, out_width, out_height, 0, 4, 3, 0) {
      stbi_image_free(pixels);
      result.pixels = scaled;
      result.scaled = true;
      result.width, result.height = out_width, out_height;
    } else {
      free(scaled,, context.default_allocator);
      stbi_image_free(pixels);
      result.pixels = null;
      result.failed = true;
      log_error("[viewer] Cannot scale down the % x % image in %.", width, height, path);
    }
  }
  return result;
}

viewer_free_pixels :: (d: *Decoded) {
  if d.pixels {
    if d.scaled  free(d.pixels,, context.default_allocator);
    else         stbi_image_free(d.pixels);
  }
  d.pixels = null;
}

// ---------------------------------------------------------------- each frame: GPU work

// Inside the frame, before anything is drawn: the device can only have changed before this point,
// and nothing released here has been drawn yet this frame.
viewer_frame_gpu_work :: () {
  viewer.frame += 1;

  if viewer.texture_generation != g_texture_generation {
    viewer.texture_generation = g_texture_generation;
    // Made on the old device: released, so it can go, and made again as they come into view.
    for viewer.with_texture  gpu_texture_release(*viewer.captures[it].texture);
    array_reset_keeping_memory(*viewer.with_texture);
    if viewer.full_texture.handle {
      gpu_texture_release(*viewer.full_texture);
      viewer_request_full(viewer.full_index);
    }
  }

  // The full-size image, when one has arrived.
  lock(*viewer.mutex);
  full := viewer.full_done;
  arrived := viewer.full_done_request != 0 && viewer.full_done_request == viewer.full_request;
  if viewer.full_done_request != 0 {
    viewer.full_done = .{};
    viewer.full_done_request = 0;
  }
  unlock(*viewer.mutex);
  if arrived {
    gpu_texture_release(*viewer.full_texture);
    viewer.full_loading = false;
    viewer.full_failed  = full.failed;
    if !full.failed {
      viewer.full_texture = gpu_texture_create(full.pixels, full.width, full.height);
      viewer.full_failed  = !viewer.full_texture.handle;
      viewer.full_source_width, viewer.full_source_height = full.source_width, full.source_height;
    }
  }
  viewer_free_pixels(*full);

  // Previews, until the frame's budget is spent; the rest wait for the next frame.
  start := current_time_monotonic();
  while to_float64_seconds(current_time_monotonic() - start) * 1000.0 < VIEWER_UPLOAD_BUDGET_MS {
    lock(*viewer.mutex);
    if !viewer.done.count { unlock(*viewer.mutex); break; }
    d := pop(*viewer.done);
    viewer.captures[d.index].in_flight = false;
    unlock(*viewer.mutex);

    capture := *viewer.captures[d.index];
    if d.failed {
      capture.failed = true;
    } else if !capture.texture.handle {
      capture.texture = gpu_texture_create(d.pixels, d.width, d.height);
      if capture.texture.handle  array_add(*viewer.with_texture, d.index);
      else                       capture.failed = true;
    }
    viewer_free_pixels(*d);
  }

  viewer_evict();
}

// When a preview was last in view or wanted ahead of it. A prefetched preview has not been drawn
// yet; going by the frame it was drawn alone, it would be released first and decoded again at once.
viewer_last_used :: (capture: Viewer_Capture) -> u64 {
  return max(capture.last_drawn_frame, capture.last_wanted_frame);
}

// Past VIEWER_MAX_TEXTURES, the previews used longest ago are released, never one in view or
// wanted around it as of the last frame.
viewer_evict :: () {
  if viewer.with_texture.count <= VIEWER_MAX_TEXTURES  return;
  quick_sort(viewer.with_texture, (a: s32, b: s32) -> s64 {
    fa, fb := viewer_last_used(viewer.captures[a]), viewer_last_used(viewer.captures[b]);
    if fa == fb  return 0;
    return ifx fa < fb then -1 else 1;
  });
  target := VIEWER_MAX_TEXTURES * 3 / 4;
  removed := 0;
  for viewer.with_texture {
    if viewer.with_texture.count - removed <= target  break;
    capture := *viewer.captures[it];
    if viewer_last_used(capture.*) >= viewer.frame - 1  break;   // still wanted: the rest are newer still
    gpu_texture_release(*capture.texture);
    removed += 1;
  }
  for 0..viewer.with_texture.count - removed - 1  viewer.with_texture[it] = viewer.with_texture[it + removed];
  viewer.with_texture.count -= removed;
}

// What the workers should decode next: the rows in view, then the rows around them, nearest first.
viewer_request_previews :: (first_row: s64, last_row: s64, columns: s64) {
  if !viewer.captures.count || columns <= 0  return;
  rows := (viewer.captures.count + columns - 1) / columns;

  // Rows ahead only as far as the textures they need leave half of VIEWER_MAX_TEXTURES spare, so
  // what is wanted always fits and nothing wanted is ever released to make room.
  visible_rows   := last_row - first_row + 1;
  rows_that_fit  := VIEWER_MAX_TEXTURES / 2 / columns - visible_rows;
  prefetch_rows  := clamp(rows_that_fit / 2, 0, VIEWER_PREFETCH_ROWS);

  wanted: [..] s32;
  wanted.allocator = temp;
  add_row :: (wanted: *[..] s32, row: s64, columns: s64) {
    for column: 0..columns - 1 {
      index := row * columns + column;
      if index >= viewer.captures.count  break;
      capture := *viewer.captures[index];
      capture.last_wanted_frame = viewer.frame;   // kept by viewer_evict, texture or not
      if !capture.texture.handle && !capture.failed  array_add(wanted, cast(s32) index);
    }
  }
  for row: max(first_row, 0)..min(last_row, rows - 1)  add_row(*wanted, row, columns);
  for distance: 1..prefetch_rows {
    if last_row + distance < rows    add_row(*wanted, last_row + distance, columns);   // scrolling down is likelier
    if first_row - distance >= 0     add_row(*wanted, first_row - distance, columns);
  }

  added := 0;
  lock(*viewer.mutex);
  // Whatever was queued and not yet taken is let go; what is still wanted is queued again.
  for viewer.queue_next..viewer.queue.count - 1  viewer.captures[viewer.queue[it]].in_flight = false;
  array_reset_keeping_memory(*viewer.queue);
  viewer.queue_next = 0;
  for wanted {
    capture := *viewer.captures[it];
    if capture.in_flight  continue;   // a worker has it
    capture.in_flight = true;
    array_add(*viewer.queue, it);
    added += 1;
  }
  unlock(*viewer.mutex);
  for 1..min(added, VIEWER_WORKER_COUNT)  signal(*viewer.wake);
}

viewer_request_full :: (index: s32) {
  if index < 0 || index >= viewer.captures.count  return;
  viewer.full_index   = index;
  viewer.full_open    = true;
  viewer.full_loading = true;
  viewer.full_failed  = false;
  lock(*viewer.mutex);
  viewer.full_wanted = index;
  viewer.full_request += 1;
  unlock(*viewer.mutex);
  signal(*viewer.wake);
}

// ---------------------------------------------------------------- drawing

viewer_ui :: () { using ImGui;
  viewer_frame_gpu_work();

  if viewer.scan_error { TextColored(app_theme.palette.warning, "%", viewer.scan_error); return; }

  loaded := viewer.with_texture.count;
  Text("% captures in %", viewer.captures.count, VIEWER_FOLDER);
  SameLine();
  TextColored(app_theme.palette.dim, "(% previews loaded)", loaded);
  SetNextItemWidth(200);
  SliderFloat("Size", *viewer.thumb_size, 64, MCAP_PREVIEW_MAX_SIDE, "%.0f px");

  BeginChild("grid");
  viewer_grid();
  EndChild();

  viewer_full_window();
}

viewer_grid :: () { using ImGui;
  style := GetStyle();
  cell  := viewer.thumb_size;
  label := GetTextLineHeight();
  available := GetContentRegionAvail().x;
  columns := max(1, cast(s64) ((available + style.ItemSpacing.x) / (cell + style.ItemSpacing.x)));
  rows    := (viewer.captures.count + columns - 1) / columns;
  row_height := cell + label + style.ItemSpacing.y;

  draw := GetWindowDrawList();
  placeholder := GetColorU32(.FrameBg);
  outline     := GetColorU32(.ButtonHovered);
  selected    := GetColorU32(.CheckMark);
  text_color  := GetColorU32(.Text);

  first_row, last_row: s64 = 0x7FFF_FFFF_FFFF_FFFF, -1;
  clipper: ListClipper;
  clipper.Constructor(*clipper);
  defer clipper.Destructor(*clipper);
  clipper.Begin(*clipper, xx rows, row_height);
  while clipper.Step(*clipper) {
    first_row = min(first_row, cast(s64) clipper.DisplayStart);
    last_row  = max(last_row,  cast(s64) clipper.DisplayEnd - 1);
    for row: clipper.DisplayStart..clipper.DisplayEnd - 1 {
      for column: 0..columns - 1 {
        index := row * columns + column;
        if index >= viewer.captures.count  break;
        capture := *viewer.captures[index];
        if column > 0  SameLine();

        top_left := GetCursorScreenPos();
        PushID(cast(s32) index);
        clicked := InvisibleButton("cell", .{cell, cell + label});
        hovered := IsItemHovered();
        PopID();

        bottom_right := ImVec2.{top_left.x + cell, top_left.y + cell};
        if capture.texture.handle {
          // Fitted inside the square, keeping its shape.
          w, h := cast(float) capture.texture.width, cast(float) capture.texture.height;
          scale := min(cell / w, cell / h);
          size := ImVec2.{w * scale, h * scale};
          offset := ImVec2.{top_left.x + (cell - size.x) / 2, top_left.y + (cell - size.y) / 2};
          draw.AddImage(draw, gpu_texture_ref(capture.texture), offset, .{offset.x + size.x, offset.y + size.y});
          capture.last_drawn_frame = viewer.frame;
        } else {
          draw.AddRectFilled(draw, top_left, bottom_right, placeholder);
          if capture.failed  draw.AddText(draw, .{top_left.x + 6, top_left.y + 6}, text_color, "Cannot read");
        }
        if index == viewer.full_index  draw.AddRect(draw, top_left, bottom_right, selected, 0, 3);
        else if hovered                draw.AddRect(draw, top_left, bottom_right, outline, 0, 2);

        label_min := ImVec2.{top_left.x, top_left.y + cell};
        label_max := ImVec2.{top_left.x + cell, top_left.y + cell + label};
        draw.PushClipRect(draw, label_min, label_max, true);
        draw.AddText(draw, label_min, text_color, capture.name);
        draw.PopClipRect(draw);

        if hovered  SetTooltip("%", capture.name);
        if clicked  viewer_request_full(cast(s32) index);
      }
    }
  }
  clipper.End(*clipper);

  if last_row >= 0  viewer_request_previews(first_row, last_row, columns);
}

viewer_full_window :: () { using ImGui;
  if !viewer.full_open {
    if viewer.full_texture.handle  gpu_texture_release(*viewer.full_texture);   // only while shown
    return;
  }
  SetNextWindowSize(.{1000, 700}, .FirstUseEver);
  if Begin(VIEWER_FULL_WINDOW, *viewer.full_open) {
    capture := *viewer.captures[viewer.full_index];
    Text("%", capture.name);
    if viewer.full_texture.handle && !viewer.full_loading {
      SameLine();
      TextColored(app_theme.palette.dim, "% x %", viewer.full_source_width, viewer.full_source_height);
    }
    if viewer.full_loading { SameLine(); TextColored(app_theme.palette.dim, "Loading..."); }
    TextColored(app_theme.palette.dim, "Left / Right: previous / next");

    // Zoom is a percentage of the size the capture was taken at. A fit sets it from the window, larger
    // or smaller, and the slider shows what it works out to; moving the slider turns the fit off,
    // and the zoom then stays for the next capture.
    fit := viewer.full_fit;
    Text("Fit:");
    SameLine(); if RadioButton("Window", fit == .WINDOW)  fit = .WINDOW;
    SameLine(); if RadioButton("Width",  fit == .WIDTH)   fit = .WIDTH;
    SameLine(); if RadioButton("Height", fit == .HEIGHT)  fit = .HEIGHT;
    SameLine(); if RadioButton("Off",    fit == .OFF)     fit = .OFF;
    viewer.full_fit = fit;
    SameLine();
    SetNextItemWidth(260);
    if SliderFloat("Zoom", *viewer.full_zoom, VIEWER_MIN_ZOOM, VIEWER_MAX_ZOOM, "%.0f%%", .Logarithmic)  viewer.full_fit = .OFF;
    SameLine();
    if SmallButton("100%") { viewer.full_fit = .OFF; viewer.full_zoom = 100; }

    // Left and Right step through the captures in grid order while this window has focus.
    if IsWindowFocused(.RootAndChildWindows) {
      if IsKeyPressed(.LeftArrow)  && viewer.full_index > 0                           viewer_request_full(viewer.full_index - 1);
      if IsKeyPressed(.RightArrow) && viewer.full_index < viewer.captures.count - 1   viewer_request_full(viewer.full_index + 1);
    }

    if viewer.full_failed {
      TextColored(app_theme.palette.warning, "This capture cannot be shown; see the Event Log.");
    } else if viewer.full_texture.handle {
      // Each fit keeps the same scrollbars from frame to frame. Otherwise fitting the width of a tall
      // image would bring up a vertical scrollbar, narrowing the area, changing the fit, and so on.
      flags := WindowFlags.HorizontalScrollbar;
      if viewer.full_fit == {
        case .WINDOW;  flags = .NoScrollbar;
        case .WIDTH;   flags |= .AlwaysVerticalScrollbar;
        case .HEIGHT;  flags |= .AlwaysHorizontalScrollbar;
      }
      BeginChild("image", .{0, 0}, .None, flags);
      area := GetContentRegionAvail();

      // Sized from the captured size, so an image scaled down to fit a texture is drawn back up.
      source_w, source_h := cast(float) viewer.full_source_width, cast(float) viewer.full_source_height;
      if source_w > 0 && source_h > 0 && area.x > 0 && area.y > 0 {
        if viewer.full_fit == {
          case .WINDOW;  viewer.full_zoom = 100 * min(area.x / source_w, area.y / source_h);
          case .WIDTH;   viewer.full_zoom = 100 * area.x / source_w;
          case .HEIGHT;  viewer.full_zoom = 100 * area.y / source_h;
        }
      }
      zoom := clamp(viewer.full_zoom, VIEWER_MIN_ZOOM, VIEWER_MAX_ZOOM) / 100;
      Image(gpu_texture_ref(viewer.full_texture), .{source_w * zoom, source_h * zoom});
      EndChild();
      // The previous image stays up while the next one loads, so stepping through does not flicker.
    }
  }
  End();
}