// record-probe: phase 1 of video-capture-plan.html. Records a screen, or a region of one, with
// modules/Screen_Record (DXGI Desktop Duplication and FFmpeg) and says how it went: which screens
// there are and on which adapter, which duplication call and encoder worked, frames captured, written
// and dropped, time per stage, and the file's size. It also answers what the plan leaves to phase 1:
//
//   -legacy         the plain DuplicateOutput, not DuplicateOutput1: does it fail for a process aware
//                   of each monitor's DPI, as mcap is (this probe is, unless -dpi-unaware)?
//   -cuda           makes a CUDA context first, as onnxruntime does in mcap: does duplication still work?
//   -exclude-test   puts two squares on the region, one kept out of captures with
//                   WDA_EXCLUDEFROMCAPTURE: is it out of the recording, and the other in?
//   -animate        moves a square across the region the whole time, so the screen changes every
//                   frame: for measuring throughput.
//   -check FILE     reads FILE back with FFmpeg and says what it holds, and nothing else: for a
//                   recording cut short by a crash.
//   -snapshot DIR   the last frame captured (captured.png) and a frame decoded back from the file
//                   (decoded.png), to compare by eye.
//
// Build: jai first.jai - record_probe
// Run:   bin/record-probe [-list] [-screen N] [-region X Y W H] [-seconds S] [-fps F] [-encoder NAME]
//                         [-out FILE.mp4] [-no-cursor] [-legacy] [-cuda] [-exclude-test] [-snapshot DIR]
//                         [-animate] [-dpi-unaware]
// Exit code: 0 recorded, 1 failed, 3 Desktop Duplication cannot run here (the fallback's case).

