// Screen_Record: a region of one screen recorded to H.264 in an MP4, for the mcap app's screen
// recording and tools/record_probe.jai (video-capture-plan.html). Windows only.
//
// Two threads. The capture thread takes the screen with DXGI Desktop Duplication (dxgi_source.jai),
// or with Windows.Graphics.Capture (wgc_source.jai) where duplication cannot run, each on a D3D11
// device of its own; copies the region on the GPU; reads it back through a ring of staging textures,
// so Map never waits on the GPU (capture_gpu.jai); and, for duplication, draws the cursor in
// (cursor.jai). The encoder thread hands the frames to FFmpeg (encoder.jai). Between them is a pool of
// RECORD_BUFFERS frames; when the encoder falls RECORD_QUEUE frames behind, the oldest waiting is
// dropped and counted.
//
// A frame is taken at most fps times a second, and only when the screen or the cursor changed, with
// one a second while nothing does, so players can seek. Each carries the time it was taken.
//
//   recorder: Recorder;
//   ok, error, fall_back := record_start(*recorder, .{path = "out.mp4", screen = 0});
//   ...
//   ok, error = record_stop(*recorder);
//
// With capture .AUTO (the default), Desktop Duplication is tried first and Windows.Graphics.Capture
// takes over when it cannot run: on a laptop with two GPUs, on a rotated screen, or with too many
// duplications of the screen running. fall_back is only returned for .DXGI, which does not fall back.

RECORD_MIN_WIDTH  :: 240;   // a smaller region is refused
RECORD_MIN_HEIGHT :: 180;
RECORD_QUEUE      :: 8;     // frames waiting for the encoder at most
RECORD_BUFFERS    :: RECORD_QUEUE + 2;   // and the one being filled, and the one being encoded

Record_Screen :: struct {
  adapter_index, output_index: u32;
  adapter: string;    // "NVIDIA GeForce RTX 5080 Laptop GPU"
  name: string;       // "\\.\DISPLAY1"
  desktop: RECT;      // on the virtual desktop
  rotation: DXGI_MODE_ROTATION;
  monitor: HMONITOR;
}

Record_Capture :: enum u8 {
  AUTO;   // Desktop Duplication, else Windows.Graphics.Capture
  DXGI;   // Desktop Duplication only
  WGC;    // Windows.Graphics.Capture only
}

// For record-probe: paths a machine may not take by itself.
Record_Testing :: enum_flags u8 {
  DXGI_UNSUPPORTED;   // Desktop Duplication fails as on a laptop with two GPUs (DXGI_ERROR_UNSUPPORTED)
  WGC_POLL;           // Windows.Graphics.Capture polled, as when it will not take the FrameArrived handler
  KEEP_BORDER;        // Windows.Graphics.Capture's yellow border left on
}

Record_Options :: struct {
  path: string;                // the .mp4 written
  screen: s32 = -1;            // into record_list_screens(); -1: the one holding region
  region: RECT;                // on the virtual desktop; empty: the whole screen
  fps: s32 = 30;
  cursor := true;
  encoder: string;             // "" tries RECORD_ENCODERS in turn
  capture: Record_Capture;
  legacy_duplicate: bool;      // the plain DuplicateOutput, not DuplicateOutput1: for record-probe
  testing: Record_Testing;
}

Record_Stats :: struct {
  source: string;              // "DXGI DuplicateOutput1 on <adapter>", "Windows.Graphics.Capture on <adapter>, border off"
  fell_back_from: string;      // why Desktop Duplication could not run, when it fell back
  encoder: string;
  width, height: s32;
  captured: s64;               // new images of the screen
  pointer_moves: s64;          // cursor changes with no new image (Desktop Duplication)
  emitted: s64;                // frames sent to be encoded
  repeats: s64;                // of them, the once-a-second repeat of an unchanged screen
  written: s64;                // encoded
  dropped: s64;                // the encoder was too far behind
  lost: s64;                   // times the duplication was lost and made again
  copy_ticks, readback_ticks, encode_ticks: s64;   // totals, in QueryPerformanceCounter ticks
  bytes: s64;                  // encoded
  seconds: float64;            // from start to stop
}

