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)

VERTEX_SHADER :: `
@varying
struct VertexOutput
	@builtin(position) position: vec4
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),
	}
end
`

FRAGMENT_SHADER :: `
struct FragmentOutput
	color: vec4
end

@entry(fragment)
function main() -> FragmentOutput
	return FragmentOutput {
		color = vec4(240.0 / 255.0, 48.0 / 255.0, 96.0 / 255.0, 1.0),
	}
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, 0}},
		{position = {x1, y0, 0}, normal = n, tex_coord = {1, 0}},
		{position = {x1, y1, 0}, normal = n, tex_coord = {1, 1}},
		{position = {x0, y0, 0}, normal = n, tex_coord = {0, 0}},
		{position = {x1, y1, 0}, normal = n, tex_coord = {1, 1}},
		{position = {x0, y1, 0}, normal = n, tex_coord = {0, 1}},
	}
}

write_vertices :: proc(backend: ^bk.Backend, buffer: bk.Buffer_Handle, vertices: []resource.Mesh_Vertex_PNU) -> bool {
	ptr := backend.map_buffer(buffer)
	if ptr == nil {
		return false
	}
	defer backend.unmap_buffer(buffer)
	mem.copy(ptr, raw_data(vertices), len(vertices) * size_of(resource.Mesh_Vertex_PNU))
	return true
}

write_ppm_rgba8 :: proc(path: string, rgba: []u8, width, height: u32) -> bool {
	if len(path) == 0 || len(rgba) < int(width * height * 4) {
		return false
	}
	header := "P6\n256 256\n255\n"
	out := make([]u8, len(header) + int(width * height * 3))
	defer delete(out)
	for i in 0..<len(header) {
		out[i] = header[i]
	}
	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]
	}
	err := os.write_entire_file(path, out)
	return err == nil
}

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

make_frame :: proc(negative: bool) -> (
	frame: ir.Frame_IR,
	color: ir.Resource_Handle,
	geometry: ir.Resource_Handle,
	coverage: ir.Resource_Handle,
) {
	frame = ir.init_frame_ir()
	color = ir.add_texture(&frame, "proof-color", WIDTH, HEIGHT, .R8G8B8A8_UNORM, {.Color_Attachment, .Sampled}, .Transient)
	depth := ir.add_texture(&frame, "proof-depth-stencil", WIDTH, HEIGHT, .D24_UNORM_S8_UINT, {.Depth_Stencil_Attachment}, .Transient)
	geometry = ir.add_buffer(
		&frame,
		"foreground-geometry",
		u64(6 * size_of(resource.Mesh_Vertex_PNU)),
		{.Vertex},
		{.Host_Visible, .Host_Coherent},
		.Transient,
	)
	coverage = ir.add_buffer(
		&frame,
		"mask-coverage",
		u64(6 * size_of(resource.Mesh_Vertex_PNU)),
		{.Vertex},
		{.Host_Visible, .Host_Coherent},
		.Transient,
	)
	pass_desc := ir.Pass_Desc{kind = .Render, name = "stencil-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_set_depth_target(&pass_desc, ir.make_depth_target(depth, .D24_UNORM_S8_UINT, .Clear, .Store, 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 = "stencil-proof-target",
			width = WIDTH,
			height = HEIGHT,
			pass = pass,
			color_format = .R8G8B8A8_UNORM,
			depth_format = .D24_UNORM_S8_UINT,
			color_resource = color,
			depth_resource = depth,
			clear_color = {16.0 / 255.0, 16.0 / 255.0, 24.0 / 255.0, 1},
			clear_depth = 1,
		},
	)
	vs := ir.add_shader_source(&frame, "parity_stencil_clip.vert", .Vertex, VERTEX_SHADER)
	fs := ir.add_shader_source(&frame, "parity_stencil_clip.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
	state.enable_blending = false
	pipeline := ir.add_graphics_pipeline(&frame, "foreground", state)
	mask := ir.add_mask(
		&frame,
		{
			target = target,
			name = "center-mask",
			source_kind = .Geometry_Coverage,
			layout_variant = .PNU,
			coverage_resource = coverage,
			vertex_count = 6,
		},
	)
	clip := ir.INVALID_CLIP
	if !negative {
		clip = ir.add_clip(&frame, {target = target, kind = .Mask, mask = mask})
	}
	_ = ir.add_draw_packet(
		&frame,
		pass,
		pipeline,
		.Vertex_Stream,
		{
			target = target,
			clip = clip,
			geometry = {kind = .Vertex_Stream, resource = geometry, vertex_count = 6},
			order = .Strict,
		},
	)
	return
}

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

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

	mode := os.get_env("GPU_PROOF_MODE", context.allocator)
	readback_path := os.get_env("GPU_PROOF_READBACK", context.allocator)
	profile_path := os.get_env("GPU_PROFILE_JSON", context.allocator)
	negative := mode == "negative"

	frame, color, geometry, coverage := make_frame(negative)
	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_stencil_clip: failed to prepare planned frame")
	}
	foreground := quad_vertices(-1, -1, 1, 1)
	mask := quad_vertices(-0.5, -0.5, 0.5, 0.5)
	foreground_buffer, fg_ok := compiler.prepared_buffer(&exec, geometry)
	mask_buffer, mask_ok := compiler.prepared_buffer(&exec, coverage)
	if !fg_ok || !mask_ok {
		fail("parity_stencil_clip: failed to resolve prepared buffers")
	}
	if !write_vertices(backend, foreground_buffer, foreground[:]) ||
	   !write_vertices(backend, mask_buffer, mask[:]) {
		fail("parity_stencil_clip: failed to upload vertices")
	}
	if !compiler.execute_prepared_planned_frame(&exec, &frame, {0, 0, 0, 1}) {
		fail("parity_stencil_clip: failed to execute planned frame")
	}

	texture, tex_ok := compiler.prepared_texture(&exec, color)
	if !tex_ok {
		fail("parity_stencil_clip: 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_stencil_clip: readback failed")
	}
	if !write_ppm_rgba8(readback_path, rgba, WIDTH, HEIGHT) {
		fail("parity_stencil_clip: failed to write readback")
	}
	if len(profile_path) > 0 {
		_ = gpu.write_profile(profile_path, 0, 1)
	}
}