package main

import "core:fmt"
import "core:mem"
import "core:os"

import bk "../../backend"
import compiler "../../compiler"
import gpu "../.."
import proof "../parity_common"
import ir "../../render_ir"
import resource "../../resource"

WIDTH :: u32(256)
HEIGHT :: u32(256)

Instance_PNUC :: struct {
	position: [3]f32,
	normal:   [3]f32,
	uv:       [2]f32,
	color:    [4]f32,
}

VERTEX_SHADER :: `
@varying
struct VertexOutput
	@builtin(position) position: vec4
	@location(0) uv: vec2
end

@entry(vertex)
function main(@location(0) pos: vec3, @location(1) normal: vec3, @location(2) uv: vec2) -> VertexOutput
	return VertexOutput {
		position = vec4(pos, 1.0),
		uv = uv,
	}
end
`

INSTANCE_VERTEX_SHADER :: `
@varying
struct InstanceVertexOutput
	@builtin(position) position: vec4
	@location(0) color: vec4
end

@entry(vertex)
function main(
	@location(0) pos: vec3,
	@location(1) normal: vec3,
	@location(2) uv: vec2,
	@location(3) instance_pos: vec3,
	@location(4) instance_normal: vec3,
	@location(5) instance_uv: vec2,
	@location(6) instance_color: vec4
) -> InstanceVertexOutput
	return InstanceVertexOutput {
		position = vec4(pos + instance_pos, 1.0),
		color = instance_color,
	}
end
`

RED_FRAGMENT_SHADER :: `
struct FragmentOutput
	color: vec4
end

@entry(fragment)
function main() -> FragmentOutput
	return FragmentOutput {
		color = vec4(240.0 / 255.0, 64.0 / 255.0, 64.0 / 255.0, 1.0),
	}
end
`

INSTANCE_FRAGMENT_SHADER :: `
struct FragmentOutput
	color: vec4
end

@entry(fragment)
function main(@location(0) color: vec4) -> FragmentOutput
	return FragmentOutput {
		color = color,
	}
end
`

BLUE_FRAGMENT_SHADER :: `
struct FragmentOutput
	color: vec4
end

@entry(fragment)
function main() -> FragmentOutput
	return FragmentOutput {
		color = vec4(64.0 / 255.0, 64.0 / 255.0, 240.0 / 255.0, 1.0),
	}
end
`

SAMPLE_FRAGMENT_SHADER :: `
struct FragmentOutput
	color: vec4
end

group frame = 0
  @binding(0) texture offscreen_color: sampler2D
end

@entry(fragment)
function main(@location(0) uv: vec2) -> FragmentOutput
	return FragmentOutput {
		color = sample(offscreen_color, uv),
	}
end
`

quad_vertices :: proc(x0, y0, x1, y1: f32) -> [6]resource.Mesh_Vertex_PNU {
	n := [3]f32{0, 0, 1}
	return {
		{position = {x0, y0, 0}, normal = n, tex_coord = {0, 1}},
		{position = {x1, y0, 0}, normal = n, tex_coord = {1, 1}},
		{position = {x1, y1, 0}, normal = n, tex_coord = {1, 0}},
		{position = {x0, y0, 0}, normal = n, tex_coord = {0, 1}},
		{position = {x1, y1, 0}, normal = n, tex_coord = {1, 0}},
		{position = {x0, y1, 0}, normal = n, tex_coord = {0, 0}},
	}
}

write_bytes :: proc(backend: ^bk.Backend, buffer: bk.Buffer_Handle, data: rawptr, size: int) -> bool {
	mapped := backend.map_buffer(buffer)
	if mapped == nil {
		return false
	}
	mem.copy(mapped, data, size)
	backend.unmap_buffer(buffer)
	return true
}

write_ppm_rgba8 :: proc(path: string, rgba: []u8, width, height: u32) -> bool {
	if len(path) == 0 {
		return false
	}
	header := fmt.aprintf("P6\n%d %d\n255\n", width, height, allocator = context.temp_allocator)
	out := make([]u8, len(header) + int(width * height * 3))
	defer delete(out)
	copy(out[:len(header)], transmute([]u8)header)
	dst := len(header)
	for i in 0 ..< int(width * height) {
		out[dst + i * 3 + 0] = rgba[i * 4 + 0]
		out[dst + i * 3 + 1] = rgba[i * 4 + 1]
		out[dst + i * 3 + 2] = rgba[i * 4 + 2]
	}
	return os.write_entire_file(path, out) == nil
}

fail :: proc(message: string) {
	fmt.println(message)
	os.exit(1)
}