Recorder :: struct {
  options: Record_Options;
  screen: Record_Screen;
  local: RECT;                 // the region, in the screen's pixels
  dxgi: Dxgi_Source;           // the capture thread's: one of these is open
  wgc: Wgc_Source;
  using_wgc: bool;
  encoder: Encoder;
  width, height: s32;
  frequency: s64;              // QueryPerformanceCounter's
  start_ticks: s64;

  capture_thread, encoder_thread: Thread;
  capture_started, encoder_started: bool;
  running: bool;
  stopping: bool;              // set by record_stop; the capture thread ends, then the encoder
  capture_ended: bool;

  opened: Semaphore;           // the capture thread has opened its source, or failed to
  open_ok, open_fall_back: bool;
  open_error: string;

  mutex: Mutex;
  ready: Semaphore;            // signalled for each frame queued, and once when capture ends
  frames: [RECORD_BUFFERS] Record_Frame;
  free: [RECORD_BUFFERS] s32;
  free_count: s32;
  queue: [RECORD_BUFFERS] s32; // oldest first
  queue_count: s32;
  last_frame: [] u8;           // the last frame queued, BGRA, for record_last_frame
  stats: Record_Stats;
  error: string;               // why it stopped by itself; heap
}

Record_Frame :: struct {
  pixels: [] u8;               // BGRA, rows packed
  ticks: s64;
}

// Opens the screen, the encoder and the file, and starts recording.
record_start :: (r: *Recorder, options: Record_Options) -> ok: bool, error: string, fall_back: bool {
  r.* = .{};
  r.options = options;
  r.options.path = copy_string(options.path);
  r.options.encoder = copy_string(options.encoder);
  if r.options.fps <= 0  r.options.fps = 30;

  screens := record_list_screens();
  if !screens.count  return false, "no screen to record", false;
  screen_index := options.screen;
  region := options.region;
  has_region := region.right > region.left && region.bottom > region.top;
  if screen_index < 0 {
    if !has_region  screen_index = 0;
    else for screens  if rect_inside(region, it.desktop) { screen_index = xx it_index; break; }
    if screen_index < 0  return false, "the region is not on one screen", false;
  }
  if screen_index >= screens.count  return false, tprint("there is no screen %", screen_index), false;
  r.screen = screens[screen_index];
  r.screen.adapter = copy_string(r.screen.adapter);
  r.screen.name = copy_string(r.screen.name);
  if !has_region  region = r.screen.desktop;
  if !rect_inside(region, r.screen.desktop)  return false, "the region is not on that screen", false;

  // The minimum first, then even both ways, in the screen's pixels: H.264's 4:2:0 colour needs it.
  // RECORD_MIN_WIDTH and _HEIGHT are even, so what passes stays at least that size.
  #assert RECORD_MIN_WIDTH % 2 == 0 && RECORD_MIN_HEIGHT % 2 == 0;
  if region.right - region.left < RECORD_MIN_WIDTH || region.bottom - region.top < RECORD_MIN_HEIGHT
    return false, tprint("% x % is smaller than % x %", region.right - region.left, region.bottom - region.top, RECORD_MIN_WIDTH, RECORD_MIN_HEIGHT), false;
  d := r.screen.desktop;
  r.local = .{region.left - d.left, region.top - d.top, region.right - d.left, region.bottom - d.top};
  r.local.right  = r.local.left + ((r.local.right - r.local.left) & cast(s32) ~1);
  r.local.bottom = r.local.top  + ((r.local.bottom - r.local.top) & cast(s32) ~1);
  r.width  = r.local.right - r.local.left;
  r.height = r.local.bottom - r.local.top;
  r.stats.width, r.stats.height = r.width, r.height;

  init(*r.mutex, "record");
  init(*r.ready);
  init(*r.opened);
  for * r.frames  it.pixels = NewArray(cast(s64) r.width * r.height * 4, u8, initialized = false);
  for 0..RECORD_BUFFERS - 1  r.free[it] = cast(s32) it;
  r.free_count = RECORD_BUFFERS;
  r.last_frame = NewArray(cast(s64) r.width * r.height * 4, u8);
  frequency: LARGE_INTEGER;
  QueryPerformanceFrequency(*frequency);
  r.frequency = frequency.QuadPart;
  r.start_ticks = ticks_now();

  // The source is opened on the capture thread, which uses it: WinRT's objects belong to the
  // apartment they are made in.
  thread_init(*r.capture_thread, record_capture_proc, name = "record capture");
  r.capture_started = true;
  thread_start(*r.capture_thread, r);
  wait_for(*r.opened);
  if !r.open_ok {
    error, fall_back := copy_string(r.open_error,, temp), r.open_fall_back;
    record_free(r);
    return false, error, fall_back;
  }

  ok, error, tried := encoder_open(*r.encoder, r.options.path, r.width, r.height, r.options.fps, r.options.encoder);
  if tried  log("[record] Encoders that did not open: %", tried);
  if !ok {
    error = copy_string(error,, temp);
    record_free(r);
    return false, error, false;
  }
  r.stats.encoder = copy_string(r.encoder.name);

  thread_init(*r.encoder_thread, record_encoder_proc, name = "record encoder");
  r.encoder_started = true;
  r.running = true;
  thread_start(*r.encoder_thread, r);
  return true, "", false;
}

