package physics

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

MAX_CONTACTS :: 1024

Matrix3 :: # row_major matrix[3, 3]f32

Sim_State :: struct {
	bodies:                                    #soa[dynamic]Body,
	suspension_constraints:                    #soa[dynamic]Suspension_Constraint,

	// Slices. When you call get_body or get_suspension_constraint you will get a pointer to an element in this slice
	bodies_slice:                              #soa[]Body,
	suspension_constraints_slice:              #soa[]Suspension_Constraint,
	first_free_body_plus_one:                  i32,
	first_free_suspension_constraint_plus_one: i32,

	// Persistent stuff for simulation
	contact_pairs:                             [MAX_CONTACTS]Contact_Pair,
	contact_pairs_len:                         int,
	convex_container:                          Convex_Container,
}

Scene :: struct {
	simulation_states:      [2]Sim_State,
	simulation_state_index: i32,
}

Body :: struct {
	// Is this body alive (if not it doesn't exist)
	alive:              bool,
	// Pos
	x:                  rl.Vector3,
	// Linear vel
	v:                  rl.Vector3,
	// Orientation
	q:                  rl.Quaternion,
	// Angular vel (omega)
	w:                  rl.Vector3,
	// Mass
	inv_mass:           f32,
	// Moment of inertia
	inv_inertia_tensor: Matrix3,
	shape:              Collision_Shape,
	// 
	next_plus_one:      i32,
}

Shape_Sphere :: struct {
	radius: f32,
}

Shape_Box :: struct {
	size: rl.Vector3,
}

Internal_Shape_Convex :: struct {
	mesh:           Convex_Handle,
	center_of_mass: rl.Vector3,
	inertia_tensor: Matrix3,
}

Shape_Convex :: struct {
	mesh:           collision.Convex,
	center_of_mass: rl.Vector3,
	inertia_tensor: Matrix3,
}

Collision_Shape :: union {
	Shape_Box,
	Internal_Shape_Convex,
}

Suspension_Constraint :: struct {
	alive:           bool,
	// Pos relative to the body
	rel_pos:         rl.Vector3,
	// Dir relative to the body
	rel_dir:         rl.Vector3,
	// Handle of the rigid body
	body:            Body_Handle,
	// Wheel radius
	radius:          f32,
	// Rest distance
	rest:            f32,
	// Inverse stiffness
	compliance:      f32,
	// How much to damp velocity of the spring
	damping:         f32,

	// Runtime state
	hit:             bool,
	hit_point:       rl.Vector3,
	// rel_hit_point = rel_pos + rel_dir * hit_t
	hit_t:           f32,
	turn_angle:      f32,
	drive_impulse:   f32,
	brake_impulse:   f32,
	applied_impulse: rl.Vector3,

	// Free list
	next_plus_one:   i32,
}

// Index plus one, so handle 0 maps to invalid body
Body_Handle :: distinct i32
Suspension_Constraint_Handle :: distinct i32

INVALID_BODY :: Body_Handle(0)
INVALID_SUSPENSION_CONSTRAINT :: Suspension_Constraint_Handle(0)

is_body_handle_valid :: proc(handle: Body_Handle) -> bool {
	return i32(handle) > 0
}
is_suspension_constraint_handle_valid :: proc(handle: Suspension_Constraint_Handle) -> bool {
	return i32(handle) > 0
}
is_handle_valid :: proc {
	is_body_handle_valid,
	is_suspension_constraint_handle_valid,
}

Body_Ptr :: #soa^#soa[]Body
Suspension_Constraint_Ptr :: #soa^#soa[]Suspension_Constraint

_invalid_body: #soa[1]Body
_invalid_body_slice := _invalid_body[:]

_invalid_suspension_constraint: #soa[1]Suspension_Constraint
_invalid_suspension_constraint_slice := _invalid_suspension_constraint[:]

get_sim_state :: proc(scene: ^Scene) -> ^Sim_State {
	return &scene.simulation_states[scene.simulation_state_index]
}

get_prev_sim_state :: proc(scene: ^Scene) -> ^Sim_State {
	return &scene.simulation_states[(scene.simulation_state_index + 1) % 2]
}

// lol
get_next_sim_state :: get_prev_sim_state

flip_sim_state :: proc(scene: ^Scene) {
	scene.simulation_state_index = (scene.simulation_state_index + 1) % 2
}

/// Returns pointer to soa slice. NEVER STORE IT
get_body :: proc(sim_state: ^Sim_State, handle: Body_Handle) -> Body_Ptr {
	index := int(handle) - 1
	if index < 0 || index >= len(sim_state.bodies_slice) {
		return &_invalid_body_slice[0]
	}

	return &sim_state.bodies_slice[index]
}

