package bvh

import "../collision"
import "base:runtime"
import "core:container/queue"
import "core:log"
import "core:math"
import lg "core:math/linalg"
import "core:mem"
import "game:debug"
import "libs:tracy"
import rl "vendor:raylib"

_ :: log
_ :: rl
_ :: lg

Vec3 :: [3]f32

AABB :: struct {
	min, max: Vec3,
}

// Helper struct to avoid passing verts/indices separately
Mesh :: struct {
	vertices: []Vec3,
	indices:  []u16,
}

BVH :: struct {
	nodes:      []Node,
	// Triangle IDs. first_index = indices[primitive * 3]
	primitives: []u16,
	nodes_used: i32,
}

destroy_bvh :: proc(bvh: ^BVH) {
	delete(bvh.nodes)
	delete(bvh.primitives)
}

// Helper struct to store mesh data together with its bvh for convenience
// You don't have to use it
Mesh_BVH :: struct {
	bvh:  BVH,
	mesh: Mesh,
}

Node :: struct {
	aabb:                AABB,
	// Index of the left child, right child is left_child + 1
	child_or_prim_start: i32,
	prim_len:            i32,
}

// uvw
Bary :: [3]f32

is_leaf_node :: #force_inline proc(node: Node) -> bool {
	return node.prim_len > 0
}

#assert(size_of(Node) == 32)

build_bvh_from_mesh :: proc(mesh: Mesh, allocator := context.allocator) -> (mesh_bvh: Mesh_BVH) {
	tracy.Zone()

	vertices, indices := mesh.vertices, mesh.indices
	assert(len(indices) % 3 == 0)

	bvh := &mesh_bvh.bvh

	num_triangles := len(indices) / 3

	// Caller owned, allocator might be temp_allocator so do this before checkpoint below, otherwise we the result accidentally
	bvh.nodes, _ = mem.make_aligned([]Node, num_triangles * 2 - 1, size_of(Node), allocator)
	bvh.primitives = make([]u16, num_triangles, allocator)

	// Clean up after ourselves
	temp := runtime.default_temp_allocator_temp_begin()
	defer runtime.default_temp_allocator_temp_end(temp)

	// Temp stuff
	centroids := make([]Vec3, num_triangles, context.temp_allocator)
	aabbs := make([]AABB, num_triangles, context.temp_allocator)

	// Calculate centroids and aabbs
	{
		tracy.ZoneN("calculate_centroids_and_aabbs")

		for i in 0 ..< num_triangles {
			i1, i2, i3 := indices[i * 3], indices[i * 3 + 1], indices[i * 3 + 2]
			v1, v2, v3 := vertices[i1], vertices[i2], vertices[i3]

			centroids[i] = (v1 + v2 + v3) * 0.33333333333

			aabbs[i].min = Vec3 {
				min(v1.x, v2.x, v3.x),
				min(v1.y, v2.y, v3.y),
				min(v1.z, v2.z, v3.z),
			}
			aabbs[i].max = Vec3 {
				max(v1.x, v2.x, v3.x),
				max(v1.y, v2.y, v3.y),
				max(v1.z, v2.z, v3.z),
			}

			size := aabbs[i].max - aabbs[i].min
			assert(size.x >= 0)
			assert(size.y >= 0)
			assert(size.z >= 0)

			bvh.primitives[i] = u16(i)
		}
	}

	bvh.nodes_used = 1 // root

	root := &bvh.nodes[0]
	root.child_or_prim_start = 0
	root.prim_len = i32(num_triangles)

	update_node_bounds(bvh, 0, aabbs)
	subdivide(bvh, 0, centroids, aabbs)

	return
}

/// Useful for a top level accel structure
build_bvh_from_aabbs :: proc(aabbs: []AABB, allocator := context.allocator) -> (bvh: BVH) {
	bvh.nodes, _ = mem.make_aligned([]Node, len(aabbs) * 2 - 1, size_of(Node), allocator)
	bvh.primitives = make([]u16, len(aabbs), allocator)

	temp := runtime.default_temp_allocator_temp_begin()
	defer runtime.default_temp_allocator_temp_end(temp)

	// Temp stuff
	centroids := make([]Vec3, len(aabbs), context.temp_allocator)

	// Calculate centroids
	for i in 0 ..< len(aabbs) {
		centroids[i] = (aabbs[i].max + aabbs[i].min) * 0.5

		bvh.primitives[i] = u16(i)
	}

	bvh.nodes_used = 1

	root := &bvh.nodes[0]
	root.prim_len = i32(len(aabbs))

	update_node_bounds(&bvh, 0, aabbs)

	subdivide(&bvh, 0, centroids, aabbs)

	return
}

update_node_bounds :: proc(bvh: ^BVH, node_idx: i32, prim_aabbs: []AABB) {
	tracy.Zone()

	node := &bvh.nodes[node_idx]

	node.aabb.min = math.F32_MAX
	node.aabb.max = -math.F32_MAX

	for i in node.child_or_prim_start ..< node.child_or_prim_start + node.prim_len {
		prim_aabb := prim_aabbs[bvh.primitives[i]]

		node.aabb.min.x = min(node.aabb.min.x, prim_aabb.min.x)
		node.aabb.min.y = min(node.aabb.min.y, prim_aabb.min.y)
		node.aabb.min.z = min(node.aabb.min.z, prim_aabb.min.z)

		node.aabb.max.x = max(node.aabb.max.x, prim_aabb.max.x)
		node.aabb.max.y = max(node.aabb.max.y, prim_aabb.max.y)
		node.aabb.max.z = max(node.aabb.max.z, prim_aabb.max.z)
	}
}