// Stops, finishes the file and frees everything. ok is false when recording stopped by itself on an
// error, or the file could not be finished; what was recorded up to then is in the file all the same.
record_stop :: (r: *Recorder) -> ok: bool, error: string {
  if !r.running  return false, "not recording";
  atomic_swap(*r.stopping, true);
  thread_is_done(*r.capture_thread, -1);
  thread_is_done(*r.encoder_thread, -1);
  r.running = false;

  closed, close_error := encoder_close(*r.encoder);
  r.stats.seconds = cast(float64) (ticks_now() - r.start_ticks) / r.frequency;
  r.stats.bytes = r.encoder.bytes;
  error := ifx r.error then copy_string(r.error,, temp) else "";
  record_free(r);

  if error   return false, error;
  if !closed return false, close_error;
  return true, "";
}

// As they are now; strings are the recorder's.
record_stats :: (r: *Recorder) -> Record_Stats {
  if !r.running  return r.stats;
  lock(*r.mutex);
  defer unlock(*r.mutex);
  return r.stats;
}

// The last frame queued, BGRA (temp), cursor and all.
record_last_frame :: (r: *Recorder) -> pixels: [] u8, width: s32, height: s32 {
  if r.running  lock(*r.mutex);
  pixels := array_copy(r.last_frame,, temp);
  if r.running  unlock(*r.mutex);
  return pixels, r.width, r.height;
}

// Whether recording stopped by itself (the screen went away, the encoder failed); why, then.
record_failed :: (r: *Recorder) -> bool, string {
  lock(*r.mutex);
  defer unlock(*r.mutex);
  return r.error.count > 0, r.error;
}

#scope_file

Staged :: struct { pending: bool; ticks: s64; pointer: POINT; pointer_visible: bool; }

// Stops the capture thread if it runs, and frees what record_start made. The stats' strings are kept.
record_free :: (r: *Recorder) {
  atomic_swap(*r.stopping, true);
  if r.capture_started { thread_is_done(*r.capture_thread, -1); thread_deinit(*r.capture_thread); r.capture_started = false; }
  if r.encoder_started { thread_is_done(*r.encoder_thread, -1); thread_deinit(*r.encoder_thread); r.encoder_started = false; }
  if r.encoder.format  encoder_close(*r.encoder);
  for * r.frames  array_free(it.pixels);
  array_free(r.last_frame);
  r.last_frame = .[];
  destroy(*r.opened);
  destroy(*r.ready);
  destroy(*r.mutex);
  free(r.open_error);
  free(r.error);
  r.open_error, r.error = "", "";
}

// Desktop Duplication, unless asked for Windows.Graphics.Capture, which takes over when duplication
// cannot run here and the options allow.
record_open_source :: (r: *Recorder) -> bool {
  if r.options.capture != .WGC {
    ok, error, fall_back := dxgi_open(*r.dxgi, r.screen, r.local, r.options.legacy_duplicate);
    if ok && r.options.testing & .DXGI_UNSUPPORTED {
      ok, error, fall_back = false, "Desktop Duplication (testing): 0x887A0004 (DXGI_ERROR_UNSUPPORTED)", true;
    }
    if ok {
      r.stats.source = sprint("DXGI % on %", r.dxgi.method, r.screen.adapter);
      return true;
    }
    dxgi_close(*r.dxgi);
    if r.options.capture == .DXGI || !fall_back {
      r.open_error, r.open_fall_back = copy_string(error), fall_back;
      return false;
    }
    log("[record] Desktop Duplication cannot run here (%); Windows.Graphics.Capture instead.", error);
    r.stats.fell_back_from = copy_string(error);
  }
  ok, error := wgc_open(*r.wgc, r.screen, r.local, r.options.cursor, r.options.testing);
  if !ok {
    wgc_close(*r.wgc);
    r.open_error = copy_string(error);
    return false;
  }
  r.using_wgc = true;
  r.stats.source = sprint("Windows.Graphics.Capture on %, %%", r.screen.adapter, r.wgc.border, ifx r.wgc.polling then ", polled" else "");
  return true;
}

