import Common;

struct HitInfo
{
  float3 position;
  float3 normal;
  float3 barycentricCoords;
  uint instanceIndex;
  uint primitiveIndex;
};

HitInfo get_hit_info(RayQuery<RAY_FLAG_NONE> q)
{
  HitInfo info;
  // In Slang for Metal/Vulkan, these are the standard names for ray query results
  info.instanceIndex = q.CommittedInstanceID();
  info.primitiveIndex = q.CommittedPrimitiveIndex();
  float2 baryCenter = q.CommittedRayBarycentrics();
  info.barycentricCoords = float3(1.0f - baryCenter.x - baryCenter.y, baryCenter.x, baryCenter.y);
  return info;
}

Vertex interpolate_vertex(uint vertexIdx0, uint vertexIdx1, uint vertexIdx2, float3 bary)
{
  Vertex v0 = loadVertex(vertexIdx0);
  Vertex v1 = loadVertex(vertexIdx1);
  Vertex v2 = loadVertex(vertexIdx2);

  Vertex vert;
  vert.position = v0.position * bary.x + v1.position * bary.y + v2.position * bary.z;
  vert.texCoords = v0.texCoords * bary.x + v1.texCoords * bary.y + v2.texCoords * bary.z;
  vert.normal = v0.normal * bary.x + v1.normal * bary.y + v2.normal * bary.z;
  return vert;
}

