// warp_check — M2 gate 5 (Gate A) harness.
//
// Loads two eye textures (raw R8G8B8A8, eye_w x eye_h), runs the WarpPass
// into an offscreen panel surface (panel_w x panel_h, left half = left
// eye), and writes the result as raw R8G8B8A8 for the python pixel-diff
// (tools/gate5_check.py) against the S3 generator output.
//
// Usage: warp_check --config <config.json> --left <l.raw> --right <r.raw>
//                   --out <panel.raw> [--linear] [--color-mult]
//        (default point sampling + no color mult = gate-A mode)

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

#include <cstdio>
#include <cstring>
#include <string>
#include <vector>

#include "calib/hmd_config.h"
#include "render/warp_pass.h"

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

namespace {

constexpr uint32_t kPanelW = 5088, kPanelH = 2544;

bool readFile(const std::string& path, std::vector<uint8_t>* out,
              size_t expect) {
  FILE* f = fopen(path.c_str(), "rb");
  if (!f) {
    fprintf(stderr, "cannot open %s\n", path.c_str());
    return false;
  }
  out->resize(expect);
  size_t got = fread(out->data(), 1, expect, f);
  fclose(f);
  if (got != expect) {
    fprintf(stderr, "%s: %zu bytes, expected %zu\n", path.c_str(), got, expect);
    return false;
  }
  return true;
}

}  // namespace