record_capture_proc :: (thread: *Thread) -> s64 {
  r := cast(*Recorder) thread.data;
  wgc_thread_start();
  defer wgc_thread_end();
  r.open_ok = record_open_source(r);
  signal(*r.opened);
  if !r.open_ok  return 0;
  gpu: *Capture_Gpu = ifx r.using_wgc then *r.wgc.gpu else *r.dxgi.gpu;
  defer if r.using_wgc  wgc_close(*r.wgc); else dxgi_close(*r.dxgi);

  // A change is sent as soon as it comes once its turn is due: turns are a frame's time apart on a
  // schedule that moves on a frame at a time, so the rate never passes fps, and a tenth of a frame
  // early counts as on time, so a 60 Hz screen's frames, which jitter, all make a 60 fps cap.
  period := r.frequency / r.options.fps;
  early := period / 10;
  turn := ticks_now();
  last_emit: s64 = -r.frequency;
  have_image, changed := false, false;
  staged: [RECORD_STAGING] Staged;
  emits: s64;

  while !atomic_read(*r.stopping) {
    now := ticks_now();
    due := ifx changed then turn - early else last_emit + r.frequency;   // the change, or the repeat
    wait := cast(u32) clamp((due - now) * 1000 / r.frequency, 0, 100);
    next: Capture_Next;
    error: string;
    if r.using_wgc  next, error = wgc_next(*r.wgc, wait);
    else            next, error = dxgi_next(*r.dxgi, wait);
    if next == {
      case .FAILED;
        record_fail(r, error);
        break;
      case .LOST;
        // A mode change, the secure desktop, a game going fullscreen: duplicated again, and on.
        lock(*r.mutex); r.stats.lost += 1; unlock(*r.mutex);
        while !atomic_read(*r.stopping) {
          hr := dxgi_duplicate(*r.dxgi);
          if SUCCEEDED(hr)  break;
          if !hresult_is(hr, DXGI_ERROR_ACCESS_LOST) && !hresult_is(hr, 0x80070005) && !hresult_is(hr, DXGI_ERROR_NOT_CURRENTLY_AVAILABLE) {
            record_fail(r, hresult_text("the screen could not be duplicated again", hr));
            break;
          }
          sleep_milliseconds(250);
        }
        if atomic_read(*r.stopping) || r.error  break;
        continue;
      case .IMAGE; #through;
      case .IMAGE_AND_POINTER;
        have_image, changed = true, true;
        lock(*r.mutex);
        r.stats.captured += 1;
        r.stats.copy_ticks = gpu.copy_ticks;
        unlock(*r.mutex);
      case .POINTER;
        if r.options.cursor && have_image {
          changed = true;
          lock(*r.mutex); r.stats.pointer_moves += 1; unlock(*r.mutex);
        }
    }

    now = ticks_now();
    if !have_image  continue;
    repeat := !changed && now - last_emit >= r.frequency;
    if !(changed && now >= turn - early) && !repeat  continue;
    turn = max(turn + period, now - period);   // after a quiet spell, no burst to catch up

    // Copied to staging now; read back two frames on, when the GPU is done with it.
    slot := emits % RECORD_STAGING;
    if staged[slot].pending  record_read_back(r, gpu, slot, staged[slot]);
    ID3D11DeviceContext_CopyResource(gpu.device_context, gpu.staging[slot], gpu.latest);
    staged[slot] = .{pending = true, ticks = now};
    if !r.using_wgc  staged[slot].pointer, staged[slot].pointer_visible = r.dxgi.pointer_position, r.dxgi.pointer_visible;
    emits += 1;
    last_emit = now;
    changed = false;
    lock(*r.mutex);
    r.stats.emitted += 1;
    if repeat  r.stats.repeats += 1;
    unlock(*r.mutex);
  }

  // What is still staged, oldest first.
  for 0..RECORD_STAGING - 1 {
    slot := (emits + it) % RECORD_STAGING;
    if staged[slot].pending  record_read_back(r, gpu, slot, staged[slot]);
  }
  atomic_swap(*r.capture_ended, true);
  signal(*r.ready);
  return 0;
}

