#include <metal_stdlib>
using namespace metal;

struct Material {
	float3 albedo;
	float metallic;
	float roughness;
};

struct FragmentInput {
	float2 uv [[user(locn0)]];
	float3 normal [[user(locn1)]];
	float3 world_pos [[user(locn2)]];
};

struct FragmentOutput {
	float4 color [[user(locn0)]];
};

float3 fresnel_schlick(float cos_theta, float3 f0) {
	return (f0 + ((float3(1.0, 1.0, 1.0) - f0) * pow((1.0 - cos_theta), 5.0)));
}

float3 compute_lighting(float3 normal, float3 view_dir, float3 albedo) {
	float3 n = normalize(normal);
	float3 l = normalize(light_dir);
	float3 h = normalize((l + view_dir));
	float ndotl = max(dot(n, l), 0.0);
	float ndoth = max(dot(n, h), 0.0);
	float3 f0 = mix(float3(0.04, 0.04, 0.04), albedo, material.metallic);
	float3 f = fresnel_schlick(max(dot(h, view_dir), 0.0), f0);
	float3 specular = (f * pow(ndoth, ((1.0 - material.roughness) * 256.0)));
	float3 diffuse = ((albedo * ndotl) * (float3(1.0, 1.0, 1.0) - f));
	return (diffuse + specular);
}

fragment FragmentOutput fs_main(FragmentInput input [[stage_in]],
	constant Material& material [[buffer(0)]],
	constant float3& camera_pos [[buffer(1)]],
	texture2d<float> albedo_tex_tex [[texture(0)]],
	sampler albedo_tex_samp [[sampler(0)]]) {
	FragmentOutput __luma_output;
	float4 tex_color = albedo_tex_tex.sample(albedo_tex_samp, input.uv);
	float3 albedo = (tex_color.rgb * material.albedo);
	float3 view_dir = normalize((camera_pos - input.world_pos));
	float3 color = compute_lighting(input.normal, view_dir, albedo);
	__luma_output.color = float4(color, tex_color.a);
	return __luma_output;
}