import bindless;
import ibl_common;
import common_structs;

struct ModelHandles
{
    StructuredBuffer<float4>.Handle vertexBuffer;           // 8
    StructuredBuffer<VertexData>.Handle vertexData;         // 16
    StructuredBuffer<MaterialData>.Handle materialBuffer;   // 24
    StructuredBuffer<TransformData>.Handle nodeBuffer;      // 32
};

[vk::push_constant]
uniform Block pcb;

struct VS_Input
{
    uint vertexId        : SV_VertexID;
};

struct VS_Output
{
    float4 worldPosition    : POSITION;
    float4 worldNormal      : NORMAL;
    float4 vertexColor      : COLOR0;
    float2 uv0              : TEXCOORD0;
    float4 vertexPosition   : SV_Position;
};

float2 GetUV(uint vertexId)
{
    return pcb.vertexData[vertexId].uv0.xy;
}

float4 GetPosition(uint vertexId)
{
    return float4(pcb.vertexBuffer[vertexId].xyz, 1.0f);
}

float4 GetNormal(uint vertexId)
{
    return pcb.vertexData[vertexId].normal;
}

float4 GetColor(uint vertexId)
{
    return pcb.vertexData[vertexId].color0;
}

float4x4 GetNodeTransform(uint nodeIdx)
{
    return pcb.nodeBuffer[nodeIdx].transform;
}

float4x4 GetNormalTransform(uint nodeIdx)
{
    return pcb.nodeBuffer[nodeIdx].normalTransform;
}

[shader("vertex")]
VS_Output vsmain(VS_Input stageInput)
{
    VS_Output stageOutput;

    float4 worldPosition = mul(GetPosition(stageInput.vertexId), GetNodeTransform(pcb.nodeIdx));
    float4 clipSpace = mul(worldPosition, pcb.camera[0].view);

    stageOutput.vertexPosition = mul(clipSpace, pcb.camera[0].proj);
    stageOutput.worldPosition = worldPosition;
    stageOutput.uv0 = GetUV(stageInput.vertexId);
    stageOutput.vertexColor = GetColor(stageInput.vertexId);

    stageOutput.worldNormal = mul(GetNormal(stageInput.vertexId), GetNormalTransform(pcb.nodeIdx));
    return stageOutput;
}

float4 GetAlbedo(float2 uv, float4 color)
{
    float4 albedoFactor = pcb.materialBuffer[NonUniformResourceIndex(pcb.materialIdx)].albedoFactor;
    var albedoTex = pcb.materialBuffer[NonUniformResourceIndex(pcb.materialIdx)].albedoTex;

    return albedoFactor * color * (IsValid(albedoTex) ? albedoTex.Sample(uv) : 1.0f.xxxx);
}

float GetOcclusion(float2 uv)
{
    var occlusionTex = pcb.materialBuffer[NonUniformResourceIndex(pcb.materialIdx)].occlusionTex;

    return IsValid(occlusionTex) ? occlusionTex.Sample(uv).r : 1.0f;
}

float3 GetEmissive(float2 uv)
{
    float3 emissionFactor = pcb.materialBuffer[NonUniformResourceIndex(pcb.materialIdx)].emissionFactor.rgb;
    var emissionTex = pcb.materialBuffer[NonUniformResourceIndex(pcb.materialIdx)].emissionTex;

    return emissionFactor * (IsValid(emissionTex) ? emissionTex.Sample(uv).rgb : 1.0f.xxx);
}

float3 GetNormal(float3 position, float3 normal, float2 uv)
{
    float3 N = normalize(normal);

    var normalTex = pcb.materialBuffer[NonUniformResourceIndex(pcb.materialIdx)].normalTex;
    if (!IsValid(normalTex))
    {
        return N;
    }

    float3 tangentSpaceNormal = normalTex.Sample(uv).xyz * 2.0f - 1.0f;

    float3 q1 = ddx(position);
    float3 q2 = ddy(position);
    float2 st1 = ddx(uv);
    float2 st2 = ddy(uv);

    float3 T = normalize(q1 * st2.y - q2 * st1.y).xyz;
    float3 B = -normalize(cross(N, T));
    float3x3 TBN = float3x3(T, B, N); // Construction is Row by Row.

    return normalize(mul(tangentSpaceNormal, TBN)); // Post multiple to avoid transpose.
}

float2 GetMetalRough(float2 uv)
{
    var material = pcb.materialBuffer[NonUniformResourceIndex(pcb.materialIdx)];
    float2 metalRough = float2(material.metalFactor, material.roughFactor);
    var metalRoughTex = pcb.materialBuffer[NonUniformResourceIndex(pcb.materialIdx)].metalRoughTex;

    return metalRough * (IsValid(metalRoughTex) ? metalRoughTex.Sample(uv).bg : 1.0f.xx); // Metal is B, Rough is G.
}

