#include "render/warp_pass.h"

#include <d3dcompiler.h>

#include <cstdio>
#include <cstring>

using Microsoft::WRL::ComPtr;

namespace sauna {
namespace {

// Direct port of EvalDistortion (S3 poly+grow). uv is viewport-local
// (0..1 across one eye's half, v=0 top). flags.x: 1 = point sampling
// (gate A), flags.y: 1 = apply color mult.
const char kWarp[] = R"(
cbuffer C : register(b0) {
  float4 centerScale;  // cx, cy, scale, r2 cutoff
  float4 kr;           // k1,k2,k3 (red), unused
  float4 kg;
  float4 kb;
  float4 colorMult;    // rgb mult, unused
  float4 flags;        // pointSample, useColorMult, r2Clamp
};
Texture2D eyeTex : register(t0);
SamplerState linearBorder : register(s0);
SamplerState pointBorder : register(s1);

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

VSOut vsmain(uint id : SV_VertexID) {
  VSOut o;
  float2 uv = float2((id << 1) & 2, id & 2);
  o.pos = float4(uv * float2(2, -2) + float2(-1, 1), 0, 1);
  o.uv = uv;
  return o;
}

float2 srcUv(float2 t, float r2, float3 k, float2 c, float s) {
  float d = 1.0 + r2 * (k.x + r2 * (k.y + r2 * k.z));
  return 0.5 + (t * d + c) * s;
}

float sampleCh(float2 uv, int ch) {
  float4 v = (flags.x > 0.5) ? eyeTex.SampleLevel(pointBorder, uv, 0)
                             : eyeTex.SampleLevel(linearBorder, uv, 0);
  return ch == 0 ? v.r : (ch == 1 ? v.g : v.b);
}

float4 psmain(VSOut i) : SV_Target {
  float2 c = centerScale.xy;
  float s = centerScale.z;
  float2 t = 2.0 * i.uv - 1.0 - c;
  float r2 = dot(t, t);   // no cutoff clamp by default (S3)
  if (flags.z > 0.5) r2 = min(r2, centerScale.w);
  float3 outc;
  outc.r = sampleCh(srcUv(t, r2, kr.xyz, c, s), 0);
  outc.g = sampleCh(srcUv(t, r2, kg.xyz, c, s), 1);
  outc.b = sampleCh(srcUv(t, r2, kb.xyz, c, s), 2);
  if (flags.y > 0.5) outc *= colorMult.rgb;
  return float4(outc, 1.0);
}
)";

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

}  // namespace

bool WarpPass::init(ID3D12Device* dev, DXGI_FORMAT rtFormat,
                    const HmdConfig& cfg) {
  for (int e = 0; e < 2; e++) {
    const EyeCalib& eye = cfg.eye[e];
    float* c = consts_[e];
    c[0] = (float)eye.rgb[1].center_x;  // shared per eye on Beyond
    c[1] = (float)eye.rgb[1].center_y;
    c[2] = (float)(0.5 / (1.0 + eye.grow_for_undistort));
    c[3] = (float)eye.undistort_r2_cutoff;
    for (int ch = 0; ch < 3; ch++) {
      c[4 + ch * 4 + 0] = (float)eye.rgb[ch].k[0];
      c[4 + ch * 4 + 1] = (float)eye.rgb[ch].k[1];
      c[4 + ch * 4 + 2] = (float)eye.rgb[ch].k[2];
      c[4 + ch * 4 + 3] = 0;
    }
    for (int i = 0; i < 3; i++) c[16 + i] = (float)eye.color_mult[i];
    c[19] = 0;
  }

  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 = 24;  // 5 float4 consts + flags
  params[0].ShaderVisibility = D3D12_SHADER_VISIBILITY_PIXEL;
  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;
  D3D12_STATIC_SAMPLER_DESC samps[2]{};
  for (int i = 0; i < 2; i++) {
    samps[i].Filter = i == 0 ? D3D12_FILTER_MIN_MAG_MIP_LINEAR
                             : D3D12_FILTER_MIN_MAG_MIP_POINT;
    samps[i].AddressU = samps[i].AddressV = samps[i].AddressW =
        D3D12_TEXTURE_ADDRESS_MODE_BORDER;
    samps[i].BorderColor = D3D12_STATIC_BORDER_COLOR_OPAQUE_BLACK;
    samps[i].ShaderRegister = i;
    samps[i].ShaderVisibility = D3D12_SHADER_VISIBILITY_PIXEL;
  }
  D3D12_ROOT_SIGNATURE_DESC rs{};
  rs.NumParameters = 2;
  rs.pParameters = params;
  rs.NumStaticSamplers = 2;
  rs.pStaticSamplers = samps;
  ComPtr<ID3DBlob> sig, err;
  if (FAILED(D3D12SerializeRootSignature(&rs, D3D_ROOT_SIGNATURE_VERSION_1,
                                         &sig, &err))) {
    fprintf(stderr, "warp root sig: %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_GRAPHICS_PIPELINE_STATE_DESC pd{};
  pd.pRootSignature = rootSig_.Get();
  pd.VS = {vs->GetBufferPointer(), vs->GetBufferSize()};
  pd.PS = {ps->GetBufferPointer(), ps->GetBufferSize()};
  pd.BlendState.RenderTarget[0].RenderTargetWriteMask =
      D3D12_COLOR_WRITE_ENABLE_ALL;
  pd.SampleMask = UINT_MAX;
  pd.RasterizerState.FillMode = D3D12_FILL_MODE_SOLID;
  pd.RasterizerState.CullMode = D3D12_CULL_MODE_NONE;
  pd.PrimitiveTopologyType = D3D12_PRIMITIVE_TOPOLOGY_TYPE_TRIANGLE;
  pd.NumRenderTargets = 1;
  pd.RTVFormats[0] = rtFormat;
  pd.SampleDesc.Count = 1;
  return SUCCEEDED(
      dev->CreateGraphicsPipelineState(&pd, IID_PPV_ARGS(&pso_)));
}

void WarpPass::record(ID3D12GraphicsCommandList* list, int eye,
                      D3D12_GPU_DESCRIPTOR_HANDLE eyeSrv, uint32_t panelW,
                      uint32_t panelH, bool pointSample, bool colorMult,
                      bool r2Clamp) {
  const float half = panelW / 2.0f;
  D3D12_VIEWPORT vp{eye * half, 0, half, (float)panelH, 0, 1};
  D3D12_RECT sc{(LONG)(eye * half), 0, (LONG)((eye + 1) * half),
                (LONG)panelH};
  list->RSSetViewports(1, &vp);
  list->RSSetScissorRects(1, &sc);
  list->SetGraphicsRootSignature(rootSig_.Get());
  list->SetPipelineState(pso_.Get());
  float c[24];
  memcpy(c, consts_[eye], sizeof(consts_[0]));
  c[20] = pointSample ? 1.0f : 0.0f;
  c[21] = colorMult ? 1.0f : 0.0f;
  c[22] = r2Clamp ? 1.0f : 0.0f;
  c[23] = 0;
  list->SetGraphicsRoot32BitConstants(0, 24, c, 0);
  list->SetGraphicsRootDescriptorTable(1, eyeSrv);
  list->IASetPrimitiveTopology(D3D_PRIMITIVE_TOPOLOGY_TRIANGLELIST);
  list->DrawInstanced(3, 1, 0, 0);
}

}  // namespace sauna