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)

COMPUTE_SHADER :: `
struct Vertices
	data: [48]float
end

@group(0) @binding(0)
buffer vertices: Vertices

@entry(compute)
@workgroup_size(1, 1, 1)
function main(@builtin(global_invocation_id) gid: uvec3)
	let z = 0.0
	vertices.data[0] = -1.0
	vertices.data[1] = -1.0
	vertices.data[2] = z
	vertices.data[3] = 0.0
	vertices.data[4] = 0.0
	vertices.data[5] = 1.0
	vertices.data[6] = 0.0
	vertices.data[7] = 1.0
	vertices.data[8] = 0.0
	vertices.data[9] = -1.0
	vertices.data[10] = z
	vertices.data[11] = 0.0
	vertices.data[12] = 0.0
	vertices.data[13] = 1.0
	vertices.data[14] = 1.0
	vertices.data[15] = 1.0
	vertices.data[16] = 0.0
	vertices.data[17] = 1.0
	vertices.data[18] = z
	vertices.data[19] = 0.0
	vertices.data[20] = 0.0
	vertices.data[21] = 1.0
	vertices.data[22] = 1.0
	vertices.data[23] = 0.0
	vertices.data[24] = -1.0
	vertices.data[25] = -1.0
	vertices.data[26] = z
	vertices.data[27] = 0.0
	vertices.data[28] = 0.0
	vertices.data[29] = 1.0
	vertices.data[30] = 0.0
	vertices.data[31] = 1.0
	vertices.data[32] = 0.0
	vertices.data[33] = 1.0
	vertices.data[34] = z
	vertices.data[35] = 0.0
	vertices.data[36] = 0.0
	vertices.data[37] = 1.0
	vertices.data[38] = 1.0
	vertices.data[39] = 0.0
	vertices.data[40] = -1.0
	vertices.data[41] = 1.0
	vertices.data[42] = z
	vertices.data[43] = 0.0
	vertices.data[44] = 0.0
	vertices.data[45] = 1.0
	vertices.data[46] = 0.0
	vertices.data[47] = 0.0
end
`

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
`

FRAGMENT_SHADER :: `
fragment main
  out color: vec4
do
  color = vec4(48.0 / 255.0, 224.0 / 255.0, 128.0 / 255.0, 1.0)
end
`

OFFSCREEN_FRAGMENT_SHADER :: `
fragment main
  out color: vec4
do
  color = vec4(64.0 / 255.0, 128.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, 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)
	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)
	vertices := ir.add_buffer(
		&frame,
		"compute-written-vertices",
		u64(6 * size_of(resource.Mesh_Vertex_PNU)),
		{.Storage, .Vertex},
		{.Device_Local},
		.Transient,
	)
	sample_geometry = ir.add_buffer(
		&frame,
		"sample-quad",
		u64(6 * size_of(resource.Mesh_Vertex_PNU)),
		{.Vertex},
		{.Host_Visible, .Host_Coherent},
		.Transient,
	)
	compute_shader := ir.add_shader_source(&frame, "parity_resource_sync.comp", .Compute, COMPUTE_SHADER)
	compute_pipeline := ir.add_compute_pipeline(
		&frame,
		"write-vertices",
		ir.default_compute_pipeline_state(compute_shader),
	)
	compute_pass := ir.add_pass(&frame, {kind = .Compute, name = "compute"})
	storage := [?]ir.Descriptor_Resource_Binding{{
		binding = 0,
		type = .Storage_Buffer,
		resource = vertices,
		size = u64(6 * size_of(resource.Mesh_Vertex_PNU)),
	}}
	_ = ir.add_compute_dispatch(&frame, compute_pass, compute_pipeline, 0, {1, 1, 1}, storage[:])

	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,
			clear_color = {16.0 / 255.0, 16.0 / 255.0, 24.0 / 255.0, 1},
		},
	)
	vs := ir.add_shader_source(&frame, "parity_resource_sync.vert", .Vertex, VERTEX_SHADER)
	offscreen_fs := ir.add_shader_source(&frame, "parity_resource_sync_offscreen.frag", .Fragment, OFFSCREEN_FRAGMENT_SHADER)
	offscreen_state := ir.default_graphics_pipeline_state()
	offscreen_state.vertex_shader = vs
	offscreen_state.fragment_shader = offscreen_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 = vertices, vertex_count = 6},
			order = .Strict,
		},
	)

	pass_desc := ir.Pass_Desc{kind = .Render, name = "draw-compute-output"}
	_ = 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_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,
			clear_color = {16.0 / 255.0, 16.0 / 255.0, 24.0 / 255.0, 1},
		},
	)
	fs := ir.add_shader_source(&frame, "parity_resource_sync.frag", .Fragment, FRAGMENT_SHADER)
	state := ir.default_graphics_pipeline_state()
	state.vertex_shader = vs
	state.fragment_shader = fs
	state.color_format = .R8G8B8A8_UNORM
	state.cull_mode = .None
	pipeline := ir.add_graphics_pipeline(&frame, "draw", state)
	_ = ir.add_draw_packet(
		&frame,
		pass,
		pipeline,
		.Vertex_Stream,
		{
			target = target,
			geometry = {kind = .Vertex_Stream, resource = vertices, vertex_count = 6},
			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_resource_sync_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 = 1,
		},
	)
	return
}

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

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

	frame, color, 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_resource_sync: failed to prepare planned frame")
	}
	sample_buffer, sample_ok := compiler.prepared_buffer(&exec, sample_geometry)
	if !sample_ok {
		fail("parity_resource_sync: failed to resolve sample geometry")
	}
	sample_vertices := quad_vertices(0.0, -1.0, 1.0, 1.0)
	if !write_bytes(backend, sample_buffer, raw_data(sample_vertices[:]), len(sample_vertices) * size_of(resource.Mesh_Vertex_PNU)) {
		fail("parity_resource_sync: failed to upload sample geometry")
	}
	if !compiler.execute_prepared_planned_frame(&exec, &frame, {0, 0, 0, 1}) {
		fail("parity_resource_sync: failed to execute planned frame")
	}
	texture, tex_ok := compiler.prepared_texture(&exec, color)
	if !tex_ok {
		fail("parity_resource_sync: 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_resource_sync: readback failed")
	}
	readback_path := os.get_env("GPU_PROOF_READBACK", context.allocator)
	if !write_ppm_rgba8(readback_path, rgba, WIDTH, HEIGHT) {
		fail("parity_resource_sync: 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)
	}
}