// svr_vrtest: minimal OpenVR D3D11 harness for the skyrim-vr-upscale
// injection path. Renders a procedural diagnostic pattern (fine checker,
// concentric rings, radial spokes) to both eyes at the runtime's
// recommended size and submits through IVRCompositor, which is where the
// vrperfkit fork intercepts and runs the SVR kernel. After a settle
// period it drives vrperfkit's own hotkeys via SendInput to capture a
// still, cycle the upscale method, and capture again, so one headless run
// leaves SVR/FSR/NIS/CAS captures of the identical frame on disk.
//
// Not Skyrim: any OpenVR D3D11 process + the same two DLLs is the
// generalization claim, and this is that process.
#include <d3d11.h>
#include <d3dcompiler.h>
#include <openvr.h>

#include <cstdio>
#include <cstring>

#pragma comment(lib, "d3d11.lib")
#pragma comment(lib, "d3dcompiler.lib")

static const char* kShader = R"(
Texture2D dummy : register(t0);
struct VSOut { float4 pos : SV_Position; float2 uv : TEXCOORD0; };
VSOut vs(uint id : SV_VertexID) {
  VSOut o;
  float2 xy = float2((id << 1) & 2, id & 2);
  o.pos = float4(xy * 2.0 - 1.0, 0.0, 1.0);
  o.uv = xy;
  return o;
}
cbuffer CB : register(b0) { float4 misc; }  // x = eye, y = time, zw = size
float4 ps(VSOut i) : SV_Target {
  float2 px = i.uv * misc.zw;
  float2 c = misc.zw * 0.5;
  float2 d = px - c;
  float r = length(d);
  // fine checker whose pitch grows with eccentricity bands
  float pitch = 2.0 + 6.0 * step(0.25 * misc.w, r) + 8.0 * step(0.55 * misc.w, r);
  float2 q = floor(px / pitch);
  float checker = fmod(q.x + q.y, 2.0);
  // concentric rings every 64 px and radial spokes every 15 degrees
  float ring = step(fmod(r, 64.0), 2.5);
  float ang = atan2(d.y, d.x);
  float spoke = step(fmod(abs(ang) * 57.29578, 15.0), 0.35);
  float3 base = lerp(float3(0.12, 0.12, 0.14), float3(0.88, 0.88, 0.9), checker);
  base = lerp(base, float3(1.0, 0.55, 0.1), ring);
  base = lerp(base, float3(0.1, 0.8, 1.0), spoke);
  // eye tint corner marker so L/R are distinguishable in captures
  if (px.x < 40 && px.y < 40) base = misc.x < 0.5 ? float3(1, 0, 0) : float3(0, 1, 0);
  return float4(base, 1.0);
}
)";

static void keytap(WORD vk) {
  // vrperfkit polls the down-state per frame (GetAsyncKeyState high bit), so
  // the key must stay down across several frames to be seen.
  INPUT in = {};
  in.type = INPUT_KEYBOARD;
  in.ki.wVk = vk;
  SendInput(1, &in, sizeof(INPUT));
  Sleep(80);
  in.ki.dwFlags = KEYEVENTF_KEYUP;
  SendInput(1, &in, sizeof(INPUT));
  Sleep(40);
}

static void dump_staging(ID3D11Device* dev, ID3D11DeviceContext* ctx,
                         ID3D11Texture2D* tex, uint32_t w, uint32_t h,
                         const char* path) {
  D3D11_TEXTURE2D_DESC sd = {};
  sd.Width = w;
  sd.Height = h;
  sd.MipLevels = 1;
  sd.ArraySize = 1;
  sd.Format = DXGI_FORMAT_R8G8B8A8_UNORM;
  sd.SampleDesc.Count = 1;
  sd.Usage = D3D11_USAGE_STAGING;
  sd.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
  ID3D11Texture2D* stg = nullptr;
  if (FAILED(dev->CreateTexture2D(&sd, nullptr, &stg))) return;
  ctx->CopyResource(stg, tex);
  D3D11_MAPPED_SUBRESOURCE m;
  if (SUCCEEDED(ctx->Map(stg, 0, D3D11_MAP_READ, 0, &m))) {
    FILE* f = fopen(path, "wb");
    if (f) {
      fprintf(f, "P6\n%u %u\n255\n", w, h);
      for (uint32_t y = 0; y < h; ++y) {
        const unsigned char* row = (const unsigned char*)m.pData + y * m.RowPitch;
        for (uint32_t x = 0; x < w; ++x) fwrite(row + x * 4, 1, 3, f);
      }
      fclose(f);
    }
    ctx->Unmap(stg, 0);
  }
  stg->Release();
}