[shader("compute")]
[numthreads(8, 8, 1)]
void computeKernel(uint2 threadId: SV_DispatchThreadID)
{
  if (threadId.x >= pParams.cam.width || threadId.y >= pParams.cam.height)
    return;

  uint pass = pSamps.pass;
  uint samplesPerPixel = pSamps.samplesPerPixel;
  if (pass == samplesPerPixel)
    return;

  uint2 pix = threadId;
  uint imgWidth = pParams.cam.width;
  uint imgHeight = pParams.cam.height;

  // -- Camera setup --
  float3 camPos = pParams.cam.cameraPosition;
  float3 camForward = pParams.cam.cameraForward;
  float f = pParams.cam.f;
  float S_O = pParams.cam.S_O;
  float3 fogEmm = pParams.cam.fogEmm;
  float ks = pParams.cam.ks;
  float A = pParams.cam.A;
  float ka = pParams.cam.ka;
  float2 sensorSize = pParams.cam.sensorSize;

  float3 cx = -normalize(cross(camForward, abs(camForward.y) < 0.9 ? float3(0, 1, 0) : float3(0, 0, 1)));
  float3 cy = cross(camForward, cx);
  const float2 sdim = sensorSize;

  float S_I = (S_O * f) / (S_O - f);

  // -- Sample sensor --
  float3 rnd = rand01(uint3(pix, pass));
  float2 rnd2 = 2.0f * float2(rnd.xy); // tent filter
  float2 tent = float2(rnd2.x < 1 ? sqrt(rnd2.x) - 1 : 1 - sqrt(2 - rnd2.x), rnd2.y < 1 ? sqrt(rnd2.y) - 1 : 1 - sqrt(2 - rnd2.y));
  float2 s = ((float2(pix) + 0.5f * (0.5f + float2((pass / 2) % 2, pass % 2) + tent)) / float2(imgWidth, imgHeight) - 0.5f) * sdim;

  float3 lc = camPos + camForward * 0.035f; // sample on 3d sensor plane
  float3 spos = camPos + cx * s.x + cy * s.y;
  float3 rayDir = normalize(lc - spos);

  // -- Lens (Aperture) --
  float3 lensN = -camForward;
  float3 lensX = cross(lensN, float3(0, 1, 0));
  float3 lensY = cross(lensN, lensX);
  float2 rndL = rand01(uint3(pix, pass + 100)).xy;
  float3 lensSample = lc + (rndL.x - 0.5) * A * lensX + (rndL.y - 0.5) * A * lensY;

  float3 focalPoint = camPos + (S_O + S_I) * camForward;

  // Simple ray construction
  float3 rayOrg = lensSample;
  float3 rayDirFinal = normalize(focalPoint - lensSample);

  // -- Path Tracing Loop --
  float3 accumulatedRadiance = float3(0.0);
  float3 throughput = float3(1.0);

  for (int bounce = 0; bounce < 4; ++bounce)
  {
    RayQuery<RAY_FLAG_NONE> q;
    RayDesc rayDesc;
    rayDesc.Origin = rayOrg;
    rayDesc.Direction = rayDirFinal;
    rayDesc.TMin = 0.001;
    rayDesc.TMax = 1e20;
    q.TraceRayInline(pParams.scene, RAY_FLAG_NONE, 0xff, rayDesc);

    if (q.Proceed())
    {
      HitInfo hit = get_hit_info(q);

      ModelReference m = pParams.modelData[hit.instanceIndex];
      uint indexOffset = m.indicesOffset;
      uint vertexOffset = m.positionOffset;

      uint v0 = vertexOffset + pParams.indexBuffer[indexOffset + 3 * hit.primitiveIndex + 0];
      uint v1 = vertexOffset + pParams.indexBuffer[indexOffset + 3 * hit.primitiveIndex + 1];
      uint v2 = vertexOffset + pParams.indexBuffer[indexOffset + 3 * hit.primitiveIndex + 2];

      Vertex vert = interpolate_vertex(v0, v1, v2, hit.barycentricCoords);
      MaterialParameter mat = pParams.materialData[m.materialIndex];

      accumulatedRadiance += throughput * mat.emissive_type.xyz;

      // --- Direct Lighting (NEE) ---
      float3 directLight = float3(0);
      for (uint i = 0; i < pSamps.numDirectionalLights; ++i)
      {
        float3 lDir = -pParams.directionalLights[i].direction.xyz;
        RayQuery<RAY_FLAG_NONE> sq;
        RayDesc rayDesc;
        rayDesc.Origin = vert.position + vert.normal * 0.001;
        rayDesc.Direction = lDir;
        rayDesc.TMin = 0.001;
        rayDesc.TMax = 1e20;
        sq.TraceRayInline(pParams.scene, RAY_FLAG_NONE, 0xff, rayDesc);
        if (!sq.Proceed())
        {
          directLight += mat.shade(vert.normal, -rayDirFinal, lDir, pParams.directionalLights[i].color);
        }
      }
      for (uint i = 0; i < pSamps.numPointLights; ++i)
      {
        float3 lVec = pParams.pointLights[i].position - vert.position;
        float3 lDir = normalize(lVec);
        RayQuery<RAY_FLAG_NONE> sq;
        RayDesc rayDesc;
        rayDesc.Origin = vert.position + vert.normal * 0.001;
        rayDesc.Direction = lDir;
        rayDesc.TMin = 0.001;
        rayDesc.TMax = 1e20;
        sq.TraceRayInline(pParams.scene, RAY_FLAG_NONE, 0xff, rayDesc);
        if (sq.Proceed() == false || sq.CommittedRayT() > length(lVec))
        {
          directLight += mat.shade(vert.normal, -rayDirFinal, lDir, pParams.pointLights[i].color);
        }
      }
      accumulatedRadiance += throughput * directLight;

      // --- Indirect Lighting (Cosine-weighted sampling) ---
      float3 rnd = rand01(uint3(pix, pass + bounce + 200));
      float r1 = 2.0 * PI * rnd.x;
      float r2 = rnd.y;
      float r2s = sqrt(r2);

      float3 w = vert.normal;
      float3 u = normalize(cross(abs(w.x) > 0.1 ? float3(0, 1, 0) : float3(1, 0, 0), w));
      float3 v = cross(w, u);
      float3 nextDir = normalize(u * cos(r1) * r2s + v * sin(r1) * r2s + w * sqrt(1.0 - r2));

      throughput *= mat.albedo_alpha.xyz;

      rayOrg = vert.position + vert.normal * 0.001;
      rayDirFinal = nextDir;

      if (length(throughput) < 0.01)
        break;
    }
    else
    {
      accumulatedRadiance += throughput * float3(0.05, 0.05, 0.1);
      break;
    }
  }

  pParams.image[threadId] = float4(accumulatedRadiance, 1.0);
}