import Bindless;

struct VertexOut {
    float4 outPosition      : SV_Position;
    float4 worldPosition    : WorldPosition;
    float4 normal           : WorldNormal;
    float2 texCoord0        : TexCoord0;
    float2 texCoord1        : TexCoord1;
    float4 vertexColor0     : VertexColor;
};

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

struct PointLight {
    float3 position;
    float  range;
    float3 color;
    float  attenuation;
};

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;
    Sampler2D.RID textureID;
    uint _padding;
    float metallic;
    float roughness;
    float4 baseColor;
}

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

[shader("vertex")]
VertexOut VertexMain(
    uint vertexId: SV_VertexID,
    float3 position,
    float3 normal,
    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 = mul(pcb.transform, float4(normalize(normal.rgb), 0.0f));
    output.texCoord0 = texCoord0;
    output.texCoord1 = texCoord1;
    output.vertexColor0 = vertexColor0;
    return output;
}

[shader("fragment")]
float4 FragmentMain(
    float4 worldPosition : WorldPosition,
    float4 normal       : WorldNormal,
    float2 uv0          : TexCoord0,
    float2 uv1          : TexCoord1,
    float4 color        : VertexColor,
) : SV_Target0 {

    float3 diffuse = 0.0f.xxx;
    float3 specular = 0.0f.xxx;

    for (uint i = 0; i < pfd.lightData.pointLightCount; ++i) {
        PointLight pointlight = pfd.lightData.pointLights[i];

        let lightPosition = pointlight.position;
        let lightDisplace = worldPosition.xyz - lightPosition;
        let lightDistance = length(lightDisplace);
        let lightDirection = normalize(lightDisplace);
        let viewDirection = normalize(worldPosition.xyz - pfd.camera.position.xyz);
        let halfWayVector = normalize(-lightDirection + viewDirection);

        let attenuation = (1.0f / lightDistance);

        let diffuseFactor = pcb.roughness * dot(-lightDirection, normalize(normal.xyz));
        diffuse += pointlight.color * diffuseFactor;

        let specularFactor = (1.0f - pcb.roughness) * pow(max(dot(halfWayVector, viewDirection), 0.0f), 32.0f) * attenuation;

        specular += pointlight.color * specularFactor;
    }

    if (let texture = pcb.textureID) {
        return float4(texture.Sample(uv0).rgb, 1.0f) * pcb.baseColor * color * float4((diffuse + specular), 0.0f);
    } else {
        return pcb.baseColor * color * float4((diffuse + specular), 0.0f);
    }
}