// Copies a staging texture into a free frame, draws the cursor (Desktop Duplication leaves it out)
// and queues it.
record_read_back :: (r: *Recorder, gpu: *Capture_Gpu, slot: s64, staged: Staged) {
  t := ticks_now();
  lock(*r.mutex);
  if !r.free_count {   // every frame is waiting: drop the oldest
    r.free[r.free_count] = r.queue[0];
    r.free_count += 1;
    for 1..r.queue_count - 1  r.queue[it - 1] = r.queue[it];
    r.queue_count -= 1;
    r.stats.dropped += 1;
  }
  r.free_count -= 1;
  index := r.free[r.free_count];
  unlock(*r.mutex);
  frame := *r.frames[index];

  mapped: D3D11_MAPPED_SUBRESOURCE;
  hr := ID3D11DeviceContext_Map(gpu.device_context, gpu.staging[slot], 0, .READ, 0, *mapped);
  if FAILED(hr) {
    lock(*r.mutex); r.free[r.free_count] = index; r.free_count += 1; unlock(*r.mutex);
    record_fail(r, hresult_text("reading the frame back", hr));
    return;
  }
  row := cast(s64) r.width * 4;
  for y: 0..r.height - 1  memcpy(frame.pixels.data + y * row, cast(*u8) mapped.pData + cast(s64) y * mapped.RowPitch, row);
  ID3D11DeviceContext_Unmap(gpu.device_context, gpu.staging[slot], 0);
  if r.options.cursor && staged.pointer_visible {
    cursor_draw(frame.pixels.data, r.width, r.height, *r.dxgi.shape,
                staged.pointer.x - cast(s32) gpu.box.left, staged.pointer.y - cast(s32) gpu.box.top);
  }
  frame.ticks = staged.ticks;

  lock(*r.mutex);
  r.queue[r.queue_count] = index;
  r.queue_count += 1;
  memcpy(r.last_frame.data, frame.pixels.data, frame.pixels.count);
  r.stats.readback_ticks += ticks_now() - t;
  unlock(*r.mutex);
  signal(*r.ready);
}

record_encoder_proc :: (thread: *Thread) -> s64 {
  r := cast(*Recorder) thread.data;
  first: s64 = -1;
  while true {
    wait_for(*r.ready, 100);
    lock(*r.mutex);
    index: s32 = -1;
    if r.queue_count {
      index = r.queue[0];
      for 1..r.queue_count - 1  r.queue[it - 1] = r.queue[it];
      r.queue_count -= 1;
    }
    unlock(*r.mutex);
    if index < 0 {
      if atomic_read(*r.capture_ended)  break;
      continue;
    }

    frame := *r.frames[index];
    if first < 0  first = frame.ticks;
    pts := (frame.ticks - first) * 90000 / r.frequency;
    t := ticks_now();
    ok, error := encoder_write(*r.encoder, frame.pixels.data, pts);

    lock(*r.mutex);
    r.free[r.free_count] = index;
    r.free_count += 1;
    r.stats.encode_ticks += ticks_now() - t;
    if ok  r.stats.written += 1;
    r.stats.bytes = r.encoder.bytes;
    unlock(*r.mutex);
    if !ok {
      record_fail(r, error);
      // Frames still come until the capture thread sees it; they are let go.
      while !atomic_read(*r.capture_ended) {
        wait_for(*r.ready, 100);
        lock(*r.mutex);
        for 0..r.queue_count - 1 { r.free[r.free_count] = r.queue[it]; r.free_count += 1; }
        r.queue_count = 0;
        unlock(*r.mutex);
      }
      break;
    }
  }
  return 0;
}

// Ends the recording from one of its threads; the first reason is kept.
record_fail :: (r: *Recorder, error: string) {
  log_error("[record] %", error);
  lock(*r.mutex);
  if !r.error  r.error = copy_string(error);
  unlock(*r.mutex);
  atomic_swap(*r.stopping, true);
}

rect_inside :: (inner: RECT, outer: RECT) -> bool {
  return inner.left >= outer.left && inner.top >= outer.top && inner.right <= outer.right && inner.bottom <= outer.bottom;
}

#scope_export

#load "capture_gpu.jai";
#load "dxgi_source.jai";
#load "wgc_source.jai";
#load "encoder.jai";
#load "cursor.jai";

#import "Basic";
#import "String";
#import "Atomics";
#import "Thread";
#import "Windows";
#import "Windows_Utf8";
#import "dxgi";
#import "d3d11";
#import "ffmpeg";