int main(int argc, char** argv) {
  int frames = 2000;
  if (argc > 1) frames = atoi(argv[1]);
  bool selftest = (argc > 2 && strcmp(argv[2], "--selftest") == 0);

  vr::IVRSystem* sys = nullptr;
  uint32_t w = 1098, h = 1220;
  if (!selftest) {
    vr::EVRInitError err = vr::VRInitError_None;
    sys = vr::VR_Init(&err, vr::VRApplication_Scene);
    if (err != vr::VRInitError_None) {
      printf("VR_Init failed: %s\n",
             vr::VR_GetVRInitErrorAsEnglishDescription(err));
      return 1;
    }
    sys->GetRecommendedRenderTargetSize(&w, &h);
    printf("recommended per-eye: %ux%u\n", w, h);
  }

  ID3D11Device* dev = nullptr;
  ID3D11DeviceContext* ctx = nullptr;
  D3D_FEATURE_LEVEL fl;
  if (FAILED(D3D11CreateDevice(nullptr, D3D_DRIVER_TYPE_HARDWARE, nullptr, 0,
                               nullptr, 0, D3D11_SDK_VERSION, &dev, &fl,
                               &ctx))) {
    printf("D3D11CreateDevice failed\n");
    return 1;
  }

  D3D11_TEXTURE2D_DESC td = {};
  td.Width = w;
  td.Height = h;
  td.MipLevels = 1;
  td.ArraySize = 1;
  td.Format = DXGI_FORMAT_R8G8B8A8_UNORM;
  td.SampleDesc.Count = 1;
  td.Usage = D3D11_USAGE_DEFAULT;
  td.BindFlags = D3D11_BIND_RENDER_TARGET | D3D11_BIND_SHADER_RESOURCE;
  ID3D11Texture2D* eye[2] = {};
  ID3D11RenderTargetView* rtv[2] = {};
  for (int e = 0; e < 2; ++e) {
    if (FAILED(dev->CreateTexture2D(&td, nullptr, &eye[e]))) return 1;
    if (FAILED(dev->CreateRenderTargetView(eye[e], nullptr, &rtv[e]))) return 1;
  }

  ID3DBlob *vsb = nullptr, *psb = nullptr, *errb = nullptr;
  if (FAILED(D3DCompile(kShader, strlen(kShader), nullptr, nullptr, nullptr,
                        "vs", "vs_5_0", 0, 0, &vsb, &errb))) {
    printf("vs: %s\n", errb ? (char*)errb->GetBufferPointer() : "?");
    return 1;
  }
  if (FAILED(D3DCompile(kShader, strlen(kShader), nullptr, nullptr, nullptr,
                        "ps", "ps_5_0", 0, 0, &psb, &errb))) {
    printf("ps: %s\n", errb ? (char*)errb->GetBufferPointer() : "?");
    return 1;
  }
  ID3D11VertexShader* vs = nullptr;
  ID3D11PixelShader* ps = nullptr;
  HRESULT hrv = dev->CreateVertexShader(vsb->GetBufferPointer(),
                                        vsb->GetBufferSize(), nullptr, &vs);
  HRESULT hrp = dev->CreatePixelShader(psb->GetBufferPointer(),
                                       psb->GetBufferSize(), nullptr, &ps);
  if (FAILED(hrv) || FAILED(hrp) || !vs || !ps) {
    printf("shader create failed: vs=0x%08lx ps=0x%08lx\n", hrv, hrp);
    return 1;
  }
  D3D11_BUFFER_DESC bd = {};
  bd.ByteWidth = 16;
  bd.Usage = D3D11_USAGE_DEFAULT;
  bd.BindFlags = D3D11_BIND_CONSTANT_BUFFER;
  ID3D11Buffer* cb = nullptr;
  if (FAILED(dev->CreateBuffer(&bd, nullptr, &cb))) {
    printf("cbuffer create failed\n");
    return 1;
  }
  // No culling: the fullscreen triangle is CCW in NDC, which the default
  // rasterizer (cull back) would discard entirely.
  D3D11_RASTERIZER_DESC rd = {};
  rd.FillMode = D3D11_FILL_SOLID;
  rd.CullMode = D3D11_CULL_NONE;
  ID3D11RasterizerState* rs = nullptr;
  if (FAILED(dev->CreateRasterizerState(&rd, &rs))) {
    printf("rasterizer create failed\n");
    return 1;
  }
  ctx->RSSetState(rs);

  auto render_eye = [&](int e, int f) {
    float misc[4] = {(float)e, f * 0.011f, (float)w, (float)h};
    ctx->UpdateSubresource(cb, 0, nullptr, misc, 0, 0);
    ctx->OMSetRenderTargets(1, &rtv[e], nullptr);
    D3D11_VIEWPORT vp = {0, 0, (float)w, (float)h, 0, 1};
    ctx->RSSetViewports(1, &vp);
    ctx->RSSetState(rs);
    ctx->IASetPrimitiveTopology(D3D11_PRIMITIVE_TOPOLOGY_TRIANGLELIST);
    ctx->IASetInputLayout(nullptr);
    ctx->VSSetShader(vs, nullptr, 0);
    ctx->VSSetConstantBuffers(0, 1, &cb);
    ctx->PSSetShader(ps, nullptr, 0);
    ctx->PSSetConstantBuffers(0, 1, &cb);
    ctx->Draw(3, 0);
  };

  if (selftest) {
    render_eye(0, 100);
    ctx->Flush();
    dump_staging(dev, ctx, eye[0], w, h, "D:\\svr_vrtest\\dump\\selftest.ppm");
    printf("selftest: wrote dump/selftest.ppm (%ux%u)\n", w, h);
    return 0;
  }

  vr::TrackedDevicePose_t poses[vr::k_unMaxTrackedDeviceCount];
  for (int f = 0; f < frames; ++f) {
    vr::VRCompositor()->WaitGetPoses(poses, vr::k_unMaxTrackedDeviceCount,
                                     nullptr, 0);
    for (int e = 0; e < 2; ++e) {
      render_eye(e, f);
      vr::Texture_t t = {eye[e], vr::TextureType_DirectX,
                         vr::ColorSpace_Gamma};
      vr::VRCompositor()->Submit(e == 0 ? vr::Eye_Left : vr::Eye_Right, &t);
    }
  }
  printf("done\n");
  vr::VR_Shutdown();
  return 0;
}