package app

import "core:log"
import gpu ".."
import shaderkit "../shaderkit"

DEFAULT_2D_VERT :: `
struct PC
  proj: mat4
end

@push_constant
uniform pc: PC

struct In
  @location(0) position: vec3
  @location(1) normal: vec3
  @location(2) uv: vec2
  @location(3) color: vec4
end

@varying
struct Varying
  @builtin(position) position: vec4
  @location(0) color: vec4
end

@entry(vertex)
function main(input: In) -> Varying
  return Varying{position = pc.proj * vec4(input.position, 1.0), color = input.color}
end
`

DEFAULT_COLOR_FRAG :: `
@varying
struct Varying
  @builtin(position) position: vec4
  @location(0) color: vec4
end

@entry(fragment)
function main(input: Varying) -> vec4
  return input.color
end
`

DEFAULT_3D_PNU_VERT :: `
struct PC
  proj_view: mat4
  model: mat4
end

@push_constant
uniform pc: PC

struct In
  @location(0) position: vec3
  @location(1) normal: vec3
  @location(2) uv: vec2
end

@varying
struct Varying
  @builtin(position) position: vec4
  @location(0) color: vec4
end

@entry(vertex)
function main(input: In) -> Varying
  return Varying{position = pc.proj_view * pc.model * vec4(input.position, 1.0), color = vec4(1.0, 1.0, 1.0, 1.0)}
end
`

DEFAULT_3D_PNUC_VERT :: `
struct PC
  proj_view: mat4
  model: mat4
end

@push_constant
uniform pc: PC

struct In
  @location(0) position: vec3
  @location(1) normal: vec3
  @location(2) uv: vec2
  @location(3) color: vec4
end

@varying
struct Varying
  @builtin(position) position: vec4
  @location(0) color: vec4
end

@entry(vertex)
function main(input: In) -> Varying
  return Varying{position = pc.proj_view * pc.model * vec4(input.position, 1.0), color = input.color}
end
`

DEFAULT_3D_LIT_PNU_VERT :: `
struct PC
  model: mat4
  base_color: vec4
  metallic: float
  roughness: float
  reflectance: float
  pad: float
  emissive: vec4
end

@push_constant
uniform pc: PC

struct LightUBO
  proj_view: mat4
 camera_pos: vec4
 ambient_color: vec4
 light_count: vec4
end

@group(1) @binding(0)
uniform lights: LightUBO

struct In
  @location(0) position: vec3
  @location(1) normal: vec3
  @location(2) uv: vec2
end

@varying
struct Varying
  @builtin(position) position: vec4
  @location(0) world_pos: vec3
  @location(1) normal: vec3
  @location(2) base_color: vec4
  @location(3) material: vec4
  @location(4) emissive: vec4
end

@entry(vertex)
function main(input: In) -> Varying
  let world = pc.model * vec4(input.position, 1.0)
  let world_normal = pc.model * vec4(input.normal, 0.0)
  return Varying{
    position = lights.proj_view * world,
    world_pos = world.xyz,
    normal = normalize(world_normal.xyz),
    base_color = pc.base_color,
    material = vec4(pc.metallic, pc.roughness, pc.reflectance, 0.0),
    emissive = pc.emissive,
  }
end
`

