// BVLE Voxels - Pixel Shader (Triplanar textured with PS-based height blending)
// Phase 3 v2: reads voxel data directly in PS for neighbor material lookups.
// Two independent blend axes (U/V), corner attenuation, winner-takes-all heightmap.

#include "voxelCommon.hlsli"

Texture2DArray materialTextures : register(t1);
SamplerState materialSampler : register(s0);

// Voxel data buffer (same as compute mesher uses) — bound at t3 in GPU mesh path
StructuredBuffer<uint> voxelData : register(t3);
StructuredBuffer<GPUChunkInfo> chunkInfoBuffer : register(t2);

struct PSInput {
    float4 position : SV_POSITION;
    float3 worldPos : WORLDPOS;
    float3 normal   : NORMAL;
    float2 uv       : TEXCOORD0;
    nointerpolation uint materialID : MATERIALID;
    nointerpolation uint faceID : FACEID;
    nointerpolation float debugFlag : DEBUGFLAG;
    nointerpolation uint chunkIndex : CHUNKINDEX;
};

// ── Constants ──────────────────────────────────────────────────────
static const uint CSIZE = 32;
static const uint CVOL = CSIZE * CSIZE * CSIZE;

// 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 voxelVS.hlsl faceU/faceV
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 helpers ────────────────────────────────────────

// Read material ID from voxel data (16-bit voxels packed as uint16 pairs)
// Returns high 8 bits = material ID, 0 = air
uint readVoxelMat(int3 coord, uint chunkIdx) {
    // Compute chunk-local coords and check bounds
    GPUChunkInfo info = chunkInfoBuffer[chunkIdx];
    float3 chunkOrigin = info.worldPos.xyz;

    // coord is in world voxel space — convert to chunk-local
    int3 local = coord - (int3)chunkOrigin;

    // Out of this chunk's bounds → treat as air (no cross-chunk lookup for now)
    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;

    // voxelData is laid out as: all chunks packed sequentially
    // Each chunk is CVOL/2 uints (16384 uints = 32^3 voxels / 2 per uint)
    uint bufferOffset = chunkIdx * (CVOL / 2);
    uint voxel = (voxelData[bufferOffset + pairIndex] >> shift) & 0xFFFF;
    return voxel >> 8; // high 8 bits = material ID
}

