#include "render/scene_renderer.h"

#include <d3dcompiler.h>

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

using Microsoft::WRL::ComPtr;

namespace sauna {
namespace {

const char kShader[] = R"(
cbuffer C : register(b0) { row_major float4x4 mvp; };
Texture2D tex : register(t0);
SamplerState smp : register(s0);

struct VSOut { float4 pos : SV_Position; float2 uv : TEXCOORD0; };

VSOut vsmain(float3 pos : POSITION, float2 uv : TEXCOORD0) {
  VSOut o;
  o.pos = mul(mvp, float4(pos, 1.0));
  o.uv = uv;
  return o;
}

float4 psmain(VSOut i) : SV_Target {
  return tex.Sample(smp, i.uv);
}
)";

struct Vertex {
  float pos[3];
  float uv[2];
};

// Unit quad (1x1 centered at origin, z=0); every draw supplies a full MVP
// (viewProj * model) to place it in the world. UV origin top-left (v down):
// uv.y=0 is the quad top. The M2 test-quad placement (2.133 x 1.2 m at
// z = -2 m, same angular size as the original 1.6x0.9 @ 1.5 m) lives in
// the record() wrapper's default model.
constexpr float kQuadW = kSceneQuadW, kQuadH = kSceneQuadH;
const Vertex kQuad[6] = {
    {{-0.5f, +0.5f, 0.0f}, {0, 0}},
    {{+0.5f, +0.5f, 0.0f}, {1, 0}},
    {{-0.5f, -0.5f, 0.0f}, {0, 1}},
    {{-0.5f, -0.5f, 0.0f}, {0, 1}},
    {{+0.5f, +0.5f, 0.0f}, {1, 0}},
    {{+0.5f, -0.5f, 0.0f}, {1, 1}},
};

constexpr uint32_t kTexSize = 1024;

// Procedural test card — every feature is drawn in WORLD-space coordinates
// (meters on the quad) so eyes-in geometry is self-judging: 0.1 m SQUARE
// checkers, a perfectly ROUND 0.6 m circle (any distortion/aspect error
// shows as an oval), red top band + up triangle, blue left band, center
// crosshair. The texture itself is square while the quad is 16:9, so all
// drawing maps through (meters -> texel) per-axis scales.
void makeTestCard(std::vector<uint32_t>* px) {
  px->resize(kTexSize * kTexSize);
  // texels per meter, per axis (the aspect correction).
  const float sx = kTexSize / kQuadW, sy = kTexSize / kQuadH;
  auto at = [&](uint32_t x, uint32_t y) -> uint32_t& {
    return (*px)[y * kTexSize + x];
  };
  for (uint32_t y = 0; y < kTexSize; y++)
    for (uint32_t x = 0; x < kTexSize; x++) {
      // World meters from the quad's top-left corner.
      const float wx = x / sx, wy = y / sy;
      const bool c = (((int)(wx / 0.1f) + (int)(wy / 0.1f)) & 1) != 0;
      at(x, y) = c ? 0xFFB0B0B0 : 0xFFFFFFFF;  // ABGR (R8G8B8A8)
      // Round circle: ring at radius 0.30 m +/- 0.01 m around quad center.
      const float dx = wx - kQuadW / 2, dy = wy - kQuadH / 2;
      const float r = sqrtf(dx * dx + dy * dy);
      if (r > 0.29f && r < 0.31f) at(x, y) = 0xFF000000;
    }
  const uint32_t bandY = (uint32_t)(0.04f * sy);  // 4 cm bands
  const uint32_t bandX = (uint32_t)(0.04f * sx);
  for (uint32_t y = 0; y < bandY; y++)
    for (uint32_t x = 0; x < kTexSize; x++) at(x, y) = 0xFF2020D0;  // red top
  for (uint32_t y = 0; y < kTexSize; y++)
    for (uint32_t x = 0; x < bandX; x++) at(x, y) = 0xFFD02020;  // blue left
  // Up triangle: apex 6 cm below top, 8 cm tall, world-proportioned.
  const uint32_t triTop = (uint32_t)(0.06f * sy);
  const uint32_t triBot = (uint32_t)(0.14f * sy);
  for (uint32_t y = triTop; y < triBot; y++) {
    const float frac = (float)(y - triTop) / (triBot - triTop);
    const uint32_t halfw = (uint32_t)(frac * 0.05f * sx);
    for (uint32_t x = kTexSize / 2 - halfw; x <= kTexSize / 2 + halfw; x++)
      at(x, y) = 0xFF2020D0;
  }
  // Crosshair: 16 cm arms, 6 mm thick, world-proportioned.
  const uint32_t armX = (uint32_t)(0.08f * sx), armY = (uint32_t)(0.08f * sy);
  const uint32_t thX = (uint32_t)(0.003f * sx), thY = (uint32_t)(0.003f * sy);
  for (uint32_t x = kTexSize / 2 - armX; x < kTexSize / 2 + armX; x++)
    for (uint32_t t = 0; t < 2 * thY; t++)
      at(x, kTexSize / 2 - thY + t) = 0xFF000000;
  for (uint32_t y = kTexSize / 2 - armY; y < kTexSize / 2 + armY; y++)
    for (uint32_t t = 0; t < 2 * thX; t++)
      at(kTexSize / 2 - thX + t, y) = 0xFF000000;
  for (uint32_t i = 0; i < kTexSize; i++) {  // border
    for (uint32_t t = 0; t < 8; t++) {
      at(i, t) = at(i, kTexSize - 1 - t) = 0xFF000000;
      at(t, i) = at(kTexSize - 1 - t, i) = 0xFF000000;
    }
  }
}

bool compile(const char* entry, const char* target, ComPtr<ID3DBlob>* out) {
  ComPtr<ID3DBlob> err;
  HRESULT hr = D3DCompile(kShader, sizeof(kShader) - 1, nullptr, nullptr,
                          nullptr, entry, target, 0, 0, &*out, &err);
  if (FAILED(hr)) {
    fprintf(stderr, "shader %s: %s\n", entry,
            err ? (const char*)err->GetBufferPointer() : "compile failed");
    return false;
  }
  return true;
}

}  // namespace

