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 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) {
	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
	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) {
	node := &bvh.nodes[node_idx]

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

	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)
	}

	size := node.aabb.max - node.aabb.min
	assert(size.x >= 0)
	assert(size.y >= 0)
	assert(size.z >= 0)
}

subdivide :: proc(bvh: ^BVH, node_idx: i32, centroids: []Vec3, aabbs: []AABB) {
	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,
}

traverse_bvh_ray_mesh :: proc(bvh: ^BVH, mesh: Mesh, ray: Ray, out_collision: ^Collision) -> bool {
	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]

		if !internal_ray_aabb_test(ray, node.aabb, out_collision.t) {
			return
		}

		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)
		}
	}
}

// https://tavianator.com/2022/ray_box_boundary.html
internal_ray_aabb_test :: proc(ray: Ray, box: AABB, min_t: f32) -> bool {
	_, ok := collision.intersect_ray_aabb(ray.origin, ray.dir, collision.Aabb{box.min, box.max})
	return ok

	// t1 := (box.min[0] - ray.origin[0]) * ray.dir_inv[0]
	// t2 := (box.max[0] - ray.origin[0]) * ray.dir_inv[0]

	// tmin := min(t1, t2)
	// tmax := max(t1, t2)

	// for i in 1 ..< 3 {
	// 	t1 = (box.min[i] - ray.origin[i]) * ray.dir_inv[i]
	// 	t2 = (box.max[i] - ray.origin[i]) * ray.dir_inv[i]

	// 	tmin = max(tmin, min(t1, t2))
	// 	tmax = min(tmax, max(t1, t2))
	// }

	// return tmax > max(tmin, 0.0)
}

// 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]

	t, _, barycentric, ok := collision.intersect_segment_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
	}
	// rl.DrawTriangle3D(v1, v2, v3, debug_int_to_color(i32(tri)))

	// rl_col := rl.GetRayCollisionTriangle(rl.Ray{ray.origin, ray.dir}, v1, v2, v3)

	// if rl_col.hit && rl_col.distance < col.t {
	// 	col.hit = true
	// 	col.t = lg.distance(ray.origin, rl_col.point)
	// }

	return

	//e1, e2 := v2 - v1, v3 - v1

	//h := lg.cross(ray.dir, e2)
	//a := lg.dot(e1, h)

	//// ray parallel to triangle
	//if a > -0.0001 || a < 0.0001 {
	//	return
	//}

	//f: f32 = 1.0 / a
	//s := ray.origin - v1
	//u := f * lg.dot(s, h)
	//if u < 0 || u > 1 {
	//	return
	//}
	//q := lg.cross(s, e1)
	//v := f * lg.dot(ray.dir, q)
	//if v < 0 || u + v > 1 {
	//	return
	//}

	//t := f * lg.dot(e2, q)
	//if t > 0.0001 && t < col.t {
	//	col.hit = true
	//	col.t = t
	//	col.prim = tri
	//	col.bary = Vec3{u, v, 0} // TODO: calc W
	//}
}