// VRAM smoke harness (M5 doze VRAM diagnosis). Isolates the CAPTURE pipeline's
// GPU memory across the exact transitions the app drives — capture pause/resume
// (the doze + SteamVR-release mechanism) and full teardown/recreate — with NO
// headset and NO present loop. Logs the PROCESS-LOCAL VRAM via
// IDXGIAdapter3::QueryVideoMemoryInfo (CurrentUsage), the same number Task
// Manager's per-process "GPU memory" shows — so we can tell a real sauna leak
// from total-VRAM noise (e.g. SteamVR's own usage).
//
// Build: cmake --build build --config Release --target vram_smoke
// Run:   build\Release\vram_smoke.exe [cycles] [nsrc]
#include "capture/duplication_source.h"

#include <windows.h>

#include <d3d12.h>
#include <dxgi1_4.h>
#include <wrl/client.h>

#include <chrono>
#include <cstdio>
#include <cstdlib>
#include <memory>
#include <thread>
#include <vector>

using Microsoft::WRL::ComPtr;
using namespace sauna;

static ComPtr<IDXGIAdapter3> g_adapter;

static uint64_t vramMB() {
  DXGI_QUERY_VIDEO_MEMORY_INFO info{};
  if (g_adapter && SUCCEEDED(g_adapter->QueryVideoMemoryInfo(
                       0, DXGI_MEMORY_SEGMENT_GROUP_LOCAL, &info)))
    return info.CurrentUsage / (1024 * 1024);
  return 0;
}

static void logv(const char* tag) {
  printf("[VRAM] %-30s %5llu MB\n", tag, (unsigned long long)vramMB());
}

static void sleepMs(int ms) {
  std::this_thread::sleep_for(std::chrono::milliseconds(ms));
}

int main(int argc, char** argv) {
  setbuf(stdout, nullptr);
  const int cycles = argc > 1 ? atoi(argv[1]) : 12;
  int nsrc = argc > 2 ? atoi(argv[2]) : 0;  // 0 = all monitors

  ComPtr<IDXGIFactory4> factory;
  if (FAILED(CreateDXGIFactory1(IID_PPV_ARGS(&factory)))) return 1;
  ComPtr<IDXGIAdapter1> a1;
  if (FAILED(factory->EnumAdapters1(0, &a1))) return 1;
  a1.As(&g_adapter);
  ComPtr<ID3D12Device> dev;
  if (FAILED(D3D12CreateDevice(a1.Get(), D3D_FEATURE_LEVEL_12_0,
                               IID_PPV_ARGS(&dev)))) {
    fprintf(stderr, "D3D12CreateDevice failed\n");
    return 1;
  }
  logv("baseline (D3D12 device only)");

  int outputs = 0;
  ComPtr<IDXGIOutput> o;
  for (UINT i = 0; a1->EnumOutputs(i, &o) == S_OK; i++) outputs++;
  if (nsrc <= 0 || nsrc > outputs) nsrc = outputs;
  printf("[VRAM] %d output(s) on adapter 0, using %d capture source(s)\n",
         outputs, nsrc);

  auto build = [&](std::vector<std::unique_ptr<DuplicationSource>>& v) {
    v.clear();
    for (int i = 0; i < nsrc; i++)
      v.push_back(std::make_unique<DuplicationSource>(i));
    for (auto& s : v) s->start(dev.Get());
    sleepMs(900);  // let the ring allocate + a few frames publish
  };

  std::vector<std::unique_ptr<DuplicationSource>> srcs;
  build(srcs);
  logv("after start (ring allocated)");
  const uint64_t allocated = vramMB();

  // --- Test A: pause/resume cycles (the doze + SteamVR-release mechanism).
  // setPaused only dup.Reset()s; textures persist. Expect FLAT VRAM.
  printf("\n[VRAM] === Test A: %d x pause/resume (doze mechanism) ===\n",
         cycles);
  for (int c = 0; c < cycles; c++) {
    for (auto& s : srcs) s->setPaused(true);
    sleepMs(250);
    for (auto& s : srcs) s->setPaused(false);
    sleepMs(450);
    char t[64];
    snprintf(t, sizeof t, "cycle %2d resumed", c);
    logv(t);
  }
  const uint64_t afterPause = vramMB();

  // --- Test B: full teardown + recreate (destroy the D3D11 device too).
  // Tests whether tearing the source all the way down RETURNS VRAM, and
  // whether recreate LEAKS (the suspected source of the growth).
  printf("\n[VRAM] === Test B: 4 x full destroy/recreate ===\n");
  for (int c = 0; c < 4; c++) {
    srcs.clear();  // destroy sources (stop worker, release D3D11 device + ring)
    sleepMs(400);
    char t[64];
    snprintf(t, sizeof t, "destroy %d", c);
    logv(t);
    build(srcs);
    snprintf(t, sizeof t, "recreate %d", c);
    logv(t);
  }
  const uint64_t afterRecreate = vramMB();

  // --- Test C: stop()/resume() cycles (the M5 doze VRAM lever). stop() now
  // releases the D3D11 device + ring (frees VRAM) while keeping the object;
  // resume() rebuilds. Expect VRAM to DROP on stop and return on resume, flat.
  printf("\n[VRAM] === Test C: 4 x stop()/resume() (doze lever) ===\n");
  uint64_t stoppedLo = afterRecreate;
  for (int c = 0; c < 4; c++) {
    for (auto& s : srcs) s->stop();
    sleepMs(300);
    char t[64];
    snprintf(t, sizeof t, "stop %d", c);
    logv(t);
    stoppedLo = vramMB();
    for (auto& s : srcs) s->resume();
    sleepMs(900);
    snprintf(t, sizeof t, "resume %d", c);
    logv(t);
  }
  const uint64_t afterStopResume = vramMB();

  // --- Test D: resume() on a RUNNING source must be a safe no-op. Regression
  // guard for the wake-after-SteamVR crash: start() over a live worker thread
  // = std::terminate. Sources are running here (last Test C resume).
  printf("\n[VRAM] === Test D: resume() on running source (no crash) ===\n");
  for (auto& s : srcs) s->resume();  // running -> must not double-start
  sleepMs(300);
  logv("after resume-on-running");
  printf("[VRAM] Test D survived (no crash)\n");
  printf("[VRAM] stop() floor ~%lluMB (baseline was %lluMB) — drop %lldMB\n",
         (unsigned long long)stoppedLo, (unsigned long long)vramMB(),
         (long long)afterStopResume - (long long)stoppedLo);

  srcs.clear();
  sleepMs(400);
  logv("final (all destroyed)");

  printf(
      "\n[VRAM] SUMMARY: allocated=%lluMB  afterPause=%lluMB(%+lld)  "
      "afterRecreate=%lluMB(%+lld vs allocated)\n",
      (unsigned long long)allocated, (unsigned long long)afterPause,
      (long long)afterPause - (long long)allocated,
      (unsigned long long)afterRecreate,
      (long long)afterRecreate - (long long)allocated);
  printf(
      "[VRAM] VERDICT: pause/resume %s ; destroy/recreate %s\n",
      (afterPause > allocated + 30) ? "LEAKS" : "flat",
      (afterRecreate > allocated + 60) ? "LEAKS" : "flat");
  return 0;
}