// OCR in the app: an engine made on a background thread at startup (modules/OCR), an OCR window that
// says whether it runs on the GPU and, if not, what to install, and the OCR button in the Capture
// window, which reads the capture shown and draws its lines over it.
//
// Engines take a moment to make and OCR takes a moment to run, so both happen on one background
// thread; the UI only looks at what it has finished.

OCR_WINDOW              :: "OCR";
OCR_DEFAULT_MODEL       :: "ppocr-v5-mobile";

ocr_ui: struct {
  models: [] Ocr_Model_Info;
  models_error: string;
  model_index: s32;              // the model chosen in the window

  // Owned by the main thread while thread_busy is false; by the background thread while it is true.
  engine: *Ocr_Engine;
  engine_model: string;          // the model engine runs
  engine_error: string;          // why the last engine could not be made
  thread: Thread;
  thread_started: bool;
  thread_busy: bool;             // atomic: set by the main thread, cleared by the background one
  task: Ocr_Task;

  // The last capture read, for the Capture window.
  lines: [] Ocr_Line;
  lines_index: s32 = -1;         // the capture they are for
  lines_error: string;
  lines_seconds: float64;
  show_boxes := true;

  window_open := false;
  toast_shown := false;
}

Ocr_Task_Kind :: enum { NONE; CREATE; READ; }

Ocr_Task :: struct {
  kind: Ocr_Task_Kind;
  model: string;
  path: string;                  // READ: the .mcap
  index: s32;                    // READ: its capture index
}

ocr_models_folder :: () -> string {
  return tprint("%ocr-models", path_strip_filename(get_path_of_running_executable()));
}

ocr_ui_startup :: () {
  push_allocator(context.default_allocator);
  models, error := ocr_list_models(ocr_models_folder());
  ocr_ui.models = models;
  if error  ocr_ui.models_error = copy_string(error);
  for ocr_ui.models  if it.name == OCR_DEFAULT_MODEL  ocr_ui.model_index = cast(s32) it_index;
  if ocr_ui.models.count  ocr_ui_start(.{kind = .CREATE, model = ocr_ui.models[ocr_ui.model_index].name});
}

ocr_ui_shutdown :: () {
  if ocr_ui.thread_started  thread_is_done(*ocr_ui.thread, -1);   // finishes what it is doing
  ocr_destroy(ocr_ui.engine);
}

// Starts task on the background thread; false while it is still busy with the last one.
ocr_ui_start :: (task: Ocr_Task) -> bool {
  if ocr_ui.thread_busy  return false;
  if ocr_ui.thread_started  thread_is_done(*ocr_ui.thread, -1);   // finished: thread_busy is clear
  push_allocator(context.default_allocator);
  ocr_ui.task = task;
  ocr_ui.task.model = copy_string(task.model);
  ocr_ui.task.path  = copy_string(task.path);
  ocr_ui.thread_busy = true;
  if ocr_ui.thread_started  thread_deinit(*ocr_ui.thread);
  ocr_ui.thread_started = thread_init(*ocr_ui.thread, ocr_ui_thread_proc, temporary_storage_size = 4 * 1024 * 1024, name = "ocr");
  if !ocr_ui.thread_started { ocr_ui.thread_busy = false; return false; }
  thread_start(*ocr_ui.thread);
  return true;
}

ocr_ui_thread_proc :: (thread: *Thread) -> s64 {
  context.allocator = context.default_allocator;
  task := ocr_ui.task;
  if task.kind == {
    case .CREATE;
      ocr_destroy(ocr_ui.engine);
      ocr_ui.engine = null;
      engine, ok, error := ocr_create(ocr_models_folder(), task.model);
      ocr_ui.engine = engine;
      ocr_ui.engine_model = task.model;
      ocr_ui.engine_error = ifx ok then "" else copy_string(error);
      if ok  log("[OCR] % on %.%", task.model, engine.provider, ifx engine.gpu_explanation then tprint(" %", engine.gpu_explanation) else "");
      else   log_error("[OCR] Could not start %: %", task.model, error);

    case .READ;
      start := current_time_monotonic();
      lines, error := ocr_ui_read_capture(task.path);
      ocr_ui.lines = lines;
      ocr_ui.lines_error = error;
      ocr_ui.lines_index = task.index;
      ocr_ui.lines_seconds = to_float64_seconds(current_time_monotonic() - start);
  }
  free(task.model);
  free(task.path);
  atomic_swap(*ocr_ui.thread_busy, false);
  return 0;
}