bool SceneRenderer::init(ID3D12Device* dev, DXGI_FORMAT rtFormat) {
  // Root signature: 16 root constants (mvp) + SRV table + static sampler.
  D3D12_DESCRIPTOR_RANGE range{};
  range.RangeType = D3D12_DESCRIPTOR_RANGE_TYPE_SRV;
  range.NumDescriptors = 1;
  D3D12_ROOT_PARAMETER params[2]{};
  params[0].ParameterType = D3D12_ROOT_PARAMETER_TYPE_32BIT_CONSTANTS;
  params[0].Constants.Num32BitValues = 16;
  params[0].ShaderVisibility = D3D12_SHADER_VISIBILITY_VERTEX;
  params[1].ParameterType = D3D12_ROOT_PARAMETER_TYPE_DESCRIPTOR_TABLE;
  params[1].DescriptorTable.NumDescriptorRanges = 1;
  params[1].DescriptorTable.pDescriptorRanges = &range;
  params[1].ShaderVisibility = D3D12_SHADER_VISIBILITY_PIXEL;
  // Anisotropic: the M3 desktop quad is viewed obliquely (canted cameras +
  // head rotation) and heavily minified — aniso keeps text legible where
  // trilinear smears it. Harmless for the test card.
  D3D12_STATIC_SAMPLER_DESC samp{};
  samp.Filter = D3D12_FILTER_ANISOTROPIC;
  samp.MaxAnisotropy = 8;
  samp.MaxLOD = D3D12_FLOAT32_MAX;
  samp.AddressU = samp.AddressV = samp.AddressW =
      D3D12_TEXTURE_ADDRESS_MODE_CLAMP;
  samp.ShaderVisibility = D3D12_SHADER_VISIBILITY_PIXEL;
  D3D12_ROOT_SIGNATURE_DESC rsDesc{};
  rsDesc.NumParameters = 2;
  rsDesc.pParameters = params;
  rsDesc.NumStaticSamplers = 1;
  rsDesc.pStaticSamplers = &samp;
  rsDesc.Flags = D3D12_ROOT_SIGNATURE_FLAG_ALLOW_INPUT_ASSEMBLER_INPUT_LAYOUT;
  ComPtr<ID3DBlob> sig, err;
  if (FAILED(D3D12SerializeRootSignature(&rsDesc, D3D_ROOT_SIGNATURE_VERSION_1,
                                         &sig, &err))) {
    fprintf(stderr, "root signature: %s\n",
            err ? (const char*)err->GetBufferPointer() : "failed");
    return false;
  }
  if (FAILED(dev->CreateRootSignature(0, sig->GetBufferPointer(),
                                      sig->GetBufferSize(),
                                      IID_PPV_ARGS(&rootSig_))))
    return false;

  ComPtr<ID3DBlob> vs, ps;
  if (!compile("vsmain", "vs_5_0", &vs) || !compile("psmain", "ps_5_0", &ps))
    return false;

  D3D12_INPUT_ELEMENT_DESC layout[] = {
      {"POSITION", 0, DXGI_FORMAT_R32G32B32_FLOAT, 0, 0,
       D3D12_INPUT_CLASSIFICATION_PER_VERTEX_DATA, 0},
      {"TEXCOORD", 0, DXGI_FORMAT_R32G32_FLOAT, 0, 12,
       D3D12_INPUT_CLASSIFICATION_PER_VERTEX_DATA, 0},
  };
  D3D12_GRAPHICS_PIPELINE_STATE_DESC pso{};
  pso.pRootSignature = rootSig_.Get();
  pso.VS = {vs->GetBufferPointer(), vs->GetBufferSize()};
  pso.PS = {ps->GetBufferPointer(), ps->GetBufferSize()};
  pso.BlendState.RenderTarget[0].RenderTargetWriteMask =
      D3D12_COLOR_WRITE_ENABLE_ALL;
  pso.SampleMask = UINT_MAX;
  pso.RasterizerState.FillMode = D3D12_FILL_MODE_SOLID;
  pso.RasterizerState.CullMode = D3D12_CULL_MODE_NONE;
  pso.RasterizerState.DepthClipEnable = TRUE;
  pso.InputLayout = {layout, 2};
  pso.PrimitiveTopologyType = D3D12_PRIMITIVE_TOPOLOGY_TYPE_TRIANGLE;
  pso.NumRenderTargets = 1;
  pso.RTVFormats[0] = rtFormat;
  pso.SampleDesc.Count = 1;
  if (FAILED(dev->CreateGraphicsPipelineState(&pso, IID_PPV_ARGS(&pso_))))
    return false;

  // Vertex buffer (upload heap — tiny, written once).
  D3D12_HEAP_PROPERTIES up{D3D12_HEAP_TYPE_UPLOAD};
  D3D12_RESOURCE_DESC bd{};
  bd.Dimension = D3D12_RESOURCE_DIMENSION_BUFFER;
  bd.Width = sizeof(kQuad);
  bd.Height = 1;
  bd.DepthOrArraySize = 1;
  bd.MipLevels = 1;
  bd.SampleDesc.Count = 1;
  bd.Layout = D3D12_TEXTURE_LAYOUT_ROW_MAJOR;
  if (FAILED(dev->CreateCommittedResource(&up, D3D12_HEAP_FLAG_NONE, &bd,
                                          D3D12_RESOURCE_STATE_GENERIC_READ,
                                          nullptr, IID_PPV_ARGS(&vb_))))
    return false;
  void* p = nullptr;
  vb_->Map(0, nullptr, &p);
  memcpy(p, kQuad, sizeof(kQuad));
  vb_->Unmap(0, nullptr);
  vbv_ = {vb_->GetGPUVirtualAddress(), sizeof(kQuad), sizeof(Vertex)};

  // Test-card texture: default-heap resource + one-shot upload.
  D3D12_RESOURCE_DESC td{};
  td.Dimension = D3D12_RESOURCE_DIMENSION_TEXTURE2D;
  td.Width = kTexSize;
  td.Height = kTexSize;
  td.DepthOrArraySize = 1;
  td.MipLevels = 1;
  td.Format = DXGI_FORMAT_R8G8B8A8_UNORM;
  td.SampleDesc.Count = 1;
  D3D12_HEAP_PROPERTIES def{D3D12_HEAP_TYPE_DEFAULT};
  if (FAILED(dev->CreateCommittedResource(&def, D3D12_HEAP_FLAG_NONE, &td,
                                          D3D12_RESOURCE_STATE_COPY_DEST,
                                          nullptr, IID_PPV_ARGS(&texture_))))
    return false;

  D3D12_PLACED_SUBRESOURCE_FOOTPRINT fp{};
  UINT64 uploadSize = 0;
  dev->GetCopyableFootprints(&td, 0, 1, 0, &fp, nullptr, nullptr, &uploadSize);
  ComPtr<ID3D12Resource> staging;
  bd.Width = uploadSize;
  if (FAILED(dev->CreateCommittedResource(&up, D3D12_HEAP_FLAG_NONE, &bd,
                                          D3D12_RESOURCE_STATE_GENERIC_READ,
                                          nullptr, IID_PPV_ARGS(&staging))))
    return false;
  std::vector<uint32_t> card;
  makeTestCard(&card);
  uint8_t* dst = nullptr;
  staging->Map(0, nullptr, (void**)&dst);
  for (uint32_t y = 0; y < kTexSize; y++)
    memcpy(dst + fp.Offset + y * fp.Footprint.RowPitch, &card[y * kTexSize],
           kTexSize * 4);
  staging->Unmap(0, nullptr);

  // One-shot copy queue.
  ComPtr<ID3D12CommandQueue> q;
  ComPtr<ID3D12CommandAllocator> alloc;
  ComPtr<ID3D12GraphicsCommandList> list;
  ComPtr<ID3D12Fence> fence;
  D3D12_COMMAND_QUEUE_DESC qd{D3D12_COMMAND_LIST_TYPE_DIRECT};
  if (FAILED(dev->CreateCommandQueue(&qd, IID_PPV_ARGS(&q)))) return false;
  dev->CreateCommandAllocator(D3D12_COMMAND_LIST_TYPE_DIRECT,
                              IID_PPV_ARGS(&alloc));
  dev->CreateCommandList(0, D3D12_COMMAND_LIST_TYPE_DIRECT, alloc.Get(),
                         nullptr, IID_PPV_ARGS(&list));
  D3D12_TEXTURE_COPY_LOCATION src{staging.Get(),
                                  D3D12_TEXTURE_COPY_TYPE_PLACED_FOOTPRINT};
  src.PlacedFootprint = fp;
  D3D12_TEXTURE_COPY_LOCATION dstLoc{texture_.Get(),
                                     D3D12_TEXTURE_COPY_TYPE_SUBRESOURCE_INDEX};
  list->CopyTextureRegion(&dstLoc, 0, 0, 0, &src, nullptr);
  D3D12_RESOURCE_BARRIER bar{};
  bar.Transition.pResource = texture_.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);
  list->Close();
  ID3D12CommandList* lists[] = {list.Get()};
  q->ExecuteCommandLists(1, lists);
  dev->CreateFence(0, D3D12_FENCE_FLAG_NONE, IID_PPV_ARGS(&fence));
  q->Signal(fence.Get(), 1);
  HANDLE ev = CreateEventW(nullptr, FALSE, FALSE, nullptr);
  fence->SetEventOnCompletion(1, ev);
  WaitForSingleObject(ev, 5000);
  CloseHandle(ev);

  // SRV heap for the scene texture.
  D3D12_DESCRIPTOR_HEAP_DESC hd{};
  hd.Type = D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV;
  hd.NumDescriptors = 1;
  hd.Flags = D3D12_DESCRIPTOR_HEAP_FLAG_SHADER_VISIBLE;
  if (FAILED(dev->CreateDescriptorHeap(&hd, IID_PPV_ARGS(&srvHeap_))))
    return false;
  D3D12_SHADER_RESOURCE_VIEW_DESC sv{};
  sv.Format = td.Format;
  sv.ViewDimension = D3D12_SRV_DIMENSION_TEXTURE2D;
  sv.Shader4ComponentMapping = D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING;
  sv.Texture2D.MipLevels = 1;
  dev->CreateShaderResourceView(texture_.Get(), &sv,
                                srvHeap_->GetCPUDescriptorHandleForHeapStart());
  return true;
}

