// hlsl_check: run svr_upscale.hlsl on a source PPM and write the result PPM,
// so it can be compared against the CUDA kernel at identical parameters.
// Fixed params: center (0.5,0.5), r1 = 0.65*oh, r0 = 0.55*r1, sharp = 0.4.
#include <d3d11.h>
#include <d3dcompiler.h>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#pragma comment(lib, "d3d11.lib")
#pragma comment(lib, "d3dcompiler.lib")

struct SvrCB { uint32_t sw, sh, dw, dh; float cx, cy, r0, r1; float sharp, p0, p1, p2; };

static unsigned char* read_ppm(const char* p, int* w, int* h) {
  FILE* f = fopen(p, "rb");
  if (!f) return nullptr;
  char m[3] = {}; fscanf(f, "%2s", m); fscanf(f, "%d %d", w, h);
  int mx; fscanf(f, "%d", &mx); fgetc(f);
  int n = (*w) * (*h);
  unsigned char* rgb = (unsigned char*)malloc(n * 3);
  fread(rgb, 1, n * 3, f); fclose(f);
  unsigned char* rgba = (unsigned char*)malloc((size_t)n * 4);
  for (int i = 0; i < n; ++i) {
    rgba[i*4+0]=rgb[i*3+0]; rgba[i*4+1]=rgb[i*3+1]; rgba[i*4+2]=rgb[i*3+2]; rgba[i*4+3]=255;
  }
  free(rgb);
  return rgba;
}

int main(int argc, char** argv) {
  if (argc < 5) { printf("usage: hlsl_check src.ppm ow oh out.ppm\n"); return 1; }
  int sw, sh;
  unsigned char* src = read_ppm(argv[1], &sw, &sh);
  if (!src) { printf("read fail\n"); return 1; }
  int ow = atoi(argv[2]), oh = atoi(argv[3]);

  ID3D11Device* dev = nullptr; ID3D11DeviceContext* ctx = nullptr; D3D_FEATURE_LEVEL fl;
  D3D11CreateDevice(nullptr, D3D_DRIVER_TYPE_HARDWARE, nullptr, 0, nullptr, 0,
                    D3D11_SDK_VERSION, &dev, &fl, &ctx);

  const char* cands[] = { getenv("SVR_HLSL"), "svr_upscale.hlsl",
                          "..\\game\\svr_upscale.hlsl", "game\\svr_upscale.hlsl" };
  FILE* sf = nullptr;
  for (const char* c : cands) if (c && (sf = fopen(c, "rb"))) break;
  if (!sf) { printf("cannot find svr_upscale.hlsl (set SVR_HLSL)\n"); return 1; }
  fseek(sf, 0, SEEK_END); long n = ftell(sf); fseek(sf, 0, SEEK_SET);
  char* csrc = (char*)malloc(n + 1); fread(csrc, 1, n, sf); csrc[n] = 0; fclose(sf);
  ID3DBlob *cb = nullptr, *eb = nullptr;
  if (FAILED(D3DCompile(csrc, n, 0, 0, 0, "main", "cs_5_0", 0, 0, &cb, &eb))) {
    printf("compile: %s\n", eb ? (char*)eb->GetBufferPointer() : "?"); return 1; }
  ID3D11ComputeShader* cs = nullptr;
  dev->CreateComputeShader(cb->GetBufferPointer(), cb->GetBufferSize(), 0, &cs);

  // source texture (immutable), full-res output UAV, staging readback
  D3D11_TEXTURE2D_DESC td = {};
  td.Width = sw; td.Height = sh; td.MipLevels = 1; td.ArraySize = 1;
  td.Format = DXGI_FORMAT_R8G8B8A8_UNORM; td.SampleDesc.Count = 1;
  td.Usage = D3D11_USAGE_IMMUTABLE; td.BindFlags = D3D11_BIND_SHADER_RESOURCE;
  D3D11_SUBRESOURCE_DATA sd = {src, (UINT)sw * 4, 0};
  ID3D11Texture2D* stex = nullptr; dev->CreateTexture2D(&td, &sd, &stex);
  ID3D11ShaderResourceView* srv = nullptr; dev->CreateShaderResourceView(stex, nullptr, &srv);

  td.Width = ow; td.Height = oh; td.Usage = D3D11_USAGE_DEFAULT;
  td.BindFlags = D3D11_BIND_UNORDERED_ACCESS | D3D11_BIND_SHADER_RESOURCE;
  ID3D11Texture2D* otex = nullptr; dev->CreateTexture2D(&td, nullptr, &otex);
  ID3D11UnorderedAccessView* uav = nullptr; dev->CreateUnorderedAccessView(otex, nullptr, &uav);

  td.Usage = D3D11_USAGE_STAGING; td.BindFlags = 0;
  td.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
  ID3D11Texture2D* stg = nullptr; dev->CreateTexture2D(&td, nullptr, &stg);

  SvrCB u = {(uint32_t)sw,(uint32_t)sh,(uint32_t)ow,(uint32_t)oh,
             0.5f*ow, 0.5f*oh, 0.55f*0.65f*oh, 0.65f*oh, 0.4f, 0,0,0};
  D3D11_BUFFER_DESC bd = {}; bd.ByteWidth = sizeof(SvrCB);
  bd.Usage = D3D11_USAGE_DEFAULT; bd.BindFlags = D3D11_BIND_CONSTANT_BUFFER;
  D3D11_SUBRESOURCE_DATA ud = {&u, 0, 0};
  ID3D11Buffer* ucb = nullptr; dev->CreateBuffer(&bd, &ud, &ucb);

  ctx->CSSetShader(cs, 0, 0);
  ctx->CSSetConstantBuffers(0, 1, &ucb);
  ctx->CSSetShaderResources(0, 1, &srv);
  ctx->CSSetUnorderedAccessViews(0, 1, &uav, nullptr);
  ctx->Dispatch((ow + 7) / 8, (oh + 7) / 8, 1);
  ctx->CopyResource(stg, otex);

  D3D11_MAPPED_SUBRESOURCE m;
  ctx->Map(stg, 0, D3D11_MAP_READ, 0, &m);
  FILE* of = fopen(argv[4], "wb");
  fprintf(of, "P6\n%d %d\n255\n", ow, oh);
  for (int y = 0; y < oh; ++y) {
    unsigned char* row = (unsigned char*)m.pData + y * m.RowPitch;
    for (int x = 0; x < ow; ++x) fwrite(row + x * 4, 1, 3, of);
  }
  fclose(of);
  ctx->Unmap(stg, 0);
  printf("wrote %s (%dx%d)\n", argv[4], ow, oh);
  return 0;
}