// BVLE Voxels - Smooth Surface Nets Pixel Shader (Phase 5.1)
// Per-pixel heightmap blending using the SAME neighbor verification as voxelPS.hlsl.
// Derives a dominant face axis from the smooth normal, then uses identical
// faceU/faceV/stair-priority logic as the blocky pixel shader.

#include "voxelCommon.hlsli"

Texture2DArray<float4> materialTextures : register(t1);
Texture2DArray<float4> normalTextures  : register(t7);
StructuredBuffer<GPUChunkInfo> chunkInfoBuffer : register(t2);
StructuredBuffer<uint> voxelData : register(t3);
SamplerState texSampler : register(s0);

struct PSInput {
    float4 position : SV_POSITION;
    float3 worldPos : WORLDPOS;
    float3 normal   : NORMAL;
    nointerpolation uint primaryMat   : PRIMARYMAT;
    nointerpolation uint chunkIndex   : CHUNKINDEX;
};

static const uint CSIZE = 32;
static const uint CVOL = CSIZE * CSIZE * CSIZE;

// ── Face direction tables (SAME as voxelPS.hlsl) ────────────────
// Face normals: +X, -X, +Y, -Y, +Z, -Z
static const int3 faceNormals[6] = {
    int3( 1, 0, 0), int3(-1, 0, 0),
    int3( 0, 1, 0), int3( 0,-1, 0),
    int3( 0, 0, 1), int3( 0, 0,-1)
};

// Face tangent axes (U, V) — must match voxelPS.hlsl
static const int3 faceUDirs[6] = {
    int3(0, 1, 0), int3(0, 1, 0),
    int3(1, 0, 0), int3(1, 0, 0),
    int3(1, 0, 0), int3(1, 0, 0)
};
static const int3 faceVDirs[6] = {
    int3(0, 0, 1), int3(0, 0, 1),
    int3(0, 0, 1), int3(0, 0, 1),
    int3(0, 1, 0), int3(0, 1, 0)
};

// ── Voxel data read (same as voxelPS.hlsl) ───────────────────────
uint readVoxelMat(int3 coord, uint chunkIdx) {
    GPUChunkInfo info = chunkInfoBuffer[chunkIdx];
    int3 local = coord - (int3)info.worldPos.xyz;
    if (any(local < 0) || any(local >= (int3)CSIZE))
        return 0;
    uint flatIdx = (uint)local.x + (uint)local.y * CSIZE + (uint)local.z * CSIZE * CSIZE;
    uint pairIndex = flatIdx >> 1;
    uint shift = (flatIdx & 1) * 16;
    uint voxel = (voxelData[chunkIdx * (CVOL / 2) + pairIndex] >> shift) & 0xFFFF;
    return voxel >> 8;
}

// ── Stair-priority neighbor lookup (SAME as voxelPS.hlsl) ────────
uint getNeighborMat(int3 voxelCoord, int3 edgeDir, int3 normalDir, uint chunkIdx) {
    // Stair neighbor (priority): block at edge AND offset by normal
    int3 stairPos = voxelCoord + edgeDir + normalDir;
    uint stairMat = readVoxelMat(stairPos, chunkIdx);
    if (stairMat > 0)
        return stairMat;

    // Planar neighbor (fallback): adjacent block in face plane
    int3 planarPos = voxelCoord + edgeDir;
    return readVoxelMat(planarPos, chunkIdx);
}

// ── Triplanar helpers ────────────────────────────────────────────
float3 triplanarWeights(float3 n, float sharpness) {
    float3 w = abs(n);
    w = pow(w, (float3)sharpness);
    return w / (w.x + w.y + w.z + 0.0001);
}

