import Bindless;
import Material;
import PBR;

struct VertexOut {
    float4 worldPosition    : POSITION;
    float3 normal           : NORMAL;
    float4 tangent          : TANGENT;
    float2 texCoord0        : TEXCOORD0;
    float2 texCoord1        : TEXCOORD1;
    float4 vertexColor0     : COLOR0;
    float4 outPosition      : SV_Position;
};

struct CameraData {
    float4x4 view;
    float4x4 proj;
    float4 position;
};

struct DirectionalLight {
    float3 direction;
    float  _padding0;
    float3 color;
    float  _padding1;
};

struct LightData {
    PointLight* pointLights;
    DirectionalLight* dirLights;
    uint pointLightCount;
    uint dirLightCount;

    PointLight getPointLight(uint idx) {
        if (idx >= pointLightCount) return pointLights[0]; 
        return pointLights[idx];
    }

    DirectionalLight getDirectionalLight(uint idx) {
        if (idx >= dirLightCount) return dirLights[0]; 
        return dirLights[idx];
    }
};

struct PerFrameData {
    CameraData camera;
    LightData lightData;
};

uniform ParameterBlock<PerFrameData> pfd;

struct PerInstanceData {
    float4x4 transform;
    float4x4 invTransform;
    Material material;
}

[[vk::push_constant]]
uniform ConstantBuffer<PerInstanceData> pcb;

[shader("vertex")]
VertexOut VertexMain(
    uint vertexId: SV_VertexID,
    float3 position,
    float3 normal,
    float4 tangent,
    float2 texCoord0,
    float2 texCoord1,
    float4 vertexColor0,
) {
    float4 worldPosition = mul(pcb.transform, float4(position, 1.0f));

    VertexOut output;
    output.outPosition = mul(pfd.camera.proj, mul(pfd.camera.view, worldPosition));
    output.worldPosition = worldPosition;
    output.normal = normalize(mul(float4(normal.xyz, 0.0f), pcb.invTransform).xyz);
    if (tangent.w == 0.0f) {
        output.tangent = 0.0f.xxxx;
    } else {
        output.tangent = float4(normalize(mul(float4(tangent.xyz, 0.0f), pcb.invTransform).xyz), tangent.w);
    }
    output.texCoord0 = texCoord0;
    output.texCoord1 = texCoord1;
    output.vertexColor0 = vertexColor0;
    return output;
}


float TrowbridgeReitzGGX(float3 Normal, float3 Halfway, float Roughness)
{
    float Coeff = Roughness * Roughness;
    float Alpha2 = Coeff * Coeff;
    float NdotH = max(dot(Normal, Halfway), 0.0f);
    float NdotH2 = NdotH * NdotH;
	
    float Numerator = Alpha2;
    float Denominator = NdotH2 * (Alpha2 - 1.0f) + 1.0f;
    Denominator = PI * Denominator * Denominator;
	
    return Numerator / Denominator;
}

float GeometrySchlickGGX(float NdotV, float Roughness)
{
    float R = Roughness + 1.0f;
    float K = (R * R) / 8.0f;

    float Numerator = NdotV;
    float Denominator = NdotV * (1.0f - K) + K;
	
    return Numerator / Denominator;
}

float GeometrySmith(float NdotV, float NdotL, float Roughness)
{
    float GGX1 = GeometrySchlickGGX(NdotV, Roughness);
    float GGX2 = GeometrySchlickGGX(NdotL, Roughness);
	
    return GGX1 * GGX2;
}

// 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);
        
    float CosineFactor = max(dot(Halfway, ViewDir), 0.0f);
    float NdotV = max(dot(Normal, ViewDir), 0.0f);
    float NdotL = max(dot(Normal, LightDir), 0.0f);
        
    float3 Radiance = LightColor * Attenuation;
    
    float NormalDistribution = TrowbridgeReitzGGX(Normal, Halfway, Roughness);
    float Geometry = GeometrySmith(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)
{
    if (pfd.lightData.pointLightCount == 0)
        return 0.0f.xxx;
    
    float3 ViewDir = normalize(pfd.camera.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 < pfd.lightData.pointLightCount; ++i)
    {
        PointLight Light = pfd.lightData.pointLights[i];
        
        if (Light.range < 0.0f)
            continue;
        
        float3 LightDir = float3(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, Light.color, ViewDir, Normal, Metallic, Roughness, F_0, LightDir, LightDistance);
    }
    
    return Contrib;
}

[shader("fragment")]
float4 FragmentMain(
    float4 position         : POSITION,
    float3 normal           : NORMAL,
    float4 tangent          : TANGENT,
    float2 texCoord0        : TEXCOORD0,
    float2 texCoord1        : TEXCOORD1,
    float4 vertexColor0     : COLOR0,
) : SV_Target0 {

    float3 N = pcb.material.getNormal(position.xyz, normal.xyz, tangent, texCoord0);
    float2 metalRough = pcb.material.getMetalRough(texCoord0);

    let albedo = pcb.material.getAlbedo(texCoord0, vertexColor0);
    let viewDir = normalize(position.xyz - pfd.camera.position.xyz);

    //float3 f_0 = 0.04f.xxx;
    //f_0 = lerp(f_0, albedo.rgb, metalRough.x);


    //float3 contrib = 0.0f.xxx;
    //for (uint i = 0; i < pfd.lightData.pointLightCount; ++i) {
    //    PointLight pointlight = pfd.lightData.pointLights[i];

    //    contrib += pointlight.getInfluence(albedo.rgb, metalRough, viewDir, position.xyz, N, f_0);
    //}

    let contrib = GetPointLightInfluence(albedo.rgb, metalRough, position.xyz, N);

    return float4(pcb.material.getEmissive(texCoord0) + contrib, 1.0f);
}