int main(int argc, char** argv) {
  std::string configPath, leftPath, rightPath, outPath;
  bool linear = false, colorMult = false;
  for (int i = 1; i < argc; i++) {
    if (!strcmp(argv[i], "--config") && i + 1 < argc) configPath = argv[++i];
    else if (!strcmp(argv[i], "--left") && i + 1 < argc) leftPath = argv[++i];
    else if (!strcmp(argv[i], "--right") && i + 1 < argc) rightPath = argv[++i];
    else if (!strcmp(argv[i], "--out") && i + 1 < argc) outPath = argv[++i];
    else if (!strcmp(argv[i], "--linear")) linear = true;
    else if (!strcmp(argv[i], "--color-mult")) colorMult = true;
    else {
      fprintf(stderr, "usage: warp_check --config c.json --left l.raw "
                      "--right r.raw --out p.raw [--linear] [--color-mult]\n");
      return 2;
    }
  }
  HmdConfig cfg;
  std::string err;
  if (!LoadHmdConfig(configPath, &cfg, &err)) {
    fprintf(stderr, "%s\n", err.c_str());
    return 1;
  }
  const uint32_t eyeW = cfg.eye_width_px, eyeH = cfg.eye_height_px;
  std::vector<uint8_t> eyeData[2];
  if (!readFile(leftPath, &eyeData[0], (size_t)eyeW * eyeH * 4) ||
      !readFile(rightPath, &eyeData[1], (size_t)eyeW * eyeH * 4))
    return 1;

  ComPtr<ID3D12Device> dev;
  if (FAILED(D3D12CreateDevice(nullptr, D3D_FEATURE_LEVEL_11_0,
                               IID_PPV_ARGS(&dev)))) {
    fprintf(stderr, "D3D12CreateDevice failed\n");
    return 1;
  }
  ComPtr<ID3D12CommandQueue> queue;
  D3D12_COMMAND_QUEUE_DESC qd{D3D12_COMMAND_LIST_TYPE_DIRECT};
  dev->CreateCommandQueue(&qd, IID_PPV_ARGS(&queue));
  ComPtr<ID3D12CommandAllocator> alloc;
  dev->CreateCommandAllocator(D3D12_COMMAND_LIST_TYPE_DIRECT,
                              IID_PPV_ARGS(&alloc));
  ComPtr<ID3D12GraphicsCommandList> list;
  dev->CreateCommandList(0, D3D12_COMMAND_LIST_TYPE_DIRECT, alloc.Get(),
                         nullptr, IID_PPV_ARGS(&list));

  const DXGI_FORMAT kFmt = DXGI_FORMAT_R8G8B8A8_UNORM;

  // Eye textures + upload.
  ComPtr<ID3D12Resource> eyeTex[2], staging[2];
  D3D12_RESOURCE_DESC td{};
  td.Dimension = D3D12_RESOURCE_DIMENSION_TEXTURE2D;
  td.Width = eyeW;
  td.Height = eyeH;
  td.DepthOrArraySize = 1;
  td.MipLevels = 1;
  td.Format = kFmt;
  td.SampleDesc.Count = 1;
  D3D12_HEAP_PROPERTIES def{D3D12_HEAP_TYPE_DEFAULT};
  D3D12_HEAP_PROPERTIES up{D3D12_HEAP_TYPE_UPLOAD};
  D3D12_PLACED_SUBRESOURCE_FOOTPRINT fp{};
  UINT64 upSize = 0;
  dev->GetCopyableFootprints(&td, 0, 1, 0, &fp, nullptr, nullptr, &upSize);
  for (int e = 0; e < 2; e++) {
    dev->CreateCommittedResource(&def, D3D12_HEAP_FLAG_NONE, &td,
                                 D3D12_RESOURCE_STATE_COPY_DEST, nullptr,
                                 IID_PPV_ARGS(&eyeTex[e]));
    D3D12_RESOURCE_DESC bd{};
    bd.Dimension = D3D12_RESOURCE_DIMENSION_BUFFER;
    bd.Width = upSize;
    bd.Height = 1;
    bd.DepthOrArraySize = 1;
    bd.MipLevels = 1;
    bd.SampleDesc.Count = 1;
    bd.Layout = D3D12_TEXTURE_LAYOUT_ROW_MAJOR;
    dev->CreateCommittedResource(&up, D3D12_HEAP_FLAG_NONE, &bd,
                                 D3D12_RESOURCE_STATE_GENERIC_READ, nullptr,
                                 IID_PPV_ARGS(&staging[e]));
    uint8_t* p = nullptr;
    staging[e]->Map(0, nullptr, (void**)&p);
    for (uint32_t y = 0; y < eyeH; y++)
      memcpy(p + fp.Offset + y * fp.Footprint.RowPitch,
             &eyeData[e][(size_t)y * eyeW * 4], eyeW * 4);
    staging[e]->Unmap(0, nullptr);
    D3D12_TEXTURE_COPY_LOCATION src{staging[e].Get(),
                                    D3D12_TEXTURE_COPY_TYPE_PLACED_FOOTPRINT};
    src.PlacedFootprint = fp;
    D3D12_TEXTURE_COPY_LOCATION dst{eyeTex[e].Get(),
                                    D3D12_TEXTURE_COPY_TYPE_SUBRESOURCE_INDEX};
    list->CopyTextureRegion(&dst, 0, 0, 0, &src, nullptr);
    D3D12_RESOURCE_BARRIER bar{};
    bar.Transition.pResource = eyeTex[e].Get();
    bar.Transition.StateBefore = D3D12_RESOURCE_STATE_COPY_DEST;
    bar.Transition.StateAfter = D3D12_RESOURCE_STATE_PIXEL_SHADER_RESOURCE;
    bar.Transition.Subresource = D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES;
    list->ResourceBarrier(1, &bar);
  }

  // SRV heap.
  ComPtr<ID3D12DescriptorHeap> srvHeap;
  D3D12_DESCRIPTOR_HEAP_DESC hd{};
  hd.Type = D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV;
  hd.NumDescriptors = 2;
  hd.Flags = D3D12_DESCRIPTOR_HEAP_FLAG_SHADER_VISIBLE;
  dev->CreateDescriptorHeap(&hd, IID_PPV_ARGS(&srvHeap));
  const UINT srvStep = dev->GetDescriptorHandleIncrementSize(
      D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV);
  for (int e = 0; e < 2; e++) {
    D3D12_SHADER_RESOURCE_VIEW_DESC sv{};
    sv.Format = kFmt;
    sv.ViewDimension = D3D12_SRV_DIMENSION_TEXTURE2D;
    sv.Shader4ComponentMapping = D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING;
    sv.Texture2D.MipLevels = 1;
    D3D12_CPU_DESCRIPTOR_HANDLE h = srvHeap->GetCPUDescriptorHandleForHeapStart();
    h.ptr += e * srvStep;
    dev->CreateShaderResourceView(eyeTex[e].Get(), &sv, h);
  }

  // Output panel RT.
  ComPtr<ID3D12Resource> panel;
  td.Width = kPanelW;
  td.Height = kPanelH;
  td.Flags = D3D12_RESOURCE_FLAG_ALLOW_RENDER_TARGET;
  dev->CreateCommittedResource(&def, D3D12_HEAP_FLAG_NONE, &td,
                               D3D12_RESOURCE_STATE_RENDER_TARGET, nullptr,
                               IID_PPV_ARGS(&panel));
  ComPtr<ID3D12DescriptorHeap> rtvHeap;
  D3D12_DESCRIPTOR_HEAP_DESC rh{D3D12_DESCRIPTOR_HEAP_TYPE_RTV, 1};
  dev->CreateDescriptorHeap(&rh, IID_PPV_ARGS(&rtvHeap));
  D3D12_CPU_DESCRIPTOR_HANDLE rtv = rtvHeap->GetCPUDescriptorHandleForHeapStart();
  dev->CreateRenderTargetView(panel.Get(), nullptr, rtv);

  WarpPass warp;
  if (!warp.init(dev.Get(), kFmt, cfg)) return 1;

  list->OMSetRenderTargets(1, &rtv, FALSE, nullptr);
  ID3D12DescriptorHeap* heaps[] = {srvHeap.Get()};
  list->SetDescriptorHeaps(1, heaps);
  for (int e = 0; e < 2; e++) {
    D3D12_GPU_DESCRIPTOR_HANDLE h = srvHeap->GetGPUDescriptorHandleForHeapStart();
    h.ptr += e * srvStep;
    warp.record(list.Get(), e, h, kPanelW, kPanelH, /*pointSample=*/!linear,
                colorMult);
  }

  // Readback.
  D3D12_RESOURCE_DESC pd = panel->GetDesc();
  D3D12_PLACED_SUBRESOURCE_FOOTPRINT pfp{};
  UINT64 total = 0;
  dev->GetCopyableFootprints(&pd, 0, 1, 0, &pfp, nullptr, nullptr, &total);
  ComPtr<ID3D12Resource> rb;
  D3D12_HEAP_PROPERTIES rbh{D3D12_HEAP_TYPE_READBACK};
  D3D12_RESOURCE_DESC bd{};
  bd.Dimension = D3D12_RESOURCE_DIMENSION_BUFFER;
  bd.Width = total;
  bd.Height = 1;
  bd.DepthOrArraySize = 1;
  bd.MipLevels = 1;
  bd.SampleDesc.Count = 1;
  bd.Layout = D3D12_TEXTURE_LAYOUT_ROW_MAJOR;
  dev->CreateCommittedResource(&rbh, D3D12_HEAP_FLAG_NONE, &bd,
                               D3D12_RESOURCE_STATE_COPY_DEST, nullptr,
                               IID_PPV_ARGS(&rb));
  D3D12_RESOURCE_BARRIER bar{};
  bar.Transition.pResource = panel.Get();
  bar.Transition.StateBefore = D3D12_RESOURCE_STATE_RENDER_TARGET;
  bar.Transition.StateAfter = D3D12_RESOURCE_STATE_COPY_SOURCE;
  bar.Transition.Subresource = D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES;
  list->ResourceBarrier(1, &bar);
  D3D12_TEXTURE_COPY_LOCATION src{panel.Get(),
                                  D3D12_TEXTURE_COPY_TYPE_SUBRESOURCE_INDEX};
  D3D12_TEXTURE_COPY_LOCATION dst{rb.Get(),
                                  D3D12_TEXTURE_COPY_TYPE_PLACED_FOOTPRINT};
  dst.PlacedFootprint = pfp;
  list->CopyTextureRegion(&dst, 0, 0, 0, &src, nullptr);
  list->Close();
  ID3D12CommandList* lists[] = {list.Get()};
  queue->ExecuteCommandLists(1, lists);
  ComPtr<ID3D12Fence> fence;
  dev->CreateFence(0, D3D12_FENCE_FLAG_NONE, IID_PPV_ARGS(&fence));
  queue->Signal(fence.Get(), 1);
  HANDLE ev = CreateEventW(nullptr, FALSE, FALSE, nullptr);
  fence->SetEventOnCompletion(1, ev);
  WaitForSingleObject(ev, 10000);
  CloseHandle(ev);

  uint8_t* p = nullptr;
  rb->Map(0, nullptr, (void**)&p);
  FILE* f = fopen(outPath.c_str(), "wb");
  if (!f) {
    fprintf(stderr, "cannot write %s\n", outPath.c_str());
    return 1;
  }
  for (uint32_t y = 0; y < kPanelH; y++)
    fwrite(p + y * pfp.Footprint.RowPitch, 1, (size_t)kPanelW * 4, f);
  fclose(f);
  rb->Unmap(0, nullptr);
  printf("wrote %s (%ux%u RGBA, %s sampling%s)\n", outPath.c_str(), kPanelW,
         kPanelH, linear ? "linear" : "point",
         colorMult ? ", color mult" : "");
  return 0;
}