package physics

import "core:log"
import lg "core:math/linalg"
import rl "vendor:raylib"

_ :: log

Body_Config_Inertia_Mode :: enum {
	From_Shape,
	Explicit,
}

// Immediate mode stuff for testing
Body_Config :: struct {
	initial_pos:     rl.Vector3,
	initial_rot:     rl.Quaternion,
	initial_vel:     rl.Vector3,
	initial_ang_vel: rl.Vector3,
	shape:           Collision_Shape,
	mass:            f32,
	inertia_mode:    Body_Config_Inertia_Mode,
	// Unit inertia tensor
	inertia_tensor:  Matrix3,
}

// TODO: rename to wheel
Suspension_Constraint_Config :: struct {
	rel_pos:    rl.Vector3,
	rel_dir:    rl.Vector3,
	body:       Body_Handle,
	rest:       f32,
	compliance: f32,
	damping:    f32,
	radius:     f32,
}

calculate_body_params_from_config :: proc(
	config: Body_Config,
) -> (
	inv_mass: f32,
	inv_inertia_tensor: Matrix3,
) {
	inv_mass = config.mass == 0 ? 0 : 1.0 / config.mass

	inertia_tensor: Matrix3
	if config.inertia_mode == .Explicit {
		inertia_tensor = config.inertia_tensor
	} else {
		inertia_tensor = inertia_tensor_collision_shape(config.shape)
	}
	inertia_tensor = inertia_tensor * Matrix3(config.mass)

	inv_inertia_tensor = lg.determinant(inertia_tensor) == 0 ? 0 : lg.inverse(inertia_tensor)

	return
}

initialize_body_from_config :: proc(body: ^Body, config: Body_Config) {
	body.x = config.initial_pos
	body.q = config.initial_rot
	body.v = config.initial_vel
	body.w = config.initial_ang_vel
	body.shape = config.shape

	body.inv_mass, body.inv_inertia_tensor = calculate_body_params_from_config(config)
}

update_body_from_config :: proc(body: Body_Ptr, config: Body_Config) {
	body.shape = config.shape
	body.inv_mass, body.inv_inertia_tensor = calculate_body_params_from_config(config)
}

update_suspension_constraint_from_config :: proc(
	constraint: Suspension_Constraint_Ptr,
	config: Suspension_Constraint_Config,
) {
	constraint.rel_pos = config.rel_pos
	constraint.rel_dir = config.rel_dir
	constraint.body = config.body
	constraint.rest = config.rest
	constraint.compliance = config.compliance
	constraint.damping = config.damping
	constraint.radius = config.radius
}

immediate_body :: proc(
	scene: ^Scene,
	state: ^Solver_State,
	id: u32,
	config: Body_Config,
) -> (
	handle: Body_Handle,
) {
	if id in state.immedate_bodies {
		body := &state.immedate_bodies[id]
		if body.last_ref != state.simulation_frame {
			body.last_ref = state.simulation_frame
			state.num_referenced_bodies += 1
		}
		handle = body.handle
		update_body_from_config(get_body(scene, handle), config)
	} else {
		new_body: Body
		state.num_referenced_bodies += 1
		initialize_body_from_config(&new_body, config)
		handle = add_body(scene, new_body)
		state.immedate_bodies[id] = {
			handle   = handle,
			last_ref = state.simulation_frame,
		}
	}

	return
}

immediate_suspension_constraint :: proc(
	scene: ^Scene,
	state: ^Solver_State,
	id: u32,
	config: Suspension_Constraint_Config,
) -> (
	handle: Suspension_Constraint_Handle,
) {
	if id in state.immediate_suspension_constraints {
		constraint := &state.immediate_suspension_constraints[id]
		if constraint.last_ref != state.simulation_frame {
			constraint.last_ref = state.simulation_frame
			state.num_referenced_suspension_constraints += 1
		}
		handle = constraint.handle
	} else {
		state.num_referenced_suspension_constraints += 1
		handle = add_suspension_constraint(scene, {})
		state.immediate_suspension_constraints[id] = {
			handle   = handle,
			last_ref = state.simulation_frame,
		}
	}

	update_suspension_constraint_from_config(get_suspension_constraint(scene, handle), config)

	return
}

prune_immediate :: proc(scene: ^Scene, state: ^Solver_State) {
	prune_immediate_bodies(scene, state)
	prune_immediate_suspension_constraints(scene, state)
}

// TODO: Generic version
prune_immediate_bodies :: proc(scene: ^Scene, state: ^Solver_State) {
	if int(state.num_referenced_bodies) == len(state.immedate_bodies) {
		return
	}

	num_unreferenced_bodies := len(state.immedate_bodies) - int(state.num_referenced_bodies)
	assert(num_unreferenced_bodies >= 0)

	bodies_to_remove := make([]u32, num_unreferenced_bodies, context.temp_allocator)

	i := 0
	for k, &v in state.immedate_bodies {
		if v.last_ref != state.simulation_frame {
			bodies_to_remove[i] = k
			i += 1
		}
	}

	assert(i == len(bodies_to_remove))

	for k in bodies_to_remove {
		handle := state.immedate_bodies[k].handle
		delete_key(&state.immedate_bodies, k)
		remove_body(scene, handle)
	}
}

prune_immediate_suspension_constraints :: proc(scene: ^Scene, state: ^Solver_State) {
	if int(state.num_referenced_suspension_constraints) ==
	   len(state.immediate_suspension_constraints) {
		return
	}

	num_unreferenced_constraints :=
		len(state.immediate_suspension_constraints) -
		int(state.num_referenced_suspension_constraints)
	assert(num_unreferenced_constraints >= 0)

	constraints_to_remove := make([]u32, num_unreferenced_constraints, context.temp_allocator)

	i := 0
	for k, &v in state.immediate_suspension_constraints {
		if v.last_ref != state.simulation_frame {
			constraints_to_remove[i] = k
			i += 1
		}
	}

	assert(i == len(constraints_to_remove))

	for k in constraints_to_remove {
		handle := state.immediate_suspension_constraints[k].handle
		delete_key(&state.immediate_suspension_constraints, k)
		remove_suspension_constraint(scene, handle)
	}
}