// The lines in a capture's image as the Capture window shows it: the rendered one, annotations and
// all, or the original when there is none. Heap.
ocr_ui_read_capture :: (path: string) -> [] Ocr_Line, error: string {
  if !ocr_ui.engine  return .[], copy_string("No OCR engine is running.");
  data, read_ok := read_entire_file(path, log_errors = false,, temp);
  if !read_ok  return .[], sprint("Cannot read %.", path);
  m, ok, error := mcap_parse(cast([] u8) data);
  if !ok  return .[], copy_string(error);
  png := mcap_rendered(m);
  if !png.count  png = m.original;
  w, h, channels: s32;
  pixels := stbi_load_from_memory(png.data, cast(s32) png.count, *w, *h, *channels, 4);
  if !pixels  return .[], copy_string("Its image cannot be decoded.");
  defer stbi_image_free(pixels);
  images := Ocr_Image.[.{pixels, w, h, w * 4}];
  results := ocr_process(ocr_ui.engine, images);
  defer array_free(results);
  if results[0].failed  return .[], results[0].error;
  return results[0].lines, "";
}

ocr_ui_free_lines :: () {
  for ocr_ui.lines  free(it.text);
  array_free(ocr_ui.lines);
  free(ocr_ui.lines_error);
  ocr_ui.lines = .[];
  ocr_ui.lines_error = "";
  ocr_ui.lines_index = -1;
}

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

// Each frame: the toast once the engine is up on the CPU though a GPU could run it, and the OCR window.
ocr_ui_frame :: () { using ImGui;
  busy := ocr_ui.thread_busy;
  if !busy && ocr_ui.engine && !ocr_ui.toast_shown {
    ocr_ui.toast_shown = true;
    in_use, explanation := ocr_gpu_status(ocr_ui.engine);
    if !in_use && contains(explanation, "can run it once") {
      toast_show(.Error, "OCR runs on the CPU: CUDA libraries are missing. See the OCR window for what to install.");
      ocr_ui.window_open = true;
    }
  }

  if !ocr_ui.window_open  return;
  SetNextWindowSize(.{640, 360}, .FirstUseEver);
  if Begin(OCR_WINDOW, *ocr_ui.window_open) {
    if ocr_ui.models_error {
      TextColored(app_theme.palette.warning, "%", ocr_ui.models_error);
      TextWrapped("Run `jai fetch_models.jai` in modules/OCR and `jai fetch.jai` in modules/onnxruntime, then rebuild.");
    } else {
      // The model: changing it makes a new engine on the background thread.
      preview := ifx ocr_ui.models.count then ocr_ui.models[ocr_ui.model_index].title else "";
      SetNextItemWidth(320);
      if BeginCombo("Model", temp_c_string(preview)) {
        for ocr_ui.models {
          if Selectable(temp_c_string(it.title), it_index == ocr_ui.model_index) && it_index != ocr_ui.model_index && !busy {
            ocr_ui.model_index = cast(s32) it_index;
            ocr_ui.toast_shown = true;   // said once already, or not needed
            ocr_ui_free_lines();
            ocr_ui_start(.{kind = .CREATE, model = it.name});
          }
        }
        EndCombo();
      }

      if busy && ocr_ui.task.kind == .CREATE {
        TextColored(app_theme.palette.dim, "Starting % ...", ocr_ui.task.model);
      } else if ocr_ui.engine_error {
        TextColored(app_theme.palette.warning, "OCR could not start:");
        TextWrapped("%", ocr_ui.engine_error);
      } else if ocr_ui.engine && !busy {
        in_use, explanation := ocr_gpu_status(ocr_ui.engine);
        Text("Running on: %", ifx in_use then "the GPU (CUDA)" else "the CPU");
        SameLine();
        TextColored(app_theme.palette.dim, "onnxruntime %, % worker threads", ort.version, ocr_ui.engine.workers);
        if explanation { Separator(); text_with_links(explanation); }
      }

      // The text of the capture last read, in reading order.
      if !busy && ocr_ui.lines_index >= 0 && ocr_ui.lines_index < viewer.captures.count {
        Separator();
        Text("Text of %:", viewer.captures[ocr_ui.lines_index].name);
        BeginChild("ocr text");
        ocr_ui_capture_text(ocr_ui.lines_index);
        EndChild();
      }
    }
  }
  End();
}

