-- PBR fragment shader — tests helper function calls, multiple builtins,
-- struct construction, and complex expressions.

struct FragmentInput
	uv: vec2
	normal: vec3
	world_pos: vec3
end

struct FragmentOutput
	color: vec4
end

struct Material
	albedo: vec3
	metallic: float
	roughness: float
end

uniform material: Material
uniform light_dir: vec3
uniform camera_pos: vec3
uniform albedo_tex: sampler2D

function fresnel_schlick(cos_theta: float, f0: vec3) -> vec3
	return f0 + (vec3(1.0, 1.0, 1.0) - f0) * pow(1.0 - cos_theta, 5.0)
end

function compute_lighting(normal: vec3, view_dir: vec3, albedo: vec3) -> vec3
	let n = normalize(normal)
	let l = normalize(light_dir)
	let h = normalize(l + view_dir)

	let ndotl = max(dot(n, l), 0.0)
	let ndoth = max(dot(n, h), 0.0)

	let f0 = mix(vec3(0.04, 0.04, 0.04), albedo, material.metallic)
	let f = fresnel_schlick(max(dot(h, view_dir), 0.0), f0)

	let specular = f * pow(ndoth, (1.0 - material.roughness) * 256.0)
	let diffuse = albedo * ndotl * (vec3(1.0, 1.0, 1.0) - f)

	return diffuse + specular
end

@entry(fragment)
function fs_main(input: FragmentInput) -> FragmentOutput
	let tex_color = sample(albedo_tex, input.uv)
	let albedo = tex_color.rgb * material.albedo
	let view_dir = normalize(camera_pos - input.world_pos)
	let color = compute_lighting(input.normal, view_dir, albedo)
	return FragmentOutput {
		color = vec4(color, tex_color.a),
	}
end