// Get neighbor material with stair priority:
// Check pos + edgeDir + normalDir FIRST (the stair block that visually masks the edge),
// then fallback to pos + edgeDir if stair is air.
uint getNeighborMat(int3 voxelCoord, int3 edgeDir, int3 normalDir, uint chunkIdx) {
    // Stair neighbor (priority): the block that sits at the edge AND is offset by the normal
    int3 stairPos = voxelCoord + edgeDir + normalDir;
    uint stairMat = readVoxelMat(stairPos, chunkIdx);
    if (stairMat > 0)
        return stairMat;

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

// ── Noise for transition variation ─────────────────────────────────
float hash31(float3 p) {
    float3 q = frac(p * float3(127.1, 311.7, 74.7));
    q += dot(q, q.yzx + 33.33);
    return frac((q.x + q.y) * q.z);
}

// ── Triplanar helpers ──────────────────────────────────────────────

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

// Triplanar sampling — RGB only (non-blended path)
float3 sampleTriplanar(float3 worldPos, float3 normal, uint texIndex, float tiling) {
    float3 w = triplanarWeights(normal, 4.0);
    float3 colX = materialTextures.Sample(materialSampler, float3(worldPos.yz * tiling, (float)texIndex)).rgb;
    float3 colY = materialTextures.Sample(materialSampler, float3(worldPos.xz * tiling, (float)texIndex)).rgb;
    float3 colZ = materialTextures.Sample(materialSampler, float3(worldPos.xy * tiling, (float)texIndex)).rgb;
    return colX * w.x + colY * w.y + colZ * w.z;
}

// Triplanar sampling — RGBA (includes heightmap in alpha)
float4 sampleTriplanarRGBA(float3 worldPos, float3 normal, uint texIndex, float tiling) {
    float3 w = triplanarWeights(normal, 4.0);
    float4 colX = materialTextures.Sample(materialSampler, float3(worldPos.yz * tiling, (float)texIndex));
    float4 colY = materialTextures.Sample(materialSampler, float3(worldPos.xz * tiling, (float)texIndex));
    float4 colZ = materialTextures.Sample(materialSampler, float3(worldPos.xy * tiling, (float)texIndex));
    return colX * w.x + colY * w.y + colZ * w.z;
}

// ── Debug face colors ──────────────────────────────────────────────
static const float3 faceDebugColors[6] = {
    float3(1.0, 0.2, 0.2),   // 0: +X = RED
    float3(0.5, 0.0, 0.0),   // 1: -X = DARK RED
    float3(0.2, 1.0, 0.2),   // 2: +Y = GREEN
    float3(0.0, 0.5, 0.0),   // 3: -Y = DARK GREEN
    float3(0.2, 0.2, 1.0),   // 4: +Z = BLUE
    float3(0.0, 0.0, 0.5),   // 5: -Z = DARK BLUE
};

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

// ── Main PS ────────────────────────────────────────────────────────

[RootSignature(VOXEL_ROOTSIG)]
PSOutput main(PSInput input)
{
    PSOutput output;

    // ── DEBUG MODE: face direction colors ──
    if (input.debugFlag > 0.5)
    {
        uint fid = min(input.faceID, 5u);
        float3 faceColor = faceDebugColors[fid];
        float2 checker = floor(input.worldPos.xz * 0.5);
        float check = frac((checker.x + checker.y) * 0.5) * 2.0;
        faceColor *= (0.85 + 0.15 * check);
        output.color = float4(faceColor, 1.0);
        output.normal = float4(normalize(input.normal), 0.0);
        return output;
    }

    // ── NORMAL MODE: triplanar textured with height-based blending ──
    float3 N = normalize(input.normal);
    float3 L = normalize(-sunDirection.xyz);
    float NdotL = max(dot(N, L), 0.0);

    uint texIndex = clamp(input.materialID - 1u, 0u, 5u);
    float tiling = textureTiling;
    float3 albedo;

    // ── Height-based blending via PS voxel data lookup ──
    if (blendEnabled > 0.5 && input.materialID > 0u)
    {
        uint face = min(input.faceID, 5u);
        int3 normalDir = faceNormals[face];
        int3 uDir = faceUDirs[face];
        int3 vDir = faceVDirs[face];

        // Compute voxel coordinate from world position
        // Offset inward by normal * 0.001 to handle positive faces at integer boundaries
        float3 samplePos = input.worldPos - (float3)normalDir * 0.001;
        int3 voxelCoord = (int3)floor(samplePos);

        // Fractional position within the voxel face
        // Use worldPos directly (chunk origin is integer-aligned, so frac is same)
        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: which side of the voxel face is this pixel closer to?
        int uSign = (faceFracU >= 0.5) ? 1 : -1;
        int vSign = (faceFracV >= 0.5) ? 1 : -1;

        int3 uEdgeDir = uDir * uSign;
        int3 vEdgeDir = vDir * vSign;

        // Get neighbor materials (with stair priority)
        uint uNeighborMat = getNeighborMat(voxelCoord, uEdgeDir, normalDir, input.chunkIndex);
        uint vNeighborMat = getNeighborMat(voxelCoord, vEdgeDir, normalDir, input.chunkIndex);

        // Blend zone: 0.25 voxels from each edge (covers 50% of face total)
        float blendZone = 0.25;

        // Edge distances normalized to 0..1 (0=center, 1=edge) for corner attenuation
        float uEdge = abs(faceFracU - 0.5) * 2.0; // 0 at center, 1 at edge
        float vEdge = abs(faceFracV - 0.5) * 2.0;

        // Corner attenuation — Subtractive (Unity reference style)
        // When one axis is very close to its edge (>0.80), it subtracts from the other axis
        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;

        // Only blend if neighbor has a different material AND blend flags allow it:
        // - Current material must NOT resist bleed (resistBleedMask)
        // - Neighbor material must be allowed to bleed (bleedMask)
        bool mainResists = (resistBleedMask >> input.materialID) & 1u;
        bool uNeighCanBleed = (bleedMask >> uNeighborMat) & 1u;
        bool vNeighCanBleed = (bleedMask >> vNeighborMat) & 1u;
        bool uBlend = (uNeighborMat > 0u && uNeighborMat != input.materialID && uWeight > 0.001
                       && !mainResists && uNeighCanBleed);
        bool vBlend = (vNeighborMat > 0u && vNeighborMat != input.materialID && vWeight > 0.001
                       && !mainResists && vNeighCanBleed);

        // ── DEBUG BLEND MODE (F4): show blend zones as colors ──
        if (debugBlend > 0.5) {
            float3 debugColor = float3(0.3, 0.3, 0.3); // gray = no blend
            uint selfMat = readVoxelMat(voxelCoord, input.chunkIndex);
            if (selfMat != input.materialID) {
                output.color = float4(1, 0, 0, 1); // RED = data mismatch bug
                output.normal = float4(N, 0.0);
                return output;
            }
            if (uBlend) debugColor.r = uWeight * 2.0;
            if (vBlend) debugColor.b = vWeight * 2.0;
            if (!uBlend && !vBlend) debugColor.g = 0.5;
            output.color = float4(debugColor, 1.0);
            output.normal = float4(N, 0.0);
            return output;
        }

        if (uBlend || vBlend) {
            // Sample main material (RGBA: rgb=color, a=heightmap)
            float4 mainTex = sampleTriplanarRGBA(input.worldPos, N, texIndex, tiling);
            float3 result = mainTex.rgb;

            // Winner-takes-all height blending:
            // Each material's "score" = its heightmap + a proximity bias.
            // Near the edge (weight→0.5), both have equal bias → heightmap decides.
            // Far from the edge (weight→0), main gets a large bias → always wins.
            // The highest score wins 100% — transition is SHARP but its shape is organic.
            // A small sharpness factor softens the very edge to avoid aliasing.
            float sharpness = 16.0; // higher = sharper transition (∞ = binary)

            if (uBlend) {
                uint uTexIdx = clamp(uNeighborMat - 1u, 0u, 5u);
                float4 uTex = sampleTriplanarRGBA(input.worldPos, N, uTexIdx, tiling);

                // Symmetric proximity bias: at edge (weight=0.5) bias=0 → pure heightmap.
                // Away from edge (weight=0) bias=0.5 → main always wins.
                float 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, N, vTexIdx, tiling);

                float 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, N, texIndex, tiling);
        }
    }
    else
    {
        float3 baseColor = N * 0.5 + 0.5;
        float3 texColor = sampleTriplanar(input.worldPos, N, texIndex, tiling);
        albedo = (input.materialID > 0u) ? texColor : baseColor;
    }

    // ── Lighting ──
    float hemiLerp = N.y * 0.5 + 0.5; // 0=down, 1=up
    float3 ambient = lerp(groundAmbient.rgb, skyAmbient.rgb, hemiLerp);
    float3 diffuse = sunColor.rgb * NdotL;
    float3 color = albedo * (ambient + diffuse);

    // ── 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;
}