// Spike S2, NVAPI rung: own-process DirectMode presenter for the Beyond.
// Ported from the first-party mmi_test directmode sample (BigscreenVR/mmi_test).
// Requires the NDA HMD NVAPI headers/lib from a local mmi_test checkout —
// referenced in place, NOT vendored into this repo (licensing).
//
//   s2_nvapi present <seconds> [--force-dsc] [--image <file.raw>]
//
// --image: raw B8G8R8X8 (byte order B,G,R,X) frame, W*H*4 bytes matching the
// selected mode, presented statically instead of the animated pattern. Used
// by spike S3 for the eyes-in pre-distorted grid check (R8G8B8A8 surfaces get
// a CPU swizzle; 10/16-bit surface flavors are rejected).
//
// Exit codes: 0 ok, 2 no DM display, 3 acquire/modeset failed, 10 other.

#include <windows.h>

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

#include <nvapi.h>

#include <chrono>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <vector>

using Microsoft::WRL::ComPtr;

static const NvU32 VENDOR_ID = 0x2709;  // EDID "BIG"

#define NVCHECK(call)                                                       \
  do {                                                                      \
    NvAPI_Status s_ = (call);                                               \
    if (s_ != NVAPI_OK) {                                                   \
      fprintf(stderr, "FATAL %s = 0x%x (line %d)\n", #call, s_, __LINE__);  \
      return 10;                                                            \
    }                                                                       \
  } while (0)

#define HRCHECK(call)                                                       \
  do {                                                                      \
    HRESULT h_ = (call);                                                    \
    if (FAILED(h_)) {                                                       \
      fprintf(stderr, "FATAL %s = 0x%lx (line %d)\n", #call, h_, __LINE__); \
      return 10;                                                            \
    }                                                                       \
  } while (0)

int main(int argc, char** argv) {
  setbuf(stdout, nullptr);
  setbuf(stderr, nullptr);
  double seconds = argc > 2 ? atof(argv[2]) : 15.0;
  // Defaults = the visually-validated config (2026-06-10, both rigs):
  // X8R8G8B8 mode + forced DSC 1.1/4-slice/8bpp per the device JSON.
  bool forceDsc = true, powerOff = false;
  int wantMode = -1;
  const char* imagePath = nullptr;
  for (int i = 1; i < argc; i++) {
    if (strcmp(argv[i], "--force-dsc") == 0) forceDsc = true;
    if (strcmp(argv[i], "--no-force-dsc") == 0) forceDsc = false;
    if (strcmp(argv[i], "--power-off") == 0) powerOff = true;
    if (strcmp(argv[i], "--mode") == 0 && i + 1 < argc) wantMode = atoi(argv[i + 1]);
    if (strcmp(argv[i], "--image") == 0 && i + 1 < argc) imagePath = argv[++i];
  }

  NVCHECK(NvAPI_Initialize());
  NvAPI_Status st = NvAPI_DISP_EnableDirectMode(VENDOR_ID, 0);
  printf("EnableDirectMode(0x%x) = 0x%x\n", VENDOR_ID, st);

  // --- enumerate DM displays ---
  NvU32 numDisplays = 0;
  st = NvAPI_DISP_EnumerateDirectModeDisplays(VENDOR_ID, &numDisplays, nullptr,
                                              NV_ENUM_DIRECTMODE_DISPLAY_ENABLED);
  printf("EnumerateDirectModeDisplays count = %u (status 0x%x)\n", numDisplays, st);
  if (numDisplays == 0) { printf("RESULT: no DM display\n"); return 2; }
  std::vector<NV_DIRECT_MODE_DISPLAY_HANDLE> displays(numDisplays);
  NVCHECK(NvAPI_DISP_EnumerateDirectModeDisplays(VENDOR_ID, &numDisplays,
                                                 displays.data(),
                                                 NV_ENUM_DIRECTMODE_DISPLAY_ENABLED));
  NV_DIRECT_MODE_DISPLAY_HANDLE display = displays[0];
  printf("using displayId 0x%x\n", display.displayId);

  // --- list modes ---
  NvU32 modeCount = 0;
  NVCHECK(NvAPI_D3D_DirectModeGetDisplayModes(&display, &modeCount, nullptr,
                                              NV_DIRECTMODE_GETMODES_FLAG_SUPPORTED));
  std::vector<NV_DIRECT_MODE_INFO> modes(modeCount);
  for (auto& m : modes) { memset(&m, 0, sizeof(m)); m.version = NV_DIRECT_MODE_INFO_VER; }
  NVCHECK(NvAPI_D3D_DirectModeGetDisplayModes(&display, &modeCount, modes.data(),
                                              NV_DIRECTMODE_GETMODES_FLAG_SUPPORTED));
  int best = -1;
  double bestScore = -1;
  for (NvU32 i = 0; i < modeCount; i++) {
    auto& m = modes[i];
    double hz = m.refresh.denominator
                    ? (double)m.refresh.numerator / m.refresh.denominator
                    : 0;
    printf("mode[%u]: %ux%u @ %.3f Hz format=%d dscMode=%d dscVer=%d slices=%d bppx16=%u\n",
           i, m.width, m.height, hz, (int)m.format, (int)m.dscParams.dscMode,
           (int)m.dscParams.dscVersion, (int)m.dscParams.sliceCount,
           m.dscParams.outputBPPx16);
    double score = (double)m.width * m.height - fabs(hz - 75.0) * 1e4;
    if (m.format == 22) score += 1;  // prefer X8R8G8B8 — the validated format
    if (score > bestScore) { bestScore = score; best = (int)i; }
  }
  if (wantMode >= 0 && wantMode < (int)modeCount) best = wantMode;
  NV_DIRECT_MODE_INFO mode = modes[best];
  double modeHz = mode.refresh.denominator
                      ? (double)mode.refresh.numerator / mode.refresh.denominator
                      : 0;
  printf("using mode[%d]: %ux%u @ %.3f Hz\n", best, mode.width, mode.height, modeHz);
  if (forceDsc) {
    mode.dscParams.dscMode = NV_DIRECT_MODE_DSC_MODE_FORCE_ENABLED;
    mode.dscParams.dscVersion = NV_DIRECT_MODE_DSC_VERSION_V11;
    mode.dscParams.sliceCount = NV_DIRECT_MODE_DSC_SLICE_COUNT_4;
    mode.dscParams.outputBPPx16 = 128;
    printf("forcing DSC 1.1, 4 slices, 8 bpp\n");
  }

  // --- D3D12 device on the display's adapter ---
  NvPhysicalGpuHandle gpu;
  NVCHECK(NvAPI_SYS_GetPhysicalGpuFromDisplayId(display.displayId, &gpu));
  ComPtr<IDXGIAdapter> adapter;
  NVCHECK(NvAPI_D3D_GetIDXGIAdapter(gpu, adapter.GetAddressOf()));
  ComPtr<ID3D12Device> dev;
  HRCHECK(D3D12CreateDevice(adapter.Get(), D3D_FEATURE_LEVEL_11_0,
                            IID_PPV_ARGS(&dev)));
  D3D12_COMMAND_QUEUE_DESC qd{};
  qd.Type = D3D12_COMMAND_LIST_TYPE_DIRECT;
  ComPtr<ID3D12CommandQueue> queue;
  HRCHECK(dev->CreateCommandQueue(&qd, IID_PPV_ARGS(&queue)));
  ComPtr<ID3D12CommandAllocator> alloc;
  HRCHECK(dev->CreateCommandAllocator(D3D12_COMMAND_LIST_TYPE_DIRECT,
                                      IID_PPV_ARGS(&alloc)));
  ComPtr<ID3D12GraphicsCommandList> list;
  HRCHECK(dev->CreateCommandList(0, D3D12_COMMAND_LIST_TYPE_DIRECT, alloc.Get(),
                                 nullptr, IID_PPV_ARGS(&list)));
  HRCHECK(list->Close());
  ComPtr<ID3D12Fence> fence;
  HRCHECK(dev->CreateFence(0, D3D12_FENCE_FLAG_NONE, IID_PPV_ARGS(&fence)));
  HANDLE fenceEvt = CreateEventA(nullptr, FALSE, FALSE, nullptr);
  UINT64 fenceVal = 0;

  // --- acquire ---
  st = NvAPI_D3D_AcquireDirectModeDisplay(VENDOR_ID, dev.Get(), &display);
  printf("AcquireDirectModeDisplay = 0x%x%s\n", st,
         st == NVAPI_HDCP_DISABLED ? " (HDCP disabled, continuing)" : "");
  if (st != NVAPI_OK && st != NVAPI_HDCP_DISABLED) return 3;

  // --- surfaces ---
  const int NBUF = 2;
  NV_DIRECT_MODE_SURFACE_HANDLE surf[NBUF];
  HANDLE shared[NBUF];
  ComPtr<ID3D12Resource> res[NBUF];
  for (int i = 0; i < NBUF; i++) {
    st = NvAPI_D3D_DirectModeCreateSurface(&display, &mode, &surf[i], &shared[i]);
    if (st != NVAPI_OK) { fprintf(stderr, "CreateSurface[%d] = 0x%x\n", i, st); return 3; }
    HRCHECK(dev->OpenSharedHandle(shared[i], IID_PPV_ARGS(&res[i])));
  }
  auto rdesc = res[0]->GetDesc();
  printf("surface: %llux%u format=%d\n", (unsigned long long)rdesc.Width,
         rdesc.Height, (int)rdesc.Format);
  // DM surfaces come back TYPELESS — an RTV needs an explicit typed format
  // (nullptr desc on a typeless resource silently fails in release builds:
  // clears never land and the panels scan out uninitialized VRAM).
  DXGI_FORMAT rtvFmt;
  switch (rdesc.Format) {
    case DXGI_FORMAT_B8G8R8A8_TYPELESS: rtvFmt = DXGI_FORMAT_B8G8R8A8_UNORM; break;
    case DXGI_FORMAT_B8G8R8X8_TYPELESS: rtvFmt = DXGI_FORMAT_B8G8R8X8_UNORM; break;
    case DXGI_FORMAT_R8G8B8A8_TYPELESS: rtvFmt = DXGI_FORMAT_R8G8B8A8_UNORM; break;
    case DXGI_FORMAT_R10G10B10A2_TYPELESS: rtvFmt = DXGI_FORMAT_R10G10B10A2_UNORM; break;
    case DXGI_FORMAT_R16G16B16A16_TYPELESS: rtvFmt = DXGI_FORMAT_R16G16B16A16_FLOAT; break;
    default:
      rtvFmt = rdesc.Format;
      fprintf(stderr, "WARNING: surface format %d has no typed mapping — "
                      "RTV will be invalid if it is typeless\n", (int)rdesc.Format);
      break;
  }
  printf("rtv format: %d\n", (int)rtvFmt);

  ComPtr<ID3D12DescriptorHeap> rtvHeap;
  D3D12_DESCRIPTOR_HEAP_DESC hd{};
  hd.Type = D3D12_DESCRIPTOR_HEAP_TYPE_RTV;
  hd.NumDescriptors = NBUF;
  HRCHECK(dev->CreateDescriptorHeap(&hd, IID_PPV_ARGS(&rtvHeap)));
  UINT rtvStep = dev->GetDescriptorHandleIncrementSize(D3D12_DESCRIPTOR_HEAP_TYPE_RTV);
  D3D12_CPU_DESCRIPTOR_HANDLE rtv[NBUF];
  for (int i = 0; i < NBUF; i++) {
    rtv[i] = rtvHeap->GetCPUDescriptorHandleForHeapStart();
    rtv[i].ptr += (SIZE_T)i * rtvStep;
    D3D12_RENDER_TARGET_VIEW_DESC rd{};
    rd.Format = rtvFmt;
    rd.ViewDimension = D3D12_RTV_DIMENSION_TEXTURE2D;
    rd.Texture2D.MipSlice = 0;
    rd.Texture2D.PlaneSlice = 0;
    dev->CreateRenderTargetView(res[i].Get(), &rd, rtv[i]);
  }

  // --- optional static image: load raw frame into every surface once ---
  bool haveImage = false;
  ComPtr<ID3D12Resource> upload;
  if (imagePath) {
    if (rtvFmt != DXGI_FORMAT_B8G8R8X8_UNORM && rtvFmt != DXGI_FORMAT_B8G8R8A8_UNORM &&
        rtvFmt != DXGI_FORMAT_R8G8B8A8_UNORM) {
      fprintf(stderr, "FATAL: --image only supports 8-bit BGRX/BGRA/RGBA surfaces, got rtv format %d\n",
              (int)rtvFmt);
      return 10;
    }
    UINT W = (UINT)rdesc.Width, H = rdesc.Height;
    size_t rowBytes = (size_t)W * 4;
    UINT64 pitch = (rowBytes + 255) & ~(size_t)255;
    std::vector<uint8_t> data(rowBytes * H);
    FILE* f = fopen(imagePath, "rb");
    if (!f) { fprintf(stderr, "FATAL: cannot open %s\n", imagePath); return 10; }
    size_t got = fread(data.data(), 1, data.size(), f);
    fclose(f);
    if (got != data.size()) {
      fprintf(stderr, "FATAL: %s is %zu bytes, mode wants %zu (%ux%u*4)\n",
              imagePath, got, data.size(), W, H);
      return 10;
    }
    if (rtvFmt == DXGI_FORMAT_R8G8B8A8_UNORM) {  // file is BGRX -> swizzle
      for (size_t p = 0; p < data.size(); p += 4) {
        uint8_t t = data[p]; data[p] = data[p + 2]; data[p + 2] = t;
      }
      printf("image swizzled BGRX -> RGBA for surface format\n");
    }

    D3D12_HEAP_PROPERTIES hp{};
    hp.Type = D3D12_HEAP_TYPE_UPLOAD;
    D3D12_RESOURCE_DESC bd{};
    bd.Dimension = D3D12_RESOURCE_DIMENSION_BUFFER;
    bd.Width = pitch * H;
    bd.Height = 1;
    bd.DepthOrArraySize = 1;
    bd.MipLevels = 1;
    bd.Format = DXGI_FORMAT_UNKNOWN;
    bd.SampleDesc.Count = 1;
    bd.Layout = D3D12_TEXTURE_LAYOUT_ROW_MAJOR;
    HRCHECK(dev->CreateCommittedResource(&hp, D3D12_HEAP_FLAG_NONE, &bd,
                                         D3D12_RESOURCE_STATE_GENERIC_READ,
                                         nullptr, IID_PPV_ARGS(&upload)));
    uint8_t* dst = nullptr;
    D3D12_RANGE noRead{0, 0};
    HRCHECK(upload->Map(0, &noRead, (void**)&dst));
    for (UINT y = 0; y < H; y++)
      memcpy(dst + y * pitch, data.data() + y * rowBytes, rowBytes);
    upload->Unmap(0, nullptr);

    for (int i = 0; i < NBUF; i++) {
      HRCHECK(alloc->Reset());
      HRCHECK(list->Reset(alloc.Get(), nullptr));
      D3D12_RESOURCE_BARRIER cb{};
      cb.Type = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION;
      cb.Transition.pResource = res[i].Get();
      cb.Transition.Subresource = D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES;
      cb.Transition.StateBefore = D3D12_RESOURCE_STATE_COMMON;
      cb.Transition.StateAfter = D3D12_RESOURCE_STATE_COPY_DEST;
      list->ResourceBarrier(1, &cb);
      D3D12_TEXTURE_COPY_LOCATION srcLoc{}, dstLoc{};
      srcLoc.pResource = upload.Get();
      srcLoc.Type = D3D12_TEXTURE_COPY_TYPE_PLACED_FOOTPRINT;
      srcLoc.PlacedFootprint.Footprint.Format = rtvFmt;
      srcLoc.PlacedFootprint.Footprint.Width = W;
      srcLoc.PlacedFootprint.Footprint.Height = H;
      srcLoc.PlacedFootprint.Footprint.Depth = 1;
      srcLoc.PlacedFootprint.Footprint.RowPitch = (UINT)pitch;
      dstLoc.pResource = res[i].Get();
      dstLoc.Type = D3D12_TEXTURE_COPY_TYPE_SUBRESOURCE_INDEX;
      dstLoc.SubresourceIndex = 0;
      list->CopyTextureRegion(&dstLoc, 0, 0, 0, &srcLoc, nullptr);
      cb.Transition.StateBefore = D3D12_RESOURCE_STATE_COPY_DEST;
      cb.Transition.StateAfter = D3D12_RESOURCE_STATE_COMMON;
      list->ResourceBarrier(1, &cb);
      HRCHECK(list->Close());
      ID3D12CommandList* lists[] = {list.Get()};
      queue->ExecuteCommandLists(1, lists);
      HRCHECK(queue->Signal(fence.Get(), ++fenceVal));
      if (fence->GetCompletedValue() < fenceVal) {
        fence->SetEventOnCompletion(fenceVal, fenceEvt);
        WaitForSingleObject(fenceEvt, 2000);
      }
    }
    haveImage = true;
    printf("static image %s loaded into %d surfaces\n", imagePath, NBUF);
  }

  // --- modeset + power on ---
  st = NvAPI_D3D_DirectModeSetDisplayMode(&display, &mode);
  if (st == NVAPI_HDCP_DISABLED) { printf("SetDisplayMode: HDCP disabled, ok\n"); st = NVAPI_OK; }
  printf("SetDisplayMode = 0x%x\n", st);
  if (st != NVAPI_OK) return 3;
  st = NvAPI_DISP_DirectModeDisplayControl(&display, NV_DM_DISPLAY_CONTROL_POWER_ON);
  printf("DisplayControl(POWER_ON) = 0x%x\n", st);

  HANDLE waitEvt = NULL;
  st = NvAPI_DISP_DirectModeGetPresentWaitableObject(
      &display, reinterpret_cast<NvNotifyHandle*>(&waitEvt));
  bool haveWaitable = st == NVAPI_OK;
  printf("GetPresentWaitableObject = 0x%x\n", st);

  // --- present loop: animated dim hue clear ---
  auto t0 = std::chrono::steady_clock::now();
  uint64_t frames = 0, presentErrors = 0, waitTimeouts = 0;
  while (true) {
    double t = std::chrono::duration<double>(std::chrono::steady_clock::now() - t0).count();
    if (t >= seconds) break;
    int i = (int)(frames % NBUF);

    HRCHECK(alloc->Reset());
    HRCHECK(list->Reset(alloc.Get(), nullptr));
    if (!haveImage) {
      D3D12_RESOURCE_BARRIER b{};
      b.Type = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION;
      b.Transition.pResource = res[i].Get();
      b.Transition.Subresource = D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES;
      b.Transition.StateBefore = D3D12_RESOURCE_STATE_COMMON;
      b.Transition.StateAfter = D3D12_RESOURCE_STATE_RENDER_TARGET;
      list->ResourceBarrier(1, &b);
      // structure-revealing pattern: left eye green field, right eye red field,
      // synchronized white bar sweeping across each eye. Any DSC/packing fault
      // shows immediately as wrong colors per eye or a broken/duplicated bar.
      {
        LONG W = (LONG)rdesc.Width, H = (LONG)rdesc.Height, half = W / 2;
        float lc[4] = {0.0f, 0.25f, 0.08f, 1.0f};
        float rc[4] = {0.25f, 0.0f, 0.08f, 1.0f};
        float wc[4] = {0.4f, 0.4f, 0.4f, 1.0f};
        D3D12_RECT rl{0, 0, half, H}, rr{half, 0, W, H};
        list->ClearRenderTargetView(rtv[i], lc, 1, &rl);
        list->ClearRenderTargetView(rtv[i], rc, 1, &rr);
        LONG barW = 64;
        LONG bx = (LONG)(fmod(t * 0.25, 1.0) * (half - barW));
        D3D12_RECT bars[2] = {{bx, 0, bx + barW, H},
                              {half + bx, 0, half + bx + barW, H}};
        list->ClearRenderTargetView(rtv[i], wc, 2, bars);
      }
      D3D12_RESOURCE_BARRIER b2 = b;
      b2.Transition.StateBefore = D3D12_RESOURCE_STATE_RENDER_TARGET;
      b2.Transition.StateAfter = D3D12_RESOURCE_STATE_COMMON;
      list->ResourceBarrier(1, &b2);
    }
    // --image: surfaces pre-filled; submit the empty list to keep the
    // validated present pacing structure identical.
    HRCHECK(list->Close());
    ID3D12CommandList* lists[] = {list.Get()};
    queue->ExecuteCommandLists(1, lists);
    HRCHECK(queue->Signal(fence.Get(), ++fenceVal));
    if (fence->GetCompletedValue() < fenceVal) {
      fence->SetEventOnCompletion(fenceVal, fenceEvt);
      WaitForSingleObject(fenceEvt, 1000);
    }

    st = NvAPI_D3D_DirectModePresent12(&display, surf[i],
                                       NV_DIRECTMODE_PRESENT_FLAG_VSYNC,
                                       queue.Get(), list.Get());
    if (st != NVAPI_OK) {
      presentErrors++;
      fprintf(stderr, "Present12 = 0x%x at frame %llu\n", st,
              (unsigned long long)frames);
      if (presentErrors > 10) break;
    }
    if (haveWaitable && WaitForSingleObject(waitEvt, 1000) == WAIT_TIMEOUT)
      waitTimeouts++;
    frames++;
    if (frames % 300 == 0)
      printf("  t=%6.1fs frames=%llu avg=%.2f fps\n", t,
             (unsigned long long)frames, frames / t);
  }

  double total = std::chrono::duration<double>(std::chrono::steady_clock::now() - t0).count();
  printf("PRESENT RESULT: %llu frames in %.1fs = %.2f fps avg, presentErrors %llu, waitTimeouts %llu\n",
         (unsigned long long)frames, total, frames / total,
         (unsigned long long)presentErrors, (unsigned long long)waitTimeouts);

  if (powerOff)
    NvAPI_DISP_DirectModeDisplayControl(&display, NV_DM_DISPLAY_CONTROL_POWER_OFF);
  NvAPI_D3D_ReleaseDirectModeDisplay(VENDOR_ID, &display);
  printf("display released\n");
  // NvAPI_Unload deliberately skipped: unloading before ComPtr/D3D teardown
  // crashes in driver DLL atexit order. Process exit cleans up regardless.

  double expectHz = modeHz > 1 ? modeHz : 75.0;
  bool pass = frames > 0 && presentErrors == 0 && frames / total > expectHz * 0.93;
  printf("VERDICT: %s\n", pass ? "PASS" : (frames > 0 ? "MARGINAL" : "FAIL"));
  return 0;
}