#define MAX_LIGHT_CNT 64


struct Material
{
    float3 ambient;
    float3 diffuse;
    float3 specular;
    // Sampler2D ambient_map;
    // Sampler2D diffuse_map;
    // Sampler2D specular_map;
    float shininess;
};

struct LightPhong
{
    float3 diffuse;
    float3 specular;
};

struct AttenuationProperties
{
    float constant;
    float linear;
    float quadratic;
};

struct PointLight
{
    LightPhong phong;
    float3 position;
    AttenuationProperties attenuation_props;
};

struct DirectionalLight
{
    LightPhong phong;
    float3 direction;
};

struct CalculatedLight
{
    float3 diffuse;
    float3 specular;
};

struct BlinnPhongLighting
{
    float3 view_pos;
    Material material;

    DirectionalLight directional_lights[MAX_LIGHT_CNT];
    uint directional_light_cnt;

    PointLight point_lights[MAX_LIGHT_CNT];
    uint point_light_cnt;

    float4 evaluate(in VertexData vertex_data)
    {
        float3 ambient_light =
            this.calc_ambient_light(vertex_data.texture_coords);

        float3 directional_light_sum = this.calc_dir_light_sum(vertex_data);
        float3 point_light_sum = this.calc_point_light_sum(vertex_data);

        return float4((ambient_light + directional_light_sum + point_light_sum), 1.0);
    }

    float3 calc_dir_light_sum(in VertexData vertex_data)
    {
        float3 directional_light_sum = float3(0.0, 0.0, 0.0);

        for (uint index = 0; index < this.directional_light_cnt; index++)
        {
            CalculatedLight calculated_dir_light;

            this.calc_light(
                // Negated since we want the light to point from the light direction
                normalize(-this.directional_lights[index].direction),
                this.directional_lights[index].phong,
                vertex_data,
                calculated_dir_light);

            directional_light_sum +=
                calculated_dir_light.diffuse + calculated_dir_light.specular;
        }

        return directional_light_sum;
    }

    float3 calc_point_light_sum(in VertexData vertex_data)
    {
        float3 point_light_sum = float3(0.0, 0.0, 0.0);

        for (uint index = 0; index < this.point_light_cnt; index++)
        {
            float3 light_direction =
                normalize(this.point_lights[index].position - vertex_data.world_space_pos);

            CalculatedLight calculated_point_light;

            this.calc_light(
                light_direction,
                this.point_lights[index].phong,
                vertex_data,
                calculated_point_light);

            float attenuation =
                this.calc_attenuation(this.point_lights[index], vertex_data.world_space_pos);

            calculated_point_light.diffuse *= attenuation;
            calculated_point_light.specular *= attenuation;

            point_light_sum +=
                calculated_point_light.diffuse + calculated_point_light.specular;
        }

        return point_light_sum;
    }

    void calc_light(
        in float3 light_direction,
        in LightPhong light_phong,
        in VertexData vertex_data,
        out CalculatedLight calculated_light)
    {
        float3 norm = normalize(vertex_data.world_space_normal);

        calculated_light.diffuse = this.calc_diffuse_light(
            light_phong,
            light_direction,
            norm,
            vertex_data.texture_coords);

        calculated_light.specular = this.calc_specular_light(
            light_phong,
            light_direction,
            norm,
            vertex_data.world_space_pos,
            vertex_data.texture_coords);
    }

    float3 calc_ambient_light(in float2 texture_coords)
    {
        return ambient_map.Sample(texture_coords).xyz * this.material.ambient;
        // return this.material.ambient_map.Sample(texture_coords).xyz * this.material.ambient;
    }

    float3 calc_diffuse_light(
        in LightPhong light_phong,
        in float3 light_dir,
        in float3 norm,
        in float2 texture_coords)
    {
        float diff = max(dot(norm, light_dir), 0.0);

        return light_phong.diffuse * (diff * (diffuse_map.Sample(texture_coords).xyz * this.material.diffuse));
        // return light_phong.diffuse * (diff * (this.material.diffuse_map.Sample(texture_coords).xyz * this.material.diffuse));
    }

    float3 calc_specular_light(
        in LightPhong light_phong,
        in float3 light_dir,
        in float3 norm,
        in float3 frag_pos,
        in float2 texture_coords)
    {
        float3 view_direction = normalize(this.view_pos - frag_pos);

        float3 halfway_direction = normalize(light_dir + view_direction);

        float spec =
            pow(max(dot(norm, halfway_direction), 0.0), this.material.shininess);

        return light_phong.specular * (spec * (specular_map.Sample(texture_coords).xyz * this.material.specular));
        // return light_phong.specular * (spec * (this.material.specular_map.Sample(texture_coords).xyz * this.material.specular));
    }

    float calc_attenuation(in PointLight point_light, in float3 position)
    {
        float light_distance = length(point_light.position - position);

        return 1.0 / (point_light.attenuation_props.constant + point_light.attenuation_props.linear * light_distance + point_light.attenuation_props.quadratic * pow(light_distance, 2));
    }
};

struct Model3D
{
    float4x4 model;
    float4x4 model_inverted;
    float4x4 view;
    float4x4 projection;
}

// ParameterBlock<BlinnPhongLighting> blinn_phong_lighting;

// ConstantBuffer<BlinnPhongLighting> blinn_phong_lighting;

// ConstantBuffer<Model3D> model_3d;

// ParameterBlock<Model3D> model_3d;

cbuffer Uniforms {
	Model3D model_3d;
	BlinnPhongLighting lighting;
}

Sampler2D ambient_map;
Sampler2D diffuse_map;
Sampler2D specular_map;


struct VertexData
{
    float2 texture_coords;
    float3 world_space_pos;
    float3 world_space_normal;
};

struct VertexStageOutput
{
    VertexData vertex_data : VertexData;
    float4 sv_position : SV_Position;
};

struct Vertex
{
    float3 pos;
    float2 texture_coords;
    float3 normal;
};

struct Fragment
{
    float4 color;
};

[shader("vertex")]
VertexStageOutput vertex_main(
	Vertex vertex,
    // uniform ConstantBuffer<Model3D> model_3d
)
{
    VertexStageOutput stage_output;

	// TODO: Investigate why mul arguments need to be ordered this way.
	//		 The mul arguments are reordered in the GLSL output

	// float4x4 proj_view = mul(model_3d.projection, model_3d.view);
	float4x4 proj_view = mul(model_3d.view, model_3d.projection);

    // float4x4 proj_view_model =
    //     mul(proj_view, model_3d.model);

    float4x4 proj_view_model =
        mul(model_3d.model, proj_view);

    // stage_output.sv_position = mul(proj_view_model, float4(vertex.pos, 1.0));

    stage_output.sv_position = mul(float4(vertex.pos, 1.0), proj_view_model);

    float4 vertex_pos = float4(vertex.pos, 1.0);

    stage_output.vertex_data.world_space_pos =
        float3(mul(model_3d.model, vertex_pos).xyz);

    stage_output.vertex_data.texture_coords = vertex.texture_coords;

    stage_output.vertex_data.world_space_normal =
        mul(float3x3(transpose(model_3d.model_inverted)), vertex.normal);

    return stage_output;
}

[shader("fragment")]
Fragment fragment_main(
    VertexData vertex_data: VertexData,
    // uniform ConstantBuffer<BlinnPhongLighting> lighting
) : SV_Target
{
    Fragment fragment;

    fragment.color = lighting.evaluate(vertex_data);

	// fragment.color = float4(1.0, 1.0, 1.0, 1.0);

    return fragment;
}