float3 SampleIrradiance(float3 direction)
{
    var diffuseIrradiance = pcb.lights[0].diffuseIrradiance;
    if (IsValid(diffuseIrradiance))
    {
        return diffuseIrradiance.Sample(direction).rgb;
    }
    return 0.04f.xxx;
}

float3 SamplePrefiltered(float3 direction, float roughness)
{
    var prefilter = pcb.lights[0].prefilter;
    if (!IsValid(prefilter))
        return 0.0f.xxx;

    const float MAX_MIP_LEVEL = 5.0f;
    float mip = MAX_MIP_LEVEL * roughness;
    
    return prefilter.SampleLevel(direction, mip).rgb;
}

float2 SampleBrdfLut(float nDotV, float roughness)
{
    var brdfLut = pcb.lights[0].brdfLUT;
    if (IsValid(brdfLut))
    {
        return brdfLut.Sample(float2(nDotV, roughness));
    }
    return 0.0f.xx;
}

// https://en.wikipedia.org/wiki/Schlick%27s_approximation
float3 FresnelSchlick(float cosine, float3 F_0)
{
    return F_0 + (1.0f - F_0) * pow(clamp(1.0f - cosine, 0.0f, 1.0f), 5.0f); // Clamp to avoid artifacts.
}

// Sebastian Lagarde
float3 FresnelSchlickRoughness(float cosine, float3 F_0, float Roughness)
{
    return F_0 + (max((1.0f - Roughness).xxx, F_0) - F_0) * pow(clamp(1.0f - cosine, 0.0f, 1.0f), 5.0f); // Clamp to avoid artifacts.
}

float3 GetPBRContrib(float3 Albedo, float3 LightColor, float3 ViewDir, float3 Normal, float Metallic, float Roughness, float3 F_0, float3 LightDir, float LightDistance)
{
    float Attenuation = 1.0f / (LightDistance * LightDistance); // TODO: Controlled Attenuation
    float3 Halfway = normalize(ViewDir + LightDir);

    // If the dot is negative, then it we just us 0. 
    // Light from the back is irrelevant.
    float CosineFactor = max(dot(Halfway, ViewDir), 0.0f);
    float NdotV = max(dot(Normal, ViewDir), 0.0f);
    float NdotL = max(dot(Normal, LightDir), 0.0f);
    float NdotH = max(dot(Normal, Halfway), 0.0f);

    float3 Radiance = LightColor * Attenuation;

    float NormalDistribution = TrowbridgeReitzGGX(NdotH, Roughness);
    float Geometry = PunctualGeometrySmith(NdotV, NdotL, Roughness);
    float3 Fresnel = FresnelSchlickRoughness(CosineFactor, F_0, Roughness);

    float3 Numerator = (NormalDistribution * Geometry) * Fresnel;
    float Denominator = 4.0f * NdotV * NdotL;
    float3 Specular = Numerator / (Denominator + 0.00001f);

    float3 K_Specular = Fresnel;
    float3 K_Diffuse = 1.0f.xxx - K_Specular;

    K_Diffuse *= 1.0f - Metallic;

    return NdotL * Radiance * ((K_Diffuse * Albedo / PI) + Specular);
}

float3 GetPointLightInfluence(float3 Albedo, float2 MetalRough, float3 Position, float3 Normal)
{
    var lightData = pcb.lights[0];
    var pointLightBuffer = (StructuredBuffer<PointLight>.Handle)lightData.lights;

    if (!IsValid(pointLightBuffer))
        return 0.0f.xxx;

    uint offset = lightData.pointLightIndexer.offset;
    uint count = lightData.pointLightIndexer.count;

    float3 ViewDir = normalize(pcb.camera[0].position.xyz - Position);

    float Metallic = MetalRough.r;
    float Roughness = MetalRough.g;

    // Dielectric F_0 based on LearnOpenGL.
    // TODO: Cite
    float3 F_0 = 0.04f.xxx;
    F_0 = lerp(F_0, Albedo, Metallic);

    float3 Contrib = 0.0f;
    for (uint i = 0; i < count; ++i)
    {
        PointLight Light = pointLightBuffer[i + offset];

        if (Light.Range < 0.0f)
            continue;

        float3 LightDir = Light.Position - Position;
        float LightDistance = length(LightDir);

        if (LightDistance > Light.Range)
            continue;

        LightDir /= LightDistance; // Normalization

        // Color Unpack
        float R = (Light.Color & 0xFF000000) >> 24;
        float G = (Light.Color & 0x00FF0000) >> 16;
        float B = (Light.Color & 0x0000FF00) >> 8;

        float3 LightColor = Light.Intensity * float3(R, G, B) * 0.00392156862f; // 0.00392156862 = 1/255

        Contrib += GetPBRContrib(Albedo, LightColor, ViewDir, Normal, Metallic, Roughness, F_0, LightDir, LightDistance);
    }

    return Contrib;
}