remove_shape :: proc(sim_state: ^Sim_State, shape: Collision_Shape) {
	switch &s in shape {
	case Shape_Box:
	case Internal_Shape_Convex:
		convex_container_remove(&sim_state.convex_container, s.mesh)
	}
}

Input_Shape :: union {
	Shape_Box,
	Shape_Convex,
}

Body_Config :: struct {
	initial_pos:     rl.Vector3,
	initial_rot:     rl.Quaternion,
	initial_vel:     rl.Vector3,
	initial_ang_vel: rl.Vector3,
	shape:           Input_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
}

add_shape :: proc(sim_state: ^Sim_State, shape: Input_Shape) -> (result: Collision_Shape) {
	switch s in shape {
	case Shape_Box:
		result = s
	case Shape_Convex:
		convex: Internal_Shape_Convex
		convex.mesh = convex_container_add(&sim_state.convex_container, s.mesh)
		convex.center_of_mass = s.center_of_mass
		convex.inertia_tensor = s.inertia_tensor
		result = convex
	}

	return
}

initialize_body_from_config :: proc(sim_state: ^Sim_State, 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 = add_shape(sim_state, config.shape)

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

update_body_from_config :: proc(sim_state: ^Sim_State, body: Body_Ptr, config: Body_Config) {
	// Inefficient, but eh
	remove_shape(sim_state, body.shape)
	body.shape = add_shape(sim_state, 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
}


add_body :: proc(sim_state: ^Sim_State, config: Body_Config) -> Body_Handle {
	body: Body

	initialize_body_from_config(sim_state, &body, config)

	body.alive = true

	if sim_state.first_free_body_plus_one > 0 {
		index := sim_state.first_free_body_plus_one
		new_body := get_body(sim_state, Body_Handle(index))
		next_plus_one := new_body.next_plus_one
		new_body^ = body
		sim_state.first_free_body_plus_one = next_plus_one

		return Body_Handle(index)
	}

	append_soa(&sim_state.bodies, body)
	index := len(sim_state.bodies)

	sim_state.bodies_slice = sim_state.bodies[:]

	return Body_Handle(index)
}

remove_body :: proc(sim_state: ^Sim_State, handle: Body_Handle) {
	if int(handle) > 1 {
		body := get_body(sim_state, handle)

		body.alive = false

		remove_shape(sim_state, body.shape)

		body.next_plus_one = sim_state.first_free_body_plus_one
		sim_state.first_free_body_plus_one = i32(handle)
	}
}

/// Returns pointer to soa slice. NEVER STORE IT
get_suspension_constraint :: proc(
	sim_state: ^Sim_State,
	handle: Suspension_Constraint_Handle,
) -> Suspension_Constraint_Ptr {
	if !is_handle_valid(handle) {
		return &_invalid_suspension_constraint_slice[0]
	}

	index := int(handle) - 1
	return &sim_state.suspension_constraints_slice[index]
}

add_suspension_constraint :: proc(
	sim_state: ^Sim_State,
	constraint: Suspension_Constraint,
) -> Suspension_Constraint_Handle {
	copy := constraint

	copy.alive = true
	copy.next_plus_one = 0

	if sim_state.first_free_suspension_constraint_plus_one > 0 {
		index := sim_state.first_free_suspension_constraint_plus_one
		new_constraint := get_suspension_constraint(sim_state, Suspension_Constraint_Handle(index))
		next_plus_one := new_constraint.next_plus_one
		new_constraint^ = copy
		sim_state.first_free_suspension_constraint_plus_one = next_plus_one
		return Suspension_Constraint_Handle(index)
	}

	append_soa(&sim_state.suspension_constraints, copy)
	sim_state.suspension_constraints_slice = sim_state.suspension_constraints[:]
	index := len(sim_state.suspension_constraints)
	return Suspension_Constraint_Handle(index)
}

remove_suspension_constraint :: proc(sim_state: ^Sim_State, handle: Suspension_Constraint_Handle) {
	if is_handle_valid(handle) {
		constraint := get_suspension_constraint(sim_state, handle)

		constraint.alive = false
		constraint.next_plus_one = sim_state.first_free_suspension_constraint_plus_one
		sim_state.first_free_suspension_constraint_plus_one = i32(handle)
	}
}

_get_first_free_body :: proc(sim_state: ^Sim_State) -> i32 {
	return sim_state.first_free_body_plus_one - 1
}

destry_sim_state :: proc(sim_state: ^Sim_State) {
	delete_soa(sim_state.bodies)
	delete_soa(sim_state.suspension_constraints)
	convex_container_destroy(&sim_state.convex_container)
}

destroy_physics_scene :: proc(scene: ^Scene) {
	for &sim_state in scene.simulation_states {
		destry_sim_state(&sim_state)
	}
}