package main

import "core:fmt"
import bk "../../backend"
import compiler "../../compiler"
import gpu "../.."
import app "../../app"
import ir "../../render_ir"
import shaderkit "../../shaderkit"

Vertex :: struct {
	position: [2]f32,
	uv:       [2]f32,
}

VERTEX_SHADER :: `
vertex main
  in position: vec2 @location(0)
  in uv: vec2 @location(1)
  out position: vec4 @builtin(position)
  out uv: vec2
do
  position = vec4(position, 0.0, 1.0)
  out.uv = uv
end
`

OFFSCREEN_FRAGMENT_SHADER :: `
fragment main
  in uv: vec2
  out color: vec4
do
  let band = step(0.5, fract((uv.x + uv.y) * 8.0))
  color = mix(vec4(0.05, 0.15, 0.65, 1.0), vec4(uv.x, uv.y, 0.25, 1.0), band)
end
`

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

fragment main
  in uv: vec2
  out color: vec4
do
  color = sample(offscreen_color, uv)
end
`

compile_shader :: proc(source, name: string, stage: bk.Shader_Stage) -> (bk.Shader_Handle, bool) {
	return shaderkit.create_shader_from_source(name, stage, source, {opt_level = .Basic})
}

make_offscreen_ir :: proc() -> (frame: ir.Frame_IR, pass: ir.Pass_Handle, color: ir.Resource_Handle, sampler: ir.Resource_Handle) {
	frame = ir.init_frame_ir()
	color = ir.add_texture(&frame, "offscreen-color", 256, 256, .R8G8B8A8_UNORM, {.Color_Attachment, .Sampled}, .Transient)
	depth := ir.add_texture(&frame, "offscreen-depth", 256, 256, .D32_SFLOAT, {.Depth_Stencil_Attachment}, .Transient)
	sampler = ir.add_sampler(&frame, "offscreen-sampler", {
		mag_filter = .Nearest,
		min_filter = .Nearest,
		address_mode_u = .Clamp_To_Edge,
		address_mode_v = .Clamp_To_Edge,
		mipmap_mode = .Nearest,
	}, .Transient)

	desc := ir.Pass_Desc{kind = .Render, name = "offscreen"}
	_ = ir.pass_add_color_target(&desc, ir.make_color_target(color, .R8G8B8A8_UNORM, .Clear, .Store, {0, 0, 0, 1}))
	ir.pass_set_depth_target(&desc, ir.make_depth_target(depth, .D32_SFLOAT, .Clear, .Dont_Care, 1))
	ir.pass_set_viewport(&desc, {x = 0, y = 0, width = 256, height = 256, min_depth = 0, max_depth = 1})
	ir.pass_set_scissor(&desc, {x = 0, y = 0, width = 256, height = 256})
	pass = ir.add_pass(&frame, desc)
	return
}

create_quad_buffers :: proc(backend: ^bk.Backend) -> (bk.Buffer_Handle, bk.Buffer_Handle, bool) {
	vertices := [?]Vertex{
		{position = {-0.85, -0.85}, uv = {0, 1}},
		{position = { 0.85, -0.85}, uv = {1, 1}},
		{position = { 0.85,  0.85}, uv = {1, 0}},
		{position = {-0.85,  0.85}, uv = {0, 0}},
	}
	indices := [?]u32{0, 1, 2, 2, 3, 0}

	vb, vb_ok := backend.create_buffer_staged(raw_data(vertices[:]), len(vertices) * size_of(Vertex), {.Vertex})
	if !vb_ok {
		return {}, {}, false
	}
	ib, ib_ok := backend.create_buffer_staged(raw_data(indices[:]), len(indices) * size_of(u32), {.Index})
	if !ib_ok {
		backend.destroy_buffer(vb)
		return {}, {}, false
	}
	return vb, ib, true
}

create_pipeline :: proc(
	vs, fs: bk.Shader_Handle,
	render_pass: bk.Render_Pass_Handle,
	descriptor_layouts: []bk.Descriptor_Handle = nil,
) -> (bk.Pipeline_Handle, bool) {
	bindings := [1]bk.Vertex_Binding{{binding = 0, stride = size_of(Vertex)}}
	attrs := [2]bk.Vertex_Attribute{
		{location = 0, offset = 0, format = .Float2},
		{location = 1, offset = 8, format = .Float2},
	}
	return gpu.create_pipeline({
		vert_shader = vs,
		frag_shader = fs,
		render_pass = render_pass,
		topology = .Triangle_List,
		cull_mode = .None,
		front_face = .Counter_Clockwise,
		enable_blending = false,
		blend_mode = .Alpha,
		enable_depth_test = false,
		depth_only = false,
		descriptor_layouts = descriptor_layouts,
		vertex_bindings = bindings[:],
		vertex_attributes = attrs[:],
	})
}

main :: proc() {
	when gpu.REQUIRED_SHADER_FORMAT != .SPIRV {
		fmt.println("offscreen_ir: Vulkan runtime proof only; D3D executor parity is tracked separately")
		return
	}

	state, ok := app.init_window(800, 600, "GPU - Offscreen IR")
	if !ok {
		return
	}
	defer app.shutdown(&state)

	backend := gpu.get_backend()
	if backend == nil {
		return
	}

	frame, offscreen_pass, color_resource, sampler_resource := make_offscreen_ir()
	defer ir.destroy_frame_ir(&frame)

	exec := compiler.init_execution_state(backend)
	defer compiler.destroy_execution_state(&exec)
	if !compiler.prepare_render_targets(&exec, &frame) {
		return
	}

	offscreen_render_pass, rp_ok := compiler.prepared_render_pass(&exec, offscreen_pass)
	if !rp_ok {
		return
	}
	offscreen_texture, tex_ok := compiler.prepared_texture(&exec, color_resource)
	offscreen_sampler, sampler_ok := compiler.prepared_sampler(&exec, sampler_resource)
	if !tex_ok || !sampler_ok {
		return
	}

	vs, vs_ok := compile_shader(VERTEX_SHADER, "offscreen_ir.vert", .Vertex)
	if !vs_ok { return }
	defer gpu.destroy_shader(vs)
	fs_offscreen, fs_offscreen_ok := compile_shader(OFFSCREEN_FRAGMENT_SHADER, "offscreen_ir_offscreen.frag", .Fragment)
	if !fs_offscreen_ok { return }
	defer gpu.destroy_shader(fs_offscreen)
	fs_sample, fs_sample_ok := compile_shader(SAMPLE_FRAGMENT_SHADER, "offscreen_ir_sample.frag", .Fragment)
	if !fs_sample_ok { return }
	defer gpu.destroy_shader(fs_sample)

	offscreen_pipeline, offscreen_pipeline_ok := create_pipeline(vs, fs_offscreen, offscreen_render_pass)
	if !offscreen_pipeline_ok { return }
	defer gpu.destroy_pipeline(offscreen_pipeline)

	descriptor_layout, descriptor_layout_ok := backend.create_descriptor_set_layout({{
		binding = 0,
		type = .Combined_Image_Sampler,
		count = 1,
		stages = {.Fragment},
	}})
	if !descriptor_layout_ok { return }
	defer backend.destroy_descriptor_set_layout(descriptor_layout)

	descriptor_types := [1]bk.Descriptor_Type{.Combined_Image_Sampler}
	descriptor_counts := [1]u32{1}
	descriptor_pool, descriptor_pool_ok := backend.create_descriptor_pool(1, descriptor_types[:], descriptor_counts[:])
	if !descriptor_pool_ok { return }
	defer backend.destroy_descriptor_pool(descriptor_pool)

	descriptor_set, descriptor_set_ok := backend.allocate_descriptor_set(descriptor_pool, descriptor_layout)
	if !descriptor_set_ok { return }
	backend.update_descriptor_image(descriptor_set, 0, offscreen_texture, offscreen_sampler, .Shader_Read_Only)

	descriptor_layouts := [1]bk.Descriptor_Handle{descriptor_layout}
	sample_pipeline, sample_pipeline_ok := create_pipeline(vs, fs_sample, gpu.get_default_render_pass(), descriptor_layouts[:])
	if !sample_pipeline_ok { return }
	defer gpu.destroy_pipeline(sample_pipeline)

	vertex_buffer, index_buffer, buffers_ok := create_quad_buffers(backend)
	if !buffers_ok { return }
	defer backend.destroy_buffer(index_buffer)
	defer backend.destroy_buffer(vertex_buffer)

	wrote_offscreen := false
	for !app.window_should_close(&state) {
		app.begin_drawing(&state, gpu.BLACK)
		ctx := gpu.get_frame_context()

		if !wrote_offscreen {
			if compiler.begin_prepared_pass(&exec, ctx, offscreen_pass) {
				backend.bind_graphics_pipeline(ctx, offscreen_pipeline)
				backend.bind_vertex_buffer(ctx, vertex_buffer)
				backend.bind_index_buffer(ctx, index_buffer)
				backend.draw_indexed(ctx, 6, 1, 0, 0, 0)
				compiler.end_prepared_pass(&exec, ctx)
				wrote_offscreen = true
			}
		}

		if gpu.ensure_default_render_pass() {
			backend.bind_graphics_pipeline(ctx, sample_pipeline)
			backend.bind_descriptor_set(ctx, sample_pipeline, descriptor_set, 0)
			backend.bind_vertex_buffer(ctx, vertex_buffer)
			backend.bind_index_buffer(ctx, index_buffer)
			backend.draw_indexed(ctx, 6, 1, 0, 0, 0)
		}

		app.end_drawing(&state)
		if app.is_key_pressed(&state, .Escape) {
			break
		}
	}
}