DEFAULT_3D_LIT_FRAG :: `
struct GPU_Light
  position: vec4
  color: vec4
  params: vec4
end

struct LightUBO
  proj_view: mat4
  camera_pos: vec4
  ambient_color: vec4
  light_count: vec4
  lights: [336]GPU_Light
end

@group(1) @binding(0)
uniform lights: LightUBO

@varying
struct Varying
  @builtin(position) position: vec4
  @location(0) world_pos: vec3
  @location(1) normal: vec3
  @location(2) base_color: vec4
  @location(3) material: vec4
  @location(4) emissive: vec4
end

function pbr_light_values(light_position: vec4, light_color: vec4, light_params: vec4, world_pos: vec3, normal: vec3, view_dir: vec3, albedo: vec3, metallic: float, roughness: float, reflectance: float) -> vec3
  if light_color.w <= 0.0 then
    return vec3(0.0, 0.0, 0.0)
  end
  var dir: vec3 = vec3(0.0, 1.0, 0.0)
  var attenuation: float = 1.0
  if light_position.w < 0.5 then
    dir = normalize(-light_position.xyz)
  else
    let to_light = light_position.xyz - world_pos
    let dist = length(to_light)
    dir = normalize(to_light)
    let radius = max(light_params.x, 0.001)
    attenuation = max(0.0, 1.0 - dist / radius)
  end
  let ndl = max(dot(normal, dir), 0.0)
  let half_dir = normalize(dir + view_dir)
  let ndh = max(dot(normal, half_dir), 0.0)
  let spec_power = max(1.0, (1.0 - roughness) * 96.0)
  let specular = pow(ndh, spec_power) * reflectance
  let diffuse = albedo * (1.0 - metallic)
  return (diffuse + vec3(specular, specular, specular)) * light_color.xyz * ndl * attenuation
end

function lighting(world_pos: vec3, normal: vec3, albedo: vec3, metallic: float, roughness: float, reflectance: float) -> vec3
  let view_dir = normalize(lights.camera_pos.xyz - world_pos)
  var lit: vec3 = albedo * lights.ambient_color.xyz
  lit = lit + pbr_light_values(lights.lights[0].position, lights.lights[0].color, lights.lights[0].params, world_pos, normal, view_dir, albedo, metallic, roughness, reflectance)
  lit = lit + pbr_light_values(lights.lights[1].position, lights.lights[1].color, lights.lights[1].params, world_pos, normal, view_dir, albedo, metallic, roughness, reflectance)
  lit = lit + pbr_light_values(lights.lights[2].position, lights.lights[2].color, lights.lights[2].params, world_pos, normal, view_dir, albedo, metallic, roughness, reflectance)
  lit = lit + pbr_light_values(lights.lights[3].position, lights.lights[3].color, lights.lights[3].params, world_pos, normal, view_dir, albedo, metallic, roughness, reflectance)
  lit = lit + pbr_light_values(lights.lights[4].position, lights.lights[4].color, lights.lights[4].params, world_pos, normal, view_dir, albedo, metallic, roughness, reflectance)
  lit = lit + pbr_light_values(lights.lights[5].position, lights.lights[5].color, lights.lights[5].params, world_pos, normal, view_dir, albedo, metallic, roughness, reflectance)
  lit = lit + pbr_light_values(lights.lights[6].position, lights.lights[6].color, lights.lights[6].params, world_pos, normal, view_dir, albedo, metallic, roughness, reflectance)
  lit = lit + pbr_light_values(lights.lights[7].position, lights.lights[7].color, lights.lights[7].params, world_pos, normal, view_dir, albedo, metallic, roughness, reflectance)
  let d1 = length(lights.lights[1].position.xyz - world_pos)
  let f1 = max(0.0, 1.0 - d1 / max(lights.lights[1].params.x, 0.001))
  let d2 = length(lights.lights[2].position.xyz - world_pos)
  let f2 = max(0.0, 1.0 - d2 / max(lights.lights[2].params.x, 0.001))
  let d3 = length(lights.lights[3].position.xyz - world_pos)
  let f3 = max(0.0, 1.0 - d3 / max(lights.lights[3].params.x, 0.001))
  let d4 = length(lights.lights[4].position.xyz - world_pos)
  let f4 = max(0.0, 1.0 - d4 / max(lights.lights[4].params.x, 0.001))
  lit = lit + albedo * lights.lights[1].color.xyz * f1 * 0.25
  lit = lit + albedo * lights.lights[2].color.xyz * f2 * 0.25
  lit = lit + albedo * lights.lights[3].color.xyz * f3 * 0.25
  lit = lit + albedo * lights.lights[4].color.xyz * f4 * 0.25
  return lit
end

@entry(fragment)
function main(input: Varying) -> vec4
  let n = normalize(input.normal)
  let lit = lighting(input.world_pos, n, input.base_color.xyz, input.material.x, input.material.y, input.material.z)
  return vec4(lit * 1.4 + input.emissive.xyz, input.base_color.w)
end
`