float3 sampleTriplanar(float3 wp, float3 n, uint texIdx, float tiling) {
    float3 w = triplanarWeights(n, 4.0);
    float3 cx = materialTextures.Sample(texSampler, float3(wp.zy * tiling, (float)texIdx)).rgb;
    float3 cy = materialTextures.Sample(texSampler, float3(wp.xz * tiling, (float)texIdx)).rgb;
    float3 cz = materialTextures.Sample(texSampler, float3(wp.xy * tiling, (float)texIdx)).rgb;
    return cx * w.x + cy * w.y + cz * w.z;
}

float4 sampleTriplanarRGBA(float3 wp, float3 n, uint texIdx, float tiling) {
    float3 w = triplanarWeights(n, 4.0);
    float4 cx = materialTextures.Sample(texSampler, float3(wp.zy * tiling, (float)texIdx));
    float4 cy = materialTextures.Sample(texSampler, float3(wp.xz * tiling, (float)texIdx));
    float4 cz = materialTextures.Sample(texSampler, float3(wp.xy * tiling, (float)texIdx));
    return cx * w.x + cy * w.y + cz * w.z;
}

// ── Triplanar normal mapping (UDN blend) ────────────────────────
float3 sampleTriplanarNormal(float3 wp, float3 n, uint texIdx, float tiling) {
    float3 w = triplanarWeights(n, 4.0);
    float3 axisSign = sign(n);
    // Ben Golus UDN reference — swizzled coordinates + sign corrections
    float3 tnX = normalTextures.Sample(texSampler, float3(wp.zy * tiling, (float)texIdx)).rgb * 2.0 - 1.0;
    float3 tnY = normalTextures.Sample(texSampler, float3(wp.xz * tiling, (float)texIdx)).rgb * 2.0 - 1.0;
    float3 tnZ = normalTextures.Sample(texSampler, float3(wp.xy * tiling, (float)texIdx)).rgb * 2.0 - 1.0;
    // OpenGL normal maps: flip green channel ONLY for Y-projection
    tnY.y = -tnY.y;
    // Sign correction for back-facing projections
    tnX.x *= axisSign.x;
    tnY.x *= axisSign.y;
    tnZ.x *= axisSign.z;
    // UDN blend using RAW normal (NOT abs!) — preserves sign for negative faces
    float3 worldNX = float3(tnX.xy + n.zy, n.x).zyx;
    float3 worldNY = float3(tnY.xy + n.xz, n.y).xzy;
    float3 worldNZ = float3(tnZ.xy + n.xy, n.z);
    return normalize(worldNX * w.x + worldNY * w.y + worldNZ * w.z);
}

// ── MRT Output ──────────────────────────────────────────────────
struct PSOutput {
    float4 color  : SV_TARGET0;
    float4 normal : SV_TARGET1;
};