make_frame :: proc() -> (
	frame: ir.Frame_IR,
	color: ir.Resource_Handle,
	mrt_geometry: ir.Resource_Handle,
	indirect_geometry: ir.Resource_Handle,
	instance_data: ir.Resource_Handle,
	indirect_args: ir.Resource_Handle,
	offscreen_geometry: ir.Resource_Handle,
	sample_geometry: ir.Resource_Handle,
) {
	frame = ir.init_frame_ir()
	color = ir.add_texture(&frame, "proof-color", WIDTH, HEIGHT, .R8G8B8A8_UNORM, {.Color_Attachment, .Sampled}, .Transient)
	secondary := ir.add_texture(&frame, "mrt-secondary", WIDTH, HEIGHT, .R8G8B8A8_UNORM, {.Color_Attachment, .Sampled}, .Transient)
	offscreen := ir.add_texture(&frame, "offscreen-sampled-proof", WIDTH, HEIGHT, .R8G8B8A8_UNORM, {.Color_Attachment, .Sampled}, .Transient)
	sampler := ir.add_sampler(&frame, "nearest-sampler", {
		mag_filter = .Nearest,
		min_filter = .Nearest,
		address_mode_u = .Clamp_To_Edge,
		address_mode_v = .Clamp_To_Edge,
		mipmap_mode = .Nearest,
	}, .Transient)
	mrt_geometry = ir.add_buffer(&frame, "mrt-vertices", u64(6 * size_of(resource.Mesh_Vertex_PNU)), {.Vertex}, {.Host_Visible, .Host_Coherent}, .Transient)
	indirect_geometry = ir.add_buffer(&frame, "indirect-vertices", u64(6 * size_of(resource.Mesh_Vertex_PNU)), {.Vertex}, {.Host_Visible, .Host_Coherent}, .Transient)
	instance_data = ir.add_buffer(&frame, "instances", u64(size_of(Instance_PNUC)), {.Vertex}, {.Host_Visible, .Host_Coherent}, .Transient)
	indirect_args = ir.add_buffer(&frame, "indirect-args", u64(size_of(bk.Indirect_Draw_Args)), {.Indirect_Argument}, {.Host_Visible, .Host_Coherent}, .Transient)
	offscreen_geometry = ir.add_buffer(&frame, "offscreen-vertices", u64(6 * size_of(resource.Mesh_Vertex_PNU)), {.Vertex}, {.Host_Visible, .Host_Coherent}, .Transient)
	sample_geometry = ir.add_buffer(&frame, "sample-vertices", u64(6 * size_of(resource.Mesh_Vertex_PNU)), {.Vertex}, {.Host_Visible, .Host_Coherent}, .Transient)

	vs := ir.add_shader_source(&frame, "parity_backend_features.vert", .Vertex, VERTEX_SHADER)
	blue_fs := ir.add_shader_source(&frame, "parity_backend_features_blue.frag", .Fragment, BLUE_FRAGMENT_SHADER)
	offscreen_pass_desc := ir.Pass_Desc{kind = .Render, name = "write-offscreen-sampled"}
	_ = ir.pass_add_color_target(&offscreen_pass_desc, ir.make_color_target(offscreen, .R8G8B8A8_UNORM, .Clear, .Store, {16.0 / 255.0, 16.0 / 255.0, 24.0 / 255.0, 1}))
	ir.pass_set_viewport(&offscreen_pass_desc, {x = 0, y = 0, width = f32(WIDTH), height = f32(HEIGHT), min_depth = 0, max_depth = 1})
	ir.pass_set_scissor(&offscreen_pass_desc, {x = 0, y = 0, width = WIDTH, height = HEIGHT})
	offscreen_pass := ir.add_pass(&frame, offscreen_pass_desc)
	offscreen_target := ir.add_target(&frame, {kind = .Offscreen, name = "offscreen-sampled-target", width = WIDTH, height = HEIGHT, pass = offscreen_pass, color_format = .R8G8B8A8_UNORM, color_resource = offscreen})
	offscreen_state := ir.default_graphics_pipeline_state()
	offscreen_state.vertex_shader = vs
	offscreen_state.fragment_shader = blue_fs
	offscreen_state.color_format = .R8G8B8A8_UNORM
	offscreen_state.cull_mode = .None
	offscreen_pipeline := ir.add_graphics_pipeline(&frame, "offscreen-blue", offscreen_state)
	_ = ir.add_draw_packet(&frame, offscreen_pass, offscreen_pipeline, .Vertex_Stream, {
		target = offscreen_target,
		geometry = {kind = .Vertex_Stream, resource = offscreen_geometry, vertex_count = 6},
		order = .Strict,
	})

	pass_desc := ir.Pass_Desc{kind = .Render, name = "backend-feature-proof"}
	_ = ir.pass_add_color_target(&pass_desc, ir.make_color_target(color, .R8G8B8A8_UNORM, .Clear, .Store, {16.0 / 255.0, 16.0 / 255.0, 24.0 / 255.0, 1}))
	_ = ir.pass_add_color_target(&pass_desc, ir.make_color_target(secondary, .R8G8B8A8_UNORM, .Clear, .Store, {0, 0, 0, 1}))
	ir.pass_set_viewport(&pass_desc, {x = 0, y = 0, width = f32(WIDTH), height = f32(HEIGHT), min_depth = 0, max_depth = 1})
	ir.pass_set_scissor(&pass_desc, {x = 0, y = 0, width = WIDTH, height = HEIGHT})
	pass := ir.add_pass(&frame, pass_desc)
	target := ir.add_target(&frame, {kind = .Offscreen, name = "proof-target", width = WIDTH, height = HEIGHT, pass = pass, color_format = .R8G8B8A8_UNORM, color_resource = color})

	red_fs := ir.add_shader_source(&frame, "parity_backend_features_red.frag", .Fragment, RED_FRAGMENT_SHADER)
	red_state := ir.default_graphics_pipeline_state()
	red_state.vertex_shader = vs
	red_state.fragment_shader = red_fs
	red_state.color_format = .R8G8B8A8_UNORM
	red_state.cull_mode = .None
	red_pipeline := ir.add_graphics_pipeline(&frame, "mrt-red", red_state)
	_ = ir.add_draw_packet(&frame, pass, red_pipeline, .Vertex_Stream, {
		target = target,
		geometry = {kind = .Vertex_Stream, resource = mrt_geometry, vertex_count = 6},
		order = .Strict,
	})

	instance_vs := ir.add_shader_source(&frame, "parity_backend_features_instance.vert", .Vertex, INSTANCE_VERTEX_SHADER)
	instance_fs := ir.add_shader_source(&frame, "parity_backend_features_instance.frag", .Fragment, INSTANCE_FRAGMENT_SHADER)
	instance_state := ir.default_graphics_pipeline_state()
	instance_state.vertex_shader = instance_vs
	instance_state.fragment_shader = instance_fs
	instance_state.color_format = .R8G8B8A8_UNORM
	instance_state.cull_mode = .None
	instance_state.has_instance_layout = true
	instance_state.instance_layout_variant = .PNUC
	instance_pipeline := ir.add_graphics_pipeline(&frame, "indirect-instance", instance_state)
	_ = ir.add_draw_packet(&frame, pass, instance_pipeline, .Indirect, {
		target = target,
		geometry = {kind = .Vertex_Stream, resource = indirect_geometry, vertex_count = 6},
		instances = {
			resource = instance_data,
			indirect = indirect_args,
			count = 1,
			stride = u32(size_of(bk.Indirect_Draw_Args)),
		},
		order = .Strict,
	})
	layout_desc := ir.make_descriptor_set_layout("offscreen-sampled-layout")
	_ = ir.descriptor_layout_add_binding(&layout_desc, 0, .Combined_Image_Sampler, 1, {.Fragment})
	layout := ir.add_descriptor_set_layout(&frame, layout_desc)
	set_desc := ir.make_descriptor_set("offscreen-sampled-set", layout)
	_ = ir.descriptor_set_bind_texture_sampler(&set_desc, 0, offscreen, sampler)
	set := ir.add_descriptor_set(&frame, set_desc)
	sample_fs := ir.add_shader_source(&frame, "parity_backend_features_sample.frag", .Fragment, SAMPLE_FRAGMENT_SHADER)
	sample_state := ir.default_graphics_pipeline_state()
	sample_state.vertex_shader = vs
	sample_state.fragment_shader = sample_fs
	sample_state.color_format = .R8G8B8A8_UNORM
	sample_state.cull_mode = .None
	sample_pipeline_desc := ir.Pipeline_Desc {
		kind = .Graphics,
		name = "sample-offscreen",
		graphics = sample_state,
	}
	_ = ir.pipeline_add_descriptor_layout(&sample_pipeline_desc, layout)
	sample_pipeline := ir.add_pipeline(&frame, sample_pipeline_desc)
	material := ir.add_material(&frame, {name = "sample-offscreen", pipeline = sample_pipeline, descriptor_set = set})
	_ = ir.add_draw_packet(&frame, pass, sample_pipeline, .Vertex_Stream, {
		target = target,
		material = material,
		geometry = {kind = .Vertex_Stream, resource = sample_geometry, vertex_count = 6},
		order = .Strict,
		sort_key = 2,
	})
	return
}