DEFAULT_3D_LIT_SHADOW_FRAG :: `
struct GPU_Light
  position: vec4
  color: vec4
  params: vec4
end

struct LightUBO
  proj_view: mat4
  camera_pos: vec4
  ambient_color: vec4
  light_count: vec4
  lights: [336]GPU_Light
end

@group(1) @binding(0)
uniform lights: LightUBO

struct ShadowUBO
  light_matrices: [4]mat4
  shadow_count: uvec4
  shadow_params: vec4
  shadow_light_indices: uvec4
  atlas_regions: [4]vec4
end

@group(2) @binding(0)
uniform shadow: ShadowUBO

@group(2) @binding(1)
uniform shadow_atlas: sampler2DShadow

@varying
struct Varying
  @builtin(position) position: vec4
  @location(0) world_pos: vec3
  @location(1) normal: vec3
  @location(2) base_color: vec4
  @location(3) material: vec4
  @location(4) emissive: vec4
end

function pbr_light_values(light_position: vec4, light_color: vec4, light_params: vec4, world_pos: vec3, normal: vec3, view_dir: vec3, albedo: vec3, metallic: float, roughness: float, reflectance: float, shadow: float) -> vec3
  if light_color.w <= 0.0 then
    return vec3(0.0, 0.0, 0.0)
  end
  var dir: vec3 = vec3(0.0, 1.0, 0.0)
  var attenuation: float = 1.0
  if light_position.w < 0.5 then
    dir = normalize(-light_position.xyz)
  else
    let to_light = light_position.xyz - world_pos
    let dist = length(to_light)
    dir = normalize(to_light)
    let radius = max(light_params.x, 0.001)
    attenuation = max(0.0, 1.0 - dist / radius)
  end
  let ndl = max(dot(normal, dir), 0.0)
  let half_dir = normalize(dir + view_dir)
  let ndh = max(dot(normal, half_dir), 0.0)
  let spec_power = max(1.0, (1.0 - roughness) * 96.0)
  let specular = pow(ndh, spec_power) * reflectance
  let diffuse = albedo * (1.0 - metallic)
  return (diffuse + vec3(specular, specular, specular)) * light_color.xyz * ndl * attenuation * shadow
end

function shadow_factor(world_pos: vec3, normal: vec3) -> float
  if shadow.shadow_count.x == 0 then
    return 1.0
  end
  let shadow_pos = world_pos + normal * shadow.shadow_params.y
  let lp = shadow.light_matrices[0] * vec4(shadow_pos, 1.0)
  let ndc = lp.xyz / lp.w
  if ndc.x < -1.0 or ndc.x > 1.0 or ndc.y < -1.0 or ndc.y > 1.0 or ndc.z < 0.0 or ndc.z > 1.0 then
    return 1.0
  end
  let region = shadow.atlas_regions[0]
  let uv = (ndc.xy * 0.5 + vec2(0.5, 0.5)) * region.zw + region.xy
  let depth = ndc.z
  let visible = sample_shadow(shadow_atlas, uv, depth - shadow.shadow_params.x)
  return 0.25 + 0.75 * visible
end

function lighting(world_pos: vec3, normal: vec3, albedo: vec3, metallic: float, roughness: float, reflectance: float) -> vec3
  let view_dir = normalize(lights.camera_pos.xyz - world_pos)
  let s = shadow_factor(world_pos, normal)
  var lit: vec3 = albedo * lights.ambient_color.xyz
  lit = lit + pbr_light_values(lights.lights[0].position, lights.lights[0].color, lights.lights[0].params, world_pos, normal, view_dir, albedo, metallic, roughness, reflectance, s)
  lit = lit + pbr_light_values(lights.lights[1].position, lights.lights[1].color, lights.lights[1].params, world_pos, normal, view_dir, albedo, metallic, roughness, reflectance, 1.0)
  lit = lit + pbr_light_values(lights.lights[2].position, lights.lights[2].color, lights.lights[2].params, world_pos, normal, view_dir, albedo, metallic, roughness, reflectance, 1.0)
  lit = lit + pbr_light_values(lights.lights[3].position, lights.lights[3].color, lights.lights[3].params, world_pos, normal, view_dir, albedo, metallic, roughness, reflectance, 1.0)
  lit = lit + pbr_light_values(lights.lights[4].position, lights.lights[4].color, lights.lights[4].params, world_pos, normal, view_dir, albedo, metallic, roughness, reflectance, 1.0)
  lit = lit + pbr_light_values(lights.lights[5].position, lights.lights[5].color, lights.lights[5].params, world_pos, normal, view_dir, albedo, metallic, roughness, reflectance, 1.0)
  lit = lit + pbr_light_values(lights.lights[6].position, lights.lights[6].color, lights.lights[6].params, world_pos, normal, view_dir, albedo, metallic, roughness, reflectance, 1.0)
  lit = lit + pbr_light_values(lights.lights[7].position, lights.lights[7].color, lights.lights[7].params, world_pos, normal, view_dir, albedo, metallic, roughness, reflectance, 1.0)
  let d1 = length(lights.lights[1].position.xyz - world_pos)
  let f1 = max(0.0, 1.0 - d1 / max(lights.lights[1].params.x, 0.001))
  let d2 = length(lights.lights[2].position.xyz - world_pos)
  let f2 = max(0.0, 1.0 - d2 / max(lights.lights[2].params.x, 0.001))
  let d3 = length(lights.lights[3].position.xyz - world_pos)
  let f3 = max(0.0, 1.0 - d3 / max(lights.lights[3].params.x, 0.001))
  let d4 = length(lights.lights[4].position.xyz - world_pos)
  let f4 = max(0.0, 1.0 - d4 / max(lights.lights[4].params.x, 0.001))
  lit = lit + albedo * lights.lights[1].color.xyz * f1 * 0.25
  lit = lit + albedo * lights.lights[2].color.xyz * f2 * 0.25
  lit = lit + albedo * lights.lights[3].color.xyz * f3 * 0.25
  lit = lit + albedo * lights.lights[4].color.xyz * f4 * 0.25
  return lit
end

@entry(fragment)
function main(input: Varying) -> vec4
  let n = normalize(input.normal)
  let lit = lighting(input.world_pos, n, input.base_color.xyz, input.material.x, input.material.y, input.material.z)
  let s = shadow_factor(input.world_pos, n)
  let contrast = 0.35 + 0.65 * s
  return vec4(lit * contrast * 1.4 + input.emissive.xyz, input.base_color.w)
end
`