subdivide :: proc(bvh: ^BVH, node_idx: i32, centroids: []Vec3, aabbs: []AABB) {
	tracy.Zone()

	node := &bvh.nodes[node_idx]

	if node.prim_len <= 2 {
		return
	}

	size := node.aabb.max - node.aabb.min

	// Split along longest axis
	largest_side := size.x
	split_axis := 0
	if size.y > largest_side {
		split_axis = 1
		largest_side = size.y
	}
	if size.z > largest_side {
		split_axis = 2
	}

	split_pos := node.aabb.min[split_axis] + size[split_axis] * 0.5

	// Partition
	i := node.child_or_prim_start
	j := i + node.prim_len - 1

	for i <= j {
		prim_i := bvh.primitives[i]
		prim_j := bvh.primitives[j]
		if centroids[prim_i][split_axis] < split_pos {
			i += 1
		} else {
			bvh.primitives[i] = prim_j
			bvh.primitives[j] = prim_i
			j -= 1
		}
	}

	left_count := i - node.child_or_prim_start
	if left_count == 0 || left_count == node.prim_len {
		return
	}

	left_child := bvh.nodes_used
	right_child := bvh.nodes_used + 1
	bvh.nodes_used += 2

	prim_start := node.child_or_prim_start
	node.child_or_prim_start = left_child

	bvh.nodes[left_child] = {}
	bvh.nodes[right_child] = {}

	bvh.nodes[left_child].child_or_prim_start = prim_start
	bvh.nodes[left_child].prim_len = left_count
	bvh.nodes[right_child].child_or_prim_start = i
	bvh.nodes[right_child].prim_len = node.prim_len - left_count
	node.prim_len = 0

	update_node_bounds(bvh, left_child, aabbs)
	update_node_bounds(bvh, right_child, aabbs)

	subdivide(bvh, left_child, centroids, aabbs)
	subdivide(bvh, right_child, centroids, aabbs)
}

Ray :: struct {
	origin, dir: Vec3,
	dir_inv:     Vec3,
}

Collision :: struct {
	hit:            bool,
	t:              f32,
	// which primitive we hit
	prim:           u16,
	// Barycentric coords of the hit triangle
	bary:           Bary,
	// jStats
	aabb_tests:     int,
	triangle_tests: int,
}

traverse_bvh_ray_mesh :: proc(bvh: ^BVH, mesh: Mesh, ray: Ray, out_collision: ^Collision) -> bool {
	tracy.Zone()

	ray := ray
	ray.dir_inv.x = 1.0 / ray.dir.x
	ray.dir_inv.y = 1.0 / ray.dir.y
	ray.dir_inv.z = 1.0 / ray.dir.z

	if !out_collision.hit {
		out_collision.t = math.F32_MAX
	}

	prev_t := out_collision.t

	internal_traverse_bvh_ray_triangles(bvh, mesh, ray, out_collision)

	return out_collision.hit && out_collision.t < prev_t
}

internal_traverse_bvh_ray_triangles :: proc(
	bvh: ^BVH,
	mesh: Mesh,
	ray: Ray,
	out_collision: ^Collision,
) {
	temp := runtime.default_temp_allocator_temp_begin()
	defer runtime.default_temp_allocator_temp_end(temp)

	nodes_to_process: queue.Queue(i32)
	queue.init(&nodes_to_process, queue.DEFAULT_CAPACITY, context.temp_allocator)
	queue.push_back(&nodes_to_process, 0)

	for queue.len(nodes_to_process) > 0 {
		node_idx := queue.pop_front(&nodes_to_process)
		assert(node_idx < bvh.nodes_used)

		node := &bvh.nodes[node_idx]

		out_collision.aabb_tests += 1
		if !internal_ray_aabb_test(ray, node.aabb, out_collision.t) {
			continue
		}

		rl.DrawBoundingBox(
			{min = node.aabb.min, max = node.aabb.max},
			debug.int_to_color(node_idx),
		)

		if is_leaf_node(node^) {
			for i in node.child_or_prim_start ..< node.child_or_prim_start + node.prim_len {
				internal_ray_tri_test(ray, mesh, bvh.primitives[i], out_collision)
			}
		} else {
			left_node := node.child_or_prim_start

			queue.push_back_elems(&nodes_to_process, left_node, left_node + 1)
		}
	}
}

internal_ray_aabb_test :: proc(ray: Ray, box: AABB, min_t: f32) -> bool {
	ts, ok := collision.intersect_ray_aabb(ray.origin, ray.dir, collision.Aabb{box.min, box.max})
	return ok && ts[0] < min_t
}

// Möller–Trumbore intersection algorithm
// https://jacco.ompf2.com/2022/04/13/how-to-build-a-bvh-part-1-basics/
internal_ray_tri_test :: proc(ray: Ray, mesh: Mesh, tri: u16, col: ^Collision) {
	i1, i2, i3 := mesh.indices[tri * 3], mesh.indices[tri * 3 + 1], mesh.indices[tri * 3 + 2]
	v1, v2, v3 := mesh.vertices[i1], mesh.vertices[i2], mesh.vertices[i3]

	col.triangle_tests += 1
	// rl.DrawTriangle3D(v1, v2, v3, debug.int_to_color(i32(tri) + 1))

	t, _, barycentric, ok := collision.intersect_ray_triangle(
		{ray.origin, ray.origin + ray.dir},
		{v1, v2, v3},
	)

	if ok && t < col.t {
		col.hit = true
		col.t = t
		col.prim = tri
		col.bary = barycentric
	}
}