namespace {
void mulMvp(const float a[16], const float b[16], float out[16]) {
  float r[16];
  for (int i = 0; i < 4; i++)
    for (int j = 0; j < 4; j++) {
      r[i * 4 + j] = 0;
      for (int k = 0; k < 4; k++) r[i * 4 + j] += a[i * 4 + k] * b[k * 4 + j];
    }
  memcpy(out, r, sizeof(r));
}
}  // namespace

void SceneRenderer::record(ID3D12GraphicsCommandList* list,
                           D3D12_CPU_DESCRIPTOR_HANDLE rtv, uint32_t width,
                           uint32_t height, const float viewProj[16]) {
  record(list, rtv, width, height, viewProj, srvHeap_.Get(),
         srvHeap_->GetGPUDescriptorHandleForHeapStart());
}

void SceneRenderer::record(ID3D12GraphicsCommandList* list,
                           D3D12_CPU_DESCRIPTOR_HANDLE rtv, uint32_t width,
                           uint32_t height, const float viewProj[16],
                           ID3D12DescriptorHeap* srvHeap,
                           D3D12_GPU_DESCRIPTOR_HANDLE srv) {
  float model[16], mvp[16];
  MakeQuadModel(kQuadW, kQuadH, 0, 0, kSceneQuadZ, model);
  mulMvp(viewProj, model, mvp);
  beginPass(list, rtv, width, height);
  drawQuad(list, mvp, srvHeap, srv);
}