// The OCR row of the Capture window, above its image: the button, and how the last read went.
ocr_ui_capture_controls :: (index: s32, path: string) { using ImGui;
  busy := ocr_ui.thread_busy;
  can_read := ocr_ui.engine && !busy;
  if !can_read  BeginDisabled();
  if SmallButton("OCR") {
    ocr_ui_free_lines();
    ocr_ui_start(.{kind = .READ, path = path, index = index});
  }
  if !can_read  EndDisabled();
  SameLine();
  if busy && ocr_ui.task.kind == .READ && ocr_ui.task.index == index {
    TextColored(app_theme.palette.dim, "Reading...");
  } else if !busy && ocr_ui.lines_index == index {
    if ocr_ui.lines_error  TextColored(app_theme.palette.warning, "%", ocr_ui.lines_error);
    else                   TextColored(app_theme.palette.dim, "% lines in % s", ocr_ui.lines.count, FF(ocr_ui.lines_seconds, 2));
    SameLine();
    Checkbox("Boxes", *ocr_ui.show_boxes);
  } else if !ocr_ui.engine {
    TextColored(app_theme.palette.dim, "(OCR is not running; see the OCR window)");
  }
  SameLine();
  if SmallButton("OCR window")  ocr_ui.window_open = true;
}

// The lines read from this capture, as boxes over its image, drawn at image_min with scale; and their
// text in a list below the image area.
ocr_ui_capture_overlay :: (index: s32, image_min: ImGui.ImVec2, scale: float) { using ImGui;
  if ocr_ui.thread_busy || ocr_ui.lines_index != index || !ocr_ui.show_boxes  return;
  draw := GetWindowDrawList();
  colour := GetColorU32(.PlotHistogram);
  for ocr_ui.lines {
    p0 := ImVec2.{image_min.x + it.x * scale, image_min.y + it.y * scale};
    p1 := ImVec2.{p0.x + it.width * scale, p0.y + it.height * scale};
    draw.AddRect(draw, p0, p1, colour, 0, 1.5);
    if IsMouseHoveringRect(p0, p1)  SetTooltip("% (%)", it.text, FF(it.confidence, 2));
  }
}

// Text a line at a time, each https:// address in it drawn as a link that opens in the browser.
text_with_links :: (text: string) { using ImGui;
  for line: split(text, "\n",, temp) {
    start := find_index_from_left(line, "https://");
    if start < 0 { TextWrapped("%", line); continue; }
    end := start;
    while end < line.count && line[end] != #char " " && line[end] != #char ")" && line[end] != #char ","  end += 1;
    url := slice(line, start, end - start);
    if start > 0 { TextUnformatted(slice(line, 0, start)); SameLine(0, 0); }
    TextLinkOpenURL(temp_c_string(url));
    if IsItemHovered()  SetTooltip("Open % in the browser", url);
    if end < line.count { SameLine(0, 0); TextUnformatted(advance(line, end)); }
  }
}

ocr_ui_capture_text :: (index: s32) { using ImGui;
  if ocr_ui.thread_busy || ocr_ui.lines_index != index || !ocr_ui.lines.count  return;
  for ocr_ui.lines  TextUnformatted(it.text);
}