DEFAULT_SHADOW_PNU_VERT :: `
struct PC
  light_proj_view: mat4
  model: mat4
end

@push_constant
uniform pc: PC

struct In
  @location(0) position: vec3
  @location(1) normal: vec3
  @location(2) uv: vec2
end

struct Out
  @builtin(position) position: vec4
end

@entry(vertex)
function main(input: In) -> Out
  return Out{position = pc.light_proj_view * pc.model * vec4(input.position, 1.0)}
end
`

Default_Pipelines_State :: struct {
	pipeline_2d: gpu.Pipeline_Handle,
	shaders:     [10]gpu.Shader_Handle,
	shader_count: int,
}

default_compile_shader :: proc(state: ^Default_Pipelines_State, source, name: string, stage: gpu.Shader_Stage) -> (gpu.Shader_Handle, bool) {
	handle, ok := shaderkit.create_shader_from_source(name, stage, source, {opt_level = .Basic})
	if !ok {
		return {}, false
	}

	state.shaders[state.shader_count] = handle
	state.shader_count += 1
	return handle, true
}

default_create_pipeline_with_layouts :: proc(
	vs, fs: gpu.Shader_Handle,
	layout: gpu.Vertex_Layout,
	cull: gpu.Cull_Mode,
	blending, depth: bool,
	push_constant_size: u32,
	descriptor_layouts: []gpu.Descriptor_Handle,
	render_pass: gpu.Render_Pass_Handle = {},
	depth_only: bool = false,
) -> (gpu.Pipeline_Handle, bool) {
	layout_copy := layout
	binding, attrs, attr_count := gpu.vertex_layout_to_pipeline_attrs(&layout_copy)
	vertex_bindings := [1]gpu.Vertex_Binding{binding}
	rp := render_pass
	if rp == {} {
		rp = gpu.get_default_render_pass()
	}
	desc := gpu.Pipeline_Desc{
		vert_shader          = vs,
		frag_shader          = fs,
		render_pass          = rp,
		topology             = .Triangle_List,
		cull_mode            = cull,
		front_face           = .Counter_Clockwise,
		enable_blending      = blending,
		blend_mode           = .Alpha,
		enable_depth_test    = depth,
		depth_only           = depth_only,
		push_constant_size   = push_constant_size,
		push_constant_stages = {.Vertex},
		descriptor_layouts   = descriptor_layouts[:],
		vertex_bindings      = vertex_bindings[:],
		vertex_attributes    = attrs[:attr_count],
	}
	return gpu.create_pipeline(desc)
}

default_create_pipeline :: proc(vs, fs: gpu.Shader_Handle, layout: gpu.Vertex_Layout, cull: gpu.Cull_Mode, blending, depth: bool, push_constant_size: u32) -> (gpu.Pipeline_Handle, bool) {
	descriptor_layouts := [1]gpu.Descriptor_Handle{gpu.get_texture_descriptor_layout()}
	return default_create_pipeline_with_layouts(vs, fs, layout, cull, blending, depth, push_constant_size, descriptor_layouts[:])
}

default_pipelines_init :: proc(state: ^State) -> bool {
	vs2d, ok := default_compile_shader(&state.defaults, DEFAULT_2D_VERT, "default_2d.vert", .Vertex)
	if !ok { return false }
	fs_color, fs_color_ok := default_compile_shader(&state.defaults, DEFAULT_COLOR_FRAG, "default_color.frag", .Fragment)
	if !fs_color_ok { return false }
	fs_lit, fs_lit_ok := default_compile_shader(&state.defaults, DEFAULT_3D_LIT_FRAG, "default_3d_lit.frag", .Fragment)
	if !fs_lit_ok { return false }
	vs3d_pnu, vs3d_pnu_ok := default_compile_shader(&state.defaults, DEFAULT_3D_PNU_VERT, "default_3d_pnu.vert", .Vertex)
	if !vs3d_pnu_ok { return false }
	vs3d_pnuc, vs3d_pnuc_ok := default_compile_shader(&state.defaults, DEFAULT_3D_PNUC_VERT, "default_3d_pnuc.vert", .Vertex)
	if !vs3d_pnuc_ok { return false }
	vs3d_lit_pnu, vs3d_lit_pnu_ok := default_compile_shader(&state.defaults, DEFAULT_3D_LIT_PNU_VERT, "default_3d_lit_pnu.vert", .Vertex)
	if !vs3d_lit_pnu_ok { return false }

	p2d, p2d_ok := default_create_pipeline(vs2d, fs_color, gpu.LAYOUT_POS_NORM_UV_COLOR, .None, true, false, 64)
	if !p2d_ok { return false }
	state.defaults.pipeline_2d = p2d
	gpu.set_pipeline_2d(p2d)
	gpu.set_pipeline_2d_shaders(vs2d, fs_color)

	p3d_culled, p3d_culled_ok := default_create_pipeline(vs3d_pnu, fs_color, gpu.LAYOUT_POS_NORM_UV, .Back, false, true, 128)
	if !p3d_culled_ok { return false }
	p3d_double, p3d_double_ok := default_create_pipeline(vs3d_pnu, fs_color, gpu.LAYOUT_POS_NORM_UV, .None, false, true, 128)
	if !p3d_double_ok { return false }
	gpu.set_pipeline_3d(.PNU, p3d_culled, p3d_double)

	lit_layouts := [2]gpu.Descriptor_Handle{gpu.get_texture_descriptor_layout(), gpu.get_light_descriptor_layout()}
	p3d_lit, p3d_lit_ok := default_create_pipeline_with_layouts(vs3d_lit_pnu, fs_lit, gpu.LAYOUT_POS_NORM_UV, .Back, false, true, 128, lit_layouts[:])
	if !p3d_lit_ok { return false }
	p3d_lit_double, p3d_lit_double_ok := default_create_pipeline_with_layouts(vs3d_lit_pnu, fs_lit, gpu.LAYOUT_POS_NORM_UV, .None, false, true, 128, lit_layouts[:])
	if !p3d_lit_double_ok { return false }
	gpu.set_pipeline_3d_lit(.PNU, p3d_lit, p3d_lit_double)

	when gpu.REQUIRED_SHADER_FORMAT == .SPIRV {
		if !gpu.init_shadows() {
			log.warn("gpu/app: failed to initialize default shadows")
		} else {
			vs_shadow, vs_shadow_ok := default_compile_shader(&state.defaults, DEFAULT_SHADOW_PNU_VERT, "default_shadow_pnu.vert", .Vertex)
			if !vs_shadow_ok { return false }
			shadow_layouts := [0]gpu.Descriptor_Handle{}
			shadow_depth, shadow_depth_ok := default_create_pipeline_with_layouts(
				vs_shadow,
				fs_color,
				gpu.LAYOUT_POS_NORM_UV,
				.None,
				false,
				true,
				128,
				shadow_layouts[:],
				gpu.get_shadow_render_pass(),
				true,
			)
			if !shadow_depth_ok { return false }
			gpu.set_shadow_depth_pipeline(.PNU, shadow_depth)

			fs_shadow, fs_shadow_ok := default_compile_shader(&state.defaults, DEFAULT_3D_LIT_SHADOW_FRAG, "default_3d_lit_shadow.frag", .Fragment)
			if !fs_shadow_ok { return false }
			shadowed_layouts := [3]gpu.Descriptor_Handle{gpu.get_texture_descriptor_layout(), gpu.get_light_descriptor_layout(), gpu.get_shadow_descriptor_layout()}
			p3d_shadowed, p3d_shadowed_ok := default_create_pipeline_with_layouts(vs3d_lit_pnu, fs_shadow, gpu.LAYOUT_POS_NORM_UV, .Back, false, true, 128, shadowed_layouts[:])
			if !p3d_shadowed_ok { return false }
			p3d_shadowed_double, p3d_shadowed_double_ok := default_create_pipeline_with_layouts(vs3d_lit_pnu, fs_shadow, gpu.LAYOUT_POS_NORM_UV, .None, false, true, 128, shadowed_layouts[:])
			if !p3d_shadowed_double_ok { return false }
			gpu.set_pipeline_3d_lit_shadowed(.PNU, p3d_shadowed, p3d_shadowed_double)
		}
	}

	p3d_particles, p3d_particles_ok := default_create_pipeline(vs3d_pnuc, fs_color, gpu.LAYOUT_POS_NORM_UV_COLOR, .None, true, true, 128)
	if !p3d_particles_ok { return false }
	gpu.set_pipeline_3d(.PNUC, p3d_particles, p3d_particles)

	return true
}

default_pipelines_shutdown :: proc(state: ^State) {
	if state == nil {
		return
	}
	gpu.destroy_offscreen_2d_pipelines()
	if state.defaults.pipeline_2d != gpu.NULL_PIPELINE {
		gpu.destroy_pipeline(state.defaults.pipeline_2d)
		state.defaults.pipeline_2d = gpu.NULL_PIPELINE
	}
	for i in 0..<state.defaults.shader_count {
		gpu.destroy_shader(state.defaults.shaders[i])
		state.defaults.shaders[i] = {}
	}
	state.defaults.shader_count = 0
}