void SceneRenderer::beginPass(ID3D12GraphicsCommandList* list,
                              D3D12_CPU_DESCRIPTOR_HANDLE rtv, uint32_t width,
                              uint32_t height) {
  const float clear[4] = {0.05f, 0.05f, 0.08f, 1.0f};
  list->ClearRenderTargetView(rtv, clear, 0, nullptr);
  list->OMSetRenderTargets(1, &rtv, FALSE, nullptr);
  D3D12_VIEWPORT vp{0, 0, (float)width, (float)height, 0.0f, 1.0f};
  D3D12_RECT sc{0, 0, (LONG)width, (LONG)height};
  list->RSSetViewports(1, &vp);
  list->RSSetScissorRects(1, &sc);
  list->SetGraphicsRootSignature(rootSig_.Get());
  list->SetPipelineState(pso_.Get());
  list->IASetPrimitiveTopology(D3D_PRIMITIVE_TOPOLOGY_TRIANGLELIST);
  list->IASetVertexBuffers(0, 1, &vbv_);
}

void SceneRenderer::drawQuad(ID3D12GraphicsCommandList* list,
                             const float mvp[16],
                             ID3D12DescriptorHeap* srvHeap,
                             D3D12_GPU_DESCRIPTOR_HANDLE srv) {
  ID3D12DescriptorHeap* heaps[] = {srvHeap};
  list->SetDescriptorHeaps(1, heaps);
  list->SetGraphicsRoot32BitConstants(0, 16, mvp, 0);
  list->SetGraphicsRootDescriptorTable(1, srv);
  list->DrawInstanced(6, 1, 0, 0);
}

void SceneRenderer::drawQuad(ID3D12GraphicsCommandList* list,
                             const float mvp[16]) {
  drawQuad(list, mvp, srvHeap_.Get(),
           srvHeap_->GetGPUDescriptorHandleForHeapStart());
}

}  // namespace sauna