main :: () {
  args := get_command_line_arguments();
  options: Record_Options;
  options.path = tprint("%/record-probe.mp4", get_working_directory());
  seconds: float64 = 10;
  list_only, cuda, exclude_test, dpi_unaware, animate := false, false, false, false, false;
  snapshot_folder: string;
  i := 1;
  number :: (index: s64) -> s32 #expand { return cast(s32) string_to_int(`args[index]); }
  while i < args.count {
    a := args[i];
    more := args.count - i - 1;
    if a == "-list"                     { list_only = true; i += 1; continue; }
    if a == "-screen"   && more >= 1    { options.screen = number(i + 1); i += 2; continue; }
    if a == "-region"   && more >= 4 {
      x, y, w, h := number(i + 1), number(i + 2), number(i + 3), number(i + 4);
      options.region = .{x, y, x + w, y + h};
      i += 5; continue;
    }
    if a == "-seconds"  && more >= 1    { seconds = string_to_float64(args[i + 1]); i += 2; continue; }
    if a == "-fps"      && more >= 1    { options.fps = number(i + 1); i += 2; continue; }
    if a == "-encoder"  && more >= 1    { options.encoder = args[i + 1]; i += 2; continue; }
    if a == "-out"      && more >= 1    { options.path = args[i + 1]; i += 2; continue; }
    if a == "-snapshot" && more >= 1    { snapshot_folder = args[i + 1]; i += 2; continue; }
    if a == "-no-cursor"                { options.cursor = false; i += 1; continue; }
    if a == "-legacy"                   { options.legacy_duplicate = true; i += 1; continue; }
    if a == "-cuda"                     { cuda = true; i += 1; continue; }
    if a == "-exclude-test"             { exclude_test = true; i += 1; continue; }
    if a == "-dpi-unaware"              { dpi_unaware = true; i += 1; continue; }
    if a == "-animate"                  { animate = true; i += 1; continue; }
    if a == "-check"    && more >= 1 {
      image, info, ok, error := video_poster(args[i + 1]);
      if !ok { print("%: cannot be read: %\n", args[i + 1], error); exit(1); }
      print("%: % x %, % s, %, % fps; a frame decoded at % s\n", args[i + 1], info.width, info.height, FF(info.duration, 2), info.codec, FF(info.frame_rate, 1), FF(info.poster_time, 2));
      exit(0);
    }
    print("Unknown argument %\n", a);
    exit(2);
  }

  if !dpi_unaware {
    DPI_AWARENESS_CONTEXT_PER_MONITOR_AWARE_V2 :: cast(HANDLE) -4;
    if !SetProcessDpiAwarenessContext(DPI_AWARENESS_CONTEXT_PER_MONITOR_AWARE_V2)  print("Could not become per-monitor DPI aware (v2).\n");
  }
  print("DPI awareness: %\n", ifx dpi_unaware then "none" else "per monitor (v2), as mcap");

  screens := record_list_screens();
  print("Screens:\n");
  for screens {
    d := it.desktop;
    print("  %: % on adapter % (%): % x % at %, %%\n", it_index, it.name, it.adapter_index, it.adapter,
          d.right - d.left, d.bottom - d.top, d.left, d.top,
          ifx it.rotation == .IDENTITY || it.rotation == .UNSPECIFIED then "" else tprint(", rotated %", it.rotation));
  }
  if list_only  return;

  if cuda {
    ok, text := cuda_make_context();
    print("CUDA: %\n", text);
    if !ok  exit(1);
  }

  recorder: Recorder;
  ok, error, fall_back := record_start(*recorder, options);
  if !ok {
    print("Cannot record: %\n", error);
    if fall_back  print("Desktop Duplication cannot run here: this is the case for the Windows.Graphics.Capture fallback.\n");
    exit(ifx fall_back then cast(s32) 3 else 1);
  }
  stats := record_stats(*recorder);
  print("Recording % x % for % s at up to % fps with % and %, to %\n", stats.width, stats.height, seconds, options.fps, stats.source, stats.encoder, options.path);

  region := options.region;
  if region.right <= region.left  region = screens[ifx options.screen >= 0 then options.screen else 0].desktop;
  excluded, control: HWND;
  if exclude_test  excluded, control = exclusion_windows(region);
  mover: HWND;
  if animate  mover = square_window("record-probe animate", 0x0000A5FF, region.left, region.top, 120);

  start := current_time_monotonic();
  last_print := start;
  while to_float64_seconds(current_time_monotonic() - start) < seconds {
    pump_messages();
    if mover {   // back and forth across the region, about 240 times a second
      elapsed := to_float64_seconds(current_time_monotonic() - start);
      span := region.right - region.left - 120;
      phase := fmod_cycling(elapsed * 0.5, 2.0);
      x := region.left + cast(s32) (span * ifx phase < 1 then phase else 2 - phase);
      SetWindowPos(mover, HWND_TOPMOST, x, region.top + (region.bottom - region.top) / 2 - 60, 0, 0, SWP_NOSIZE | SWP_NOACTIVATE);
      sleep_milliseconds(4);
    } else {
      sleep_milliseconds(20);
    }
    if to_float64_seconds(current_time_monotonic() - last_print) >= 1 {
      last_print = current_time_monotonic();
      s := record_stats(*recorder);
      print("  % s: % captured, % sent (% repeats), % written, % dropped, % lost, % KB\n",
            FF(to_float64_seconds(last_print - start), 0), s.captured, s.emitted, s.repeats, s.written, s.dropped, s.lost, s.bytes / 1024);
    }
    failed, _ := record_failed(*recorder);
    if failed  break;
  }

  last, width, height := record_last_frame(*recorder);
  if mover  DestroyWindow(mover);
  stopped, stop_error := record_stop(*recorder);
  s := recorder.stats;
  print("\n%\n", ifx stopped then "Recorded." else tprint("Recording failed: %", stop_error));
  ms_each :: (ticks: s64, count: s64) -> string #expand {
    if !count  return "-";
    return tprint("% ms", FF(cast(float64) ticks / `recorder.frequency * 1000 / count, 2));
  }
  file_size := file_size_of(options.path);
  print("  source:     %\n", s.source);
  print("  encoder:    %\n", s.encoder);
  print("  size:       % x %\n", s.width, s.height);
  print("  seconds:    %\n", FF(s.seconds, 1));
  print("  captured:   % new images of the screen (% a second), % cursor-only changes\n", s.captured, FF(s.captured / s.seconds, 1), s.pointer_moves);
  print("  sent:       % frames (% a second), % of them repeats of an unchanged screen\n", s.emitted, FF(s.emitted / s.seconds, 1), s.repeats);
  print("  written:    %, dropped %, duplication lost %\n", s.written, s.dropped, s.lost);
  print("  each:       copy % (per image), read back % , encode % (per frame)\n",
        ms_each(s.copy_ticks, s.captured), ms_each(s.readback_ticks, s.emitted), ms_each(s.encode_ticks, s.written));
  print("  file:       % KB, % MB a minute\n", file_size / 1024, FF(cast(float64) file_size / (1024 * 1024) / s.seconds * 60, 1));

  if exclude_test {
    excluded_rect := window_rect_in(excluded, region);
    control_rect  := window_rect_in(control, region);
    excluded_seen := colour_fraction(last, width, height, excluded_rect, EXCLUDED_BGR);
    control_seen  := colour_fraction(last, width, height, control_rect, CONTROL_BGR);
    print("  exclusion:  the excluded square is % in the last frame, the control square %: %\n",
          percent(excluded_seen), percent(control_seen),
          ifx control_seen < 0.9 then "inconclusive (the control is not there)"
          else ifx excluded_seen < 0.01 then "WDA_EXCLUDEFROMCAPTURE is respected"
          else "WDA_EXCLUDEFROMCAPTURE is NOT respected");
    DestroyWindow(excluded);
    DestroyWindow(control);
  }

  if snapshot_folder {
    make_directory_if_it_does_not_exist(snapshot_folder, recursive = true);
    rgba := bgra_to_rgba(last);
    captured_path := tprint("%/captured.png", snapshot_folder);
    stbi_write_png(temp_c_string(captured_path), width, height, 4, rgba.data, width * 4);
    image, info, decoded, decode_error := video_poster(options.path);
    if decoded {
      decoded_path := tprint("%/decoded.png", snapshot_folder);
      stbi_write_png(temp_c_string(decoded_path), image.width, image.height, 4, image.pixels, image.width * 4);
      print("  snapshot:   % and % (decoded at % s of %: %, % fps)\n", captured_path, decoded_path, FF(info.poster_time, 2), FF(info.duration, 2), info.codec, FF(info.frame_rate, 1));
      image_free(*image);
    } else {
      print("  snapshot:   %; the file could not be decoded back: %\n", captured_path, decode_error);
    }
  }
  exit(ifx stopped then cast(s32) 0 else 1);
}

#scope_file

EXCLUDED_BGR :: u8.[255, 0, 255];   // magenta
CONTROL_BGR  :: u8.[255, 255, 0];   // cyan

// Two squares on the region, topmost: one kept out of captures, one not.
exclusion_windows :: (region: RECT) -> excluded: HWND, control: HWND {
  width  := region.right - region.left;
  height := region.bottom - region.top;
  side := min(cast(s32) 160, width / 4, height / 2);
  y := region.top + (height - side) / 2;
  excluded := square_window("record-probe excluded", 0x00FF00FF, region.left + width / 4 - side / 2, y, side);
  control  := square_window("record-probe control",   0x00FFFF00, region.left + 3 * width / 4 - side / 2, y, side);
  WDA_EXCLUDEFROMCAPTURE :: 0x11;
  if !SetWindowDisplayAffinity(excluded, WDA_EXCLUDEFROMCAPTURE)  print("SetWindowDisplayAffinity failed; the test cannot say anything.\n");
  pump_messages();
  return excluded, control;
}

// A square of one colour, topmost, that never takes the focus.
square_window :: (class_name: string, colour: u32, x: s32, y: s32, side: s32) -> HWND {
  instance := GetModuleHandleW(null);
  wc: WNDCLASSEXW;
  wc.cbSize = size_of(WNDCLASSEXW);
  wc.lpfnWndProc = xx DefWindowProcW;
  wc.hInstance = instance;
  wc.hbrBackground = CreateSolidBrush(colour);
  wc.lpszClassName = utf8_to_wide(class_name,, temp);
  RegisterClassExW(*wc);
  return CreateWindowExW(WS_EX_TOPMOST | WS_EX_TOOLWINDOW | WS_EX_NOACTIVATE, utf8_to_wide(class_name,, temp), utf8_to_wide(class_name,, temp),
                         WS_POPUP | WS_VISIBLE, x, y, side, side, null, null, instance, null);
}

// A window's rectangle, in the region's pixels, a few pixels in from its edges.
window_rect_in :: (window: HWND, region: RECT) -> RECT {
  r: RECT;
  GetWindowRect(window, *r);
  return .{r.left - region.left + 4, r.top - region.top + 4, r.right - region.left - 4, r.bottom - region.top - 4};
}

colour_fraction :: (pixels: [] u8, width: s32, height: s32, r: RECT, bgr: [3] u8) -> float64 {
  total, same: s64;
  for y: max(r.top, 0)..min(r.bottom, height) - 1  for x: max(r.left, 0)..min(r.right, width) - 1 {
    p := pixels.data + (cast(s64) y * width + x) * 4;
    total += 1;
    if abs(cast(s32) p[0] - bgr[0]) <= 8 && abs(cast(s32) p[1] - bgr[1]) <= 8 && abs(cast(s32) p[2] - bgr[2]) <= 8  same += 1;
  }
  return ifx total then cast(float64) same / total else 0;
}

pump_messages :: () {
  msg: MSG;
  while PeekMessageW(*msg, null, 0, 0, PM_REMOVE) {
    TranslateMessage(*msg);
    DispatchMessageW(*msg);
  }
}

// A CUDA context on the first device, as onnxruntime makes for OCR and labels in mcap, and kept.
cuda_make_context :: () -> ok: bool, text: string {
  library := LoadLibraryW(utf8_to_wide("nvcuda.dll",, temp));
  if !library  return false, "nvcuda.dll is not installed";
  Cu_Init     :: #type (flags: u32) -> s32 #c_call;
  Cu_Device   :: #type (device: *s32, ordinal: s32) -> s32 #c_call;
  Cu_Context  :: #type (cuda_context: **void, flags: u32, device: s32) -> s32 #c_call;
  Cu_Name     :: #type (name: *u8, length: s32, device: s32) -> s32 #c_call;
  cu_init       := cast(Cu_Init)    GetProcAddress(library, "cuInit");
  cu_device_get := cast(Cu_Device)  GetProcAddress(library, "cuDeviceGet");
  cu_ctx_create := cast(Cu_Context) GetProcAddress(library, "cuCtxCreate_v2");
  cu_name       := cast(Cu_Name)    GetProcAddress(library, "cuDeviceGetName");
  if !cu_init || !cu_device_get || !cu_ctx_create || !cu_name  return false, "nvcuda.dll lacks the calls needed";
  result := cu_init(0);
  if result  return false, tprint("cuInit failed (%)", result);
  device: s32;
  result = cu_device_get(*device, 0);
  if result  return false, tprint("cuDeviceGet failed (%)", result);
  name: [256] u8;
  cu_name(name.data, name.count, device);
  cuda_context: *void;
  result = cu_ctx_create(*cuda_context, 0, device);
  if result  return false, tprint("cuCtxCreate failed (%)", result);
  return true, tprint("a context on %, kept for the recording", to_string(name.data));
}

bgra_to_rgba :: (pixels: [] u8) -> [] u8 {
  out := NewArray(pixels.count, u8,, temp);
  for i: 0..pixels.count / 4 - 1 {
    out[i * 4 + 0] = pixels[i * 4 + 2];
    out[i * 4 + 1] = pixels[i * 4 + 1];
    out[i * 4 + 2] = pixels[i * 4 + 0];
    out[i * 4 + 3] = 255;
  }
  return out;
}

file_size_of :: (path: string) -> s64 {
  _, size, ok := file_modtime_and_size(path);
  return ifx ok then size else 0;
}

percent :: (fraction: float64) -> string {
  return join(tprint("%", FF(fraction * 100, 1)), "%",, temp);
}

FF :: (value: float64, decimals: s64) -> FormatFloat #expand {
  return formatFloat(value, trailing_width = decimals, zero_removal = .NO);
}

SetProcessDpiAwarenessContext :: (value: HANDLE) -> BOOL #foreign user32_dll;
SetWindowDisplayAffinity      :: (window: HWND, affinity: u32) -> BOOL #foreign user32_dll;
user32_dll :: #library,system "user32";

#import "Basic";
#import "String";
#import "File";
#import "File_Utilities";
#import "Math";
#import "Windows";
#import "Windows_Utf8";
#import "Screen_Record";
#import "Image_Decode";
#import "stb_image_write";