// ── Main PS ──────────────────────────────────────────────────────
[RootSignature(VOXEL_ROOTSIG)]
PSOutput main(PSInput input) {
    PSOutput output;
    float3 N = normalize(input.normal); // smooth normal (for lighting)

    // Geometric normal from screen-space derivatives of worldPos.
    // This is the true triangle face normal — use it for triplanar weights
    // to avoid texture stretching caused by smooth normal interpolation.
    float3 dpx = ddx(input.worldPos);
    float3 dpy = ddy(input.worldPos);
    float3 geoN = normalize(cross(dpx, dpy));
    // Ensure geometric normal faces same hemisphere as smooth normal
    if (dot(geoN, N) < 0.0) geoN = -geoN;

    float tiling = textureTiling;

    // ── Derive dominant face from smooth normal (same tables as blocky PS) ──
    // Find the axis with the largest absolute normal component
    float3 absN = abs(N);
    uint dominantAxis;
    if (absN.x >= absN.y && absN.x >= absN.z)
        dominantAxis = 0; // X
    else if (absN.y >= absN.z)
        dominantAxis = 1; // Y
    else
        dominantAxis = 2; // Z

    // Map to face index: axis*2 + (negative ? 1 : 0)
    uint face = dominantAxis * 2;
    if (N[dominantAxis] < 0.0) face += 1;

    int3 normalDir = faceNormals[face];
    int3 uDir = faceUDirs[face];
    int3 vDir = faceVDirs[face];

    // ── Compute voxel coordinate (SAME as blocky PS) ──
    // Tiny offset inward along dominant normal to handle integer boundaries
    float3 samplePos = input.worldPos - (float3)normalDir * 0.001;
    int3 voxelCoord = (int3)floor(samplePos);

    // Read actual material at this voxel position
    uint selfMat = readVoxelMat(voxelCoord, input.chunkIndex);
    if (selfMat == 0u) selfMat = input.primaryMat; // air fallback

    // ── Face-aligned fractional position (SAME as blocky PS) ──
    float faceFracU = frac(dot(input.worldPos, (float3)uDir));
    float faceFracV = frac(dot(input.worldPos, (float3)vDir));

    // Distance from nearest edge (0 = at edge, 0.5 = at center)
    float uDist = 0.5 - abs(faceFracU - 0.5);
    float vDist = 0.5 - abs(faceFracV - 0.5);

    // Nearest edge direction
    int uSign = (faceFracU >= 0.5) ? 1 : -1;
    int vSign = (faceFracV >= 0.5) ? 1 : -1;
    int3 uEdgeDir = uDir * uSign;
    int3 vEdgeDir = vDir * vSign;

    // ── Stair-priority neighbor lookup (SAME as blocky PS) ──
    uint uNeighborMat = getNeighborMat(voxelCoord, uEdgeDir, normalDir, input.chunkIndex);
    uint vNeighborMat = getNeighborMat(voxelCoord, vEdgeDir, normalDir, input.chunkIndex);

    // ── Blend weights (SAME params as blocky PS) ──
    float blendZone = 0.40;
    float uEdge = abs(faceFracU - 0.5) * 2.0;
    float vEdge = abs(faceFracV - 0.5) * 2.0;

    // Corner attenuation — subtractive (same as blocky PS)
    float blendStart = 1.0 - blendZone * 2.0;
    float uAdj = uEdge - saturate(vEdge - 0.80);
    float vAdj = vEdge - saturate(uEdge - 0.80);
    float uWeight = saturate((uAdj - blendStart) / (1.0 - blendStart)) * 0.5;
    float vWeight = saturate((vAdj - blendStart) / (1.0 - blendStart)) * 0.5;

    // Blend conditions (same as blocky PS, with bleed mask checks)
    bool mainResists = (resistBleedMask >> selfMat) & 1u;
    bool uNeighCanBleed = (bleedMask >> uNeighborMat) & 1u;
    bool vNeighCanBleed = (bleedMask >> vNeighborMat) & 1u;
    bool uNeighResists = (resistBleedMask >> uNeighborMat) & 1u;
    bool vNeighResists = (resistBleedMask >> vNeighborMat) & 1u;
    bool uBlend = (uNeighborMat > 0u && uNeighborMat != selfMat && uWeight > 0.001
                   && !mainResists && uNeighCanBleed);
    bool vBlend = (vNeighborMat > 0u && vNeighborMat != selfMat && vWeight > 0.001
                   && !mainResists && vNeighCanBleed);

    // ── Texturing ──
    uint selfTexIdx = clamp(selfMat - 1u, 0u, 5u);
    float3 albedo;

    if (uBlend || vBlend) {
        float4 mainTex = sampleTriplanarRGBA(input.worldPos, geoN, selfTexIdx, tiling);
        float3 result = mainTex.rgb;
        float sharpness = 16.0;

        if (uBlend) {
            uint uTexIdx = clamp(uNeighborMat - 1u, 0u, 5u);
            float4 uTex = sampleTriplanarRGBA(input.worldPos, geoN, uTexIdx, tiling);
            float bias;
            if (uNeighResists) {
                bias = 0.5 - uWeight * 1.6;
            } else {
                bias = 0.5 - uWeight;
            }
            float mainScore  = mainTex.a + bias;
            float neighScore = uTex.a   - bias;
            float blend = saturate((neighScore - mainScore) * sharpness + 0.5);
            result = lerp(result, uTex.rgb, blend);
        }

        if (vBlend) {
            uint vTexIdx = clamp(vNeighborMat - 1u, 0u, 5u);
            float4 vTex = sampleTriplanarRGBA(input.worldPos, geoN, vTexIdx, tiling);
            float bias;
            if (vNeighResists) {
                bias = 0.5 - vWeight * 1.6;
            } else {
                bias = 0.5 - vWeight;
            }
            float mainScore  = mainTex.a + bias;
            float neighScore = vTex.a   - bias;
            float blend = saturate((neighScore - mainScore) * sharpness + 0.5);
            result = lerp(result, vTex.rgb, blend);
        }

        albedo = result;
    } else {
        albedo = sampleTriplanar(input.worldPos, geoN, selfTexIdx, tiling);
    }

    // ── Normal map perturbation ──
    float3 flatN = N; // preserve for ambient
    float nmStrength = toneMapParams.z;
    if (nmStrength > 0.0) {
        float3 perturbedN = sampleTriplanarNormal(input.worldPos, geoN, selfTexIdx, tiling);
        N = normalize(lerp(N, perturbedN, nmStrength * 0.7)); // lighter on smooth
    }

    // ── Debug lighting modes (F9 cycle) ──
    uint dbgLight = (uint)toneMapParams.w;
    if (dbgLight == 2) {
        // FLAT: uniform gray, no texture, no normal map — pure lighting with geometric normal
        float flatNdotL = max(dot(flatN, normalize(-sunDirection.xyz)), 0.0);
        float flatHemi = flatN.y * 0.5 + 0.5;
        float3 flatAmb = lerp(groundAmbient.rgb, skyAmbient.rgb, flatHemi);
        float3 flatColor = float3(0.5, 0.5, 0.5) * (flatAmb + sunColor.rgb * flatNdotL);
        output.color = float4(flatColor, 1.0);
        output.normal = float4(flatN, 0.0);
        return output;
    }
    if (dbgLight == 3) {
        // ALBEDO only: texture + blend, no lighting
        output.color = float4(albedo, 1.0);
        output.normal = float4(flatN, 0.0);
        return output;
    }
    if (dbgLight == 4) {
        // NdotL only: grayscale NdotL with geometric normal (no normal map)
        float flatNdotL = max(dot(flatN, normalize(-sunDirection.xyz)), 0.0);
        output.color = float4(flatNdotL, flatNdotL, flatNdotL, 1.0);
        output.normal = float4(flatN, 0.0);
        return output;
    }
    if (dbgLight == 5) {
        // NORMAL viz: geometric normal mapped to RGB (XYZ → [0,1])
        output.color = float4(flatN * 0.5 + 0.5, 1.0);
        output.normal = float4(flatN, 0.0);
        return output;
    }

    // Lighting: flat normal for ambient (consistent), perturbed for NdotL (detail)
    float3 L = normalize(-sunDirection.xyz);
    float NdotL = max(dot(N, L), 0.0);
    float hemiLerp = flatN.y * 0.5 + 0.5;
    float3 ambient = lerp(groundAmbient.rgb, skyAmbient.rgb, hemiLerp);
    float3 color = albedo * (sunColor.rgb * NdotL + ambient);

    // ── Rim light ──
    float3 V = normalize(cameraPosition.xyz - input.worldPos);
    float NdotV = saturate(dot(N, V));
    float rim = pow(1.0 - NdotV, rimParams.x) * rimParams.y;
    color += rimColor.rgb * rim;

    // Distance fog
    float dist = length(input.worldPos - cameraPosition.xyz);
    float fogDensity = fogParams.x;
    float fog = 1.0 - exp(-dist * fogDensity);
    color = lerp(color, fogColor.rgb, saturate(fog));

    output.color = float4(color, 1.0);
    output.normal = float4(N, 0.0);
    return output;
}