float3 GetDirectionalLightInfluence(float3 Albedo, float2 MetalRough, float3 Position, float3 Normal)
{
    var lightData = pcb.lights[0];
    var dirLightBuffer = (StructuredBuffer<DirectionalLight>.Handle)lightData.lights;

    if (!IsValid(dirLightBuffer))
        return 0.0f.xxx;

    uint count = lightData.dirLightIndexer.count;

    float3 ViewDir = normalize(pcb.camera[0].position.xyz - Position);

    float Metallic = MetalRough.r;
    float Roughness = MetalRough.g;

    // Dielectric F_0 based on LearnOpenGL.
    // TODO: Cite
    float3 F_0 = 0.04f.xxx;
    F_0 = lerp(F_0, Albedo, Metallic);

    float3 Contrib = 0.0f;
    for (uint i = 0; i < count; ++i)
    {
        DirectionalLight Light = dirLightBuffer[i];

        if (Light.Validity_ < 0.0f)
            continue;

        float3 LightDir = -normalize(float3(Light.Direction));

        // Color Unpack
        float R = (Light.Color & 0xFF000000) >> 24;
        float G = (Light.Color & 0x00FF0000) >> 16;
        float B = (Light.Color & 0x0000FF00) >> 8;

        float3 LightColor = Light.Intensity * float3(R, G, B) * 0.00392156862f; // 0.00392156862 = 1/255

        Contrib += GetPBRContrib(Albedo, LightColor, ViewDir, Normal, Metallic, Roughness, F_0, LightDir, 1.0f);
    }

    return Contrib;
}

float3 GetAmbientInfluence(float3 Albedo, float2 MetalRough, float3 Position, float3 Normal, float Occlusion)
{
    float3 ViewDir = normalize(pcb.camera[0].position.xyz - Position);
    float CosineFactor = max(dot(Normal, ViewDir), 0.0f); // Normal instead of Halfway since there's no halfway in ambient.

    float Metal = MetalRough.r;
    float Roughness = MetalRough.g;

    float3 F_0 = 0.04f.xxx;
    F_0 = lerp(F_0, Albedo, MetalRough.r);
    float3 K_Specular = FresnelSchlickRoughness(CosineFactor, F_0, Roughness);
    float3 K_Diffuse = 1.0f.xxx - K_Specular;

    K_Diffuse *= 1.0f - Metal; // Metals don't have diffuse/refractions.

    float3 ReflectionDir = reflect(-ViewDir, Normal);

    float NdotV = max(dot(Normal, ViewDir), 0.0f);
    float3 PrefilteredColor = SamplePrefiltered(ReflectionDir, Roughness).rgb;
    float2 EnvBRDF = SampleBrdfLut(NdotV, Roughness);
    float3 Specular = PrefilteredColor * (K_Specular * EnvBRDF.x + EnvBRDF.y);

    float3 DiffuseIrradiance = Albedo * SampleIrradiance(Normal);
#if defined(_DEBUG)
    if (pcb.ignoreDiffuse) {
        DiffuseIrradiance = 0.0f.xxx;
    }
    if (pcb.ignoreSpecular) {
        Specular = 0.0f.xxx;
    }
#endif

    return (K_Diffuse * DiffuseIrradiance + Specular) * Occlusion;
}

[shader("fragment")]
float4 fsmain(
    float4 inPosition: POSITION,
    float4 inNormal: NORMAL,
    float4 inColor: COLOR,
    float2 inUV0: TEXCOORD0,
) : SV_Target0
{
    // TODO: This should be invalid on the CPU side.
    if (pcb.materialIdx < 0)
    {
        return float4(1.0f, 0.0f, 1.0f, 1.0f);
    }

    float3 position = inPosition.xyz;
    float3 normal = GetNormal(position, inNormal.xyz, inUV0);

    float4 albedoAlpha = GetAlbedo(inUV0, inColor);
    float3 albedo = albedoAlpha.rgb;
    float alpha = albedoAlpha.a;
    float2 metalRough = GetMetalRough(inUV0);
    float3 emission = GetEmissive(inUV0);
    float occlusion = GetOcclusion(inUV0);

    float3 directionalLightLum = GetDirectionalLightInfluence(albedo, metalRough, position, normal);
    float3 pointLightLum = GetPointLightInfluence(albedo, metalRough, position, normal);
    float3 ambientLum = GetAmbientInfluence(albedo, metalRough, position, normal, occlusion);
    float3 color = directionalLightLum + pointLightLum + ambientLum;

    return float4(Uncharted2Tonemap(color + emission), alpha);
}