main :: proc() {
	runtime, ok := proof.init_runtime(WIDTH, HEIGHT, "GPU parity backend features")
	if !ok {
		fail("parity_backend_features: failed to initialize proof runtime")
	}
	defer proof.shutdown_runtime(&runtime)

	backend := gpu.get_backend()
	if backend == nil || backend.read_texture_rgba8 == nil {
		fail("parity_backend_features: backend readback unsupported")
	}

	frame, color, mrt_geometry, indirect_geometry, instance_data, indirect_args, offscreen_geometry, sample_geometry := make_frame()
	defer ir.destroy_frame_ir(&frame)
	exec := compiler.init_execution_state(backend)
	defer compiler.destroy_execution_state(&exec)
	if !compiler.prepare_planned_frame(&exec, &frame) {
		fail("parity_backend_features: failed to prepare planned frame")
	}
	mrt_buffer, mrt_ok := compiler.prepared_buffer(&exec, mrt_geometry)
	indirect_buffer, indirect_ok := compiler.prepared_buffer(&exec, indirect_geometry)
	instance_buffer, instance_ok := compiler.prepared_buffer(&exec, instance_data)
	args_buffer, args_ok := compiler.prepared_buffer(&exec, indirect_args)
	offscreen_buffer, offscreen_ok := compiler.prepared_buffer(&exec, offscreen_geometry)
	sample_buffer, sample_ok := compiler.prepared_buffer(&exec, sample_geometry)
	if !mrt_ok || !indirect_ok || !instance_ok || !args_ok || !offscreen_ok || !sample_ok {
		fail("parity_backend_features: failed to resolve prepared buffers")
	}
	mrt_vertices := quad_vertices(-0.85, -0.5, -0.35, 0.5)
	indirect_vertices := quad_vertices(-0.25, -0.5, 0.25, 0.5)
	offscreen_vertices := quad_vertices(-1, -1, 1, 1)
	sample_vertices := quad_vertices(0.35, -0.5, 0.85, 0.5)
	instance := Instance_PNUC{normal = {0, 0, 1}, color = {64.0 / 255.0, 240.0 / 255.0, 64.0 / 255.0, 1}}
	args := bk.Indirect_Draw_Args{vertex_count = 6, instance_count = 1}
	if !write_bytes(backend, mrt_buffer, raw_data(mrt_vertices[:]), len(mrt_vertices) * size_of(resource.Mesh_Vertex_PNU)) ||
	   !write_bytes(backend, indirect_buffer, raw_data(indirect_vertices[:]), len(indirect_vertices) * size_of(resource.Mesh_Vertex_PNU)) ||
	   !write_bytes(backend, instance_buffer, &instance, size_of(Instance_PNUC)) ||
	   !write_bytes(backend, args_buffer, &args, size_of(bk.Indirect_Draw_Args)) ||
	   !write_bytes(backend, offscreen_buffer, raw_data(offscreen_vertices[:]), len(offscreen_vertices) * size_of(resource.Mesh_Vertex_PNU)) ||
	   !write_bytes(backend, sample_buffer, raw_data(sample_vertices[:]), len(sample_vertices) * size_of(resource.Mesh_Vertex_PNU)) {
		fail("parity_backend_features: failed to upload proof buffers")
	}
	if !compiler.execute_prepared_planned_frame(&exec, &frame, {0, 0, 0, 1}) {
		fail("parity_backend_features: failed to execute planned frame")
	}
	texture, tex_ok := compiler.prepared_texture(&exec, color)
	if !tex_ok {
		fail("parity_backend_features: failed to resolve readback target")
	}
	rgba := make([]u8, int(WIDTH * HEIGHT * 4))
	defer delete(rgba)
	if !backend.read_texture_rgba8({texture = texture, width = WIDTH, height = HEIGHT, current_layout = .Shader_Read_Only}, rgba) {
		fail("parity_backend_features: readback failed")
	}
	readback_path := os.get_env("GPU_PROOF_READBACK", context.allocator)
	if !write_ppm_rgba8(readback_path, rgba, WIDTH, HEIGHT) {
		fail("parity_backend_features: failed to write readback")
	}
	profile_path := os.get_env("GPU_PROFILE_JSON", context.allocator)
	if len(profile_path) > 0 {
		_ = gpu.write_profile(profile_path, 0, 1)
	}
}