package physics

import "collision"
import lg "core:math/linalg"
import "game:container/spanpool"

MAX_CONTACTS :: 1024 * 16

Vec3 :: [3]f32
Vec2 :: [2]f32
Quat :: quaternion128
Matrix3 :: # row_major matrix[3, 3]f32
AABB :: struct {
	center: Vec3,
	extent: Vec3,
}

Contact_Pair :: [2]i32

make_contact_pair :: proc(body_a: i32, body_b: i32) -> Contact_Pair {
	return {min(body_a, body_b), max(body_a, body_b)}
}

Contact_Container :: struct {
	// body index pair to contact index
	lookup:   map[Contact_Pair]i32,
	contacts: #soa[dynamic]Contact,
}

contact_container_copy :: proc(dst: ^Contact_Container, src: Contact_Container) {
	clear(&dst.lookup)
	reserve(&dst.lookup, cap(src.lookup))
	resize_soa(&dst.contacts, len(src.contacts))

	for k, v in src.lookup {
		dst.lookup[k] = v
	}

	for i in 0 ..< len(src.contacts) {
		dst.contacts[i] = src.contacts[i]
	}
}

Sim_State :: struct {
	bodies:                                    #soa[dynamic]Body,
	suspension_constraints:                    #soa[dynamic]Suspension_Constraint,
	engines:                                   [dynamic]Engine,
	// Number of alive bodies
	num_bodies:                                i32,
	num_engines:                               i32,

	// 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,
	first_free_engine_plus_one:                i32,

	// Persistent stuff for simulation
	contact_container:                         Contact_Container,
	convex_container:                          Convex_Container,

	// NOTE: kinda overkill, but it simplifies copying sim states around a lot
	// Engine array data
	rpm_torque_curves_pool:                    spanpool.Span_Pool([2]f32),
	gear_ratios_pool:                          spanpool.Span_Pool(f32),
}

Scene :: struct {
	simulation_states:      []Sim_State,
	// Speculative prediction state to find collisions
	scratch_sim_state:      Sim_State,
	simulation_state_index: i32,
}

init_physics_scene :: proc(scene: ^Scene, max_history := int(2)) {
	scene.simulation_states = make([]Sim_State, max(max_history, 2))
}

Body :: struct {
	// Is this body alive (if not it doesn't exist)
	alive:              bool,
	// Pos
	x:                  Vec3,
	// Linear vel
	v:                  Vec3,
	// Orientation
	q:                  Quat,
	// Angular vel (omega)
	w:                  Vec3,
	prev_x:             Vec3,
	prev_v:             Vec3,
	prev_q:             Quat,
	prev_w:             Vec3,
	// 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: Vec3,
}

Internal_Shape_Convex :: struct {
	mesh:           Convex_Handle,
	center_of_mass: Vec3,
	inertia_tensor: Matrix3,
	center:         Vec3,
	extent:         Vec3,
}

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

Collision_Shape :: union {
	Shape_Box,
	Internal_Shape_Convex,
}

Suspension_Constraint :: struct {
	alive:                  bool,
	// Pos relative to the body
	rel_pos:                Vec3,
	// Dir relative to the body
	rel_dir:                Vec3,
	// Handle of the rigid body
	body:                   Body_Handle,
	// Wheel radius
	radius:                 f32,
	// Wheel mass
	mass:                   f32,
	// Rest distance
	rest:                   f32,
	// Frequency of the spring, affects stiffness
	natural_frequency:      f32,
	// How much to damp velocity of the spring
	damping:                f32,
	pacejka_long:           Pacejka94_Longitudinal_Params,
	pacejka_lat:            Pacejka94_Lateral_Params,

	// Runtime state

	// Accumulated impulse
	spring_impulse:         f32,
	lateral_impulse:        f32,
	longitudinal_impulse:   f32,
	brake_friction_impulse: f32,

	// Inverse inertia along wheel rotation axis
	inv_inertia:            f32,

	// Rotation
	q:                      f32,
	// Angular velocity
	w:                      f32,
	hit:                    bool,
	hit_normal:             Vec3,
	hit_point:              Vec3,
	// rel_hit_point = rel_pos + rel_dir * hit_t
	hit_t:                  f32,
	turn_angle:             f32,
	drive_impulse:          f32,
	// Impulse of brake pad on brake disc, set outsie
	brake_impulse:          f32,
	applied_impulse:        Vec3,

	// Convenience for debug visualization to avoid recomputing
	slip_angle:             f32,
	slip_ratio:             f32,
	// Multipliers for combined friction
	slip_vec:               Vec2,

	// Free list
	next_plus_one:          i32,
}

Diff_Type :: enum {
	Open,
	Fixed,
	// TODO: LSD
}

Drive_Wheel :: struct {
	wheel:   Suspension_Constraint_Handle,
	impulse: f32,
}

Drive_Axle :: struct {
	// Params
	wheels:            [2]Drive_Wheel,
	wheel_count:       i32,
	diff_type:         Diff_Type,
	final_drive_ratio: f32,

	// State
	// Diff angular vel
	w:                 f32,
	// Impulse that constrains diff and engine motion
	engine_impulse:    f32,
	// Impulse that constrains wheels motion relative to each other
	diff_impulse:      f32,
}

Engine_Curve_Handle :: distinct spanpool.Handle
Gear_Ratios_Handle :: distinct spanpool.Handle

// This actually handles everything, engine, transmission, differential, etc. It's easier to keep it in one place
Engine :: struct {
	alive:             bool,
	// Engine Params
	rpm_torque_curve:  Engine_Curve_Handle,
	lowest_rpm:        f32,
	inertia:           f32,
	internal_friction: f32,

	// Transmission Params
	// 0 - reverse, 1 - first, etc.
	gear_ratios:       Gear_Ratios_Handle,
	axle:              Drive_Axle,

	// Engine State
	// Angular velocity, omega
	q:                 f32,
	w:                 f32,

	// Impulse applied when engine is stalling (rpm < lowest_rpm)
	unstall_impulse:   f32,
	// Friction that makes rpm go down when you're not accelerating
	friction_impulse:  f32,
	// Impulse applied from releasing throttle
	throttle_impulse:  f32,

	// Transmission State
	// -1 - reeverse, 0 - neutral, 1 - first, etc.
	gear:              i32,

	// Controls
	throttle:          f32,

	// 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
Engine_Handle :: distinct i32

INVALID_BODY :: Body_Handle(0)
INVALID_SUSPENSION_CONSTRAINT :: Suspension_Constraint_Handle(0)
INVALID_ENGINE :: Engine_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_engine_handle_valid :: proc(handle: Engine_Handle) -> bool {
	return i32(handle) > 0
}
is_handle_valid :: proc {
	is_body_handle_valid,
	is_suspension_constraint_handle_valid,
	is_engine_handle_valid,
}

index_to_body_handle :: proc(idx: int) -> Body_Handle {
	return Body_Handle(idx + 1)
}

body_handle_to_index :: proc(handle: Body_Handle) -> int {
	return int(handle) - 1
}

handle_to_index :: proc {
	body_handle_to_index,
}

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

_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_prev_sim_state_index :: proc(scene: ^Scene) -> i32 {
	return (scene.simulation_state_index - 1) %% i32(len(scene.simulation_states))
}

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[get_prev_sim_state_index(scene)]
}

get_next_sim_state :: proc(scene: ^Scene) -> ^Sim_State {
	return(
		&scene.simulation_states[(scene.simulation_state_index + 1) %% i32(len(scene.simulation_states))] \
	)
}

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

/// 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:     Vec3,
	initial_rot:     Quat,
	initial_vel:     Vec3,
	initial_ang_vel: Vec3,
	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:           Vec3,
	rel_dir:           Vec3,
	body:              Body_Handle,
	rest:              f32,
	natural_frequency: f32,
	damping:           f32,
	radius:            f32,
	mass:              f32,
	pacejka_long:      Pacejka94_Longitudinal_Params,
	pacejka_lat:       Pacejka94_Lateral_Params,
}

Drive_Axle_Config :: struct {
	wheels:            [2]Suspension_Constraint_Handle,
	wheel_count:       i32,
	diff_type:         Diff_Type,
	final_drive_ratio: f32,
}

Engine_Config :: struct {
	rpm_torque_curve:  [][2]f32,
	lowest_rpm:        f32,
	inertia:           f32,
	internal_friction: f32,

	// Transmission Params
	// 0 - reverse, 1 - first, etc.
	gear_ratios:       []f32,
	axle:              Drive_Axle_Config,
}

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 = s.mesh.center
		convex.extent = s.mesh.extent
		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.natural_frequency = config.natural_frequency
	constraint.damping = config.damping
	constraint.radius = config.radius
	constraint.inv_inertia = 1.0 / (0.5 * config.mass * config.radius * config.radius)
	constraint.pacejka_long = config.pacejka_long
	constraint.pacejka_lat = config.pacejka_lat
}

add_engine_curve :: proc(sim_state: ^Sim_State, curve: [][2]f32) -> Engine_Curve_Handle {
	handle := spanpool.allocate_elems(&sim_state.rpm_torque_curves_pool, ..curve)
	return Engine_Curve_Handle(handle)
}

remove_engine_curve :: proc(sim_state: ^Sim_State, handle: Engine_Curve_Handle) {
	spanpool.free(&sim_state.rpm_torque_curves_pool, spanpool.Handle(handle))
}

get_engine_curve :: proc(sim_state: ^Sim_State, handle: Engine_Curve_Handle) -> [][2]f32 {
	return spanpool.resolve_slice(&sim_state.rpm_torque_curves_pool, spanpool.Handle(handle))
}

add_gear_ratios :: proc(sim_state: ^Sim_State, gear_ratios: []f32) -> Gear_Ratios_Handle {
	handle := spanpool.allocate_elems(&sim_state.gear_ratios_pool, ..gear_ratios)
	return Gear_Ratios_Handle(handle)
}

remove_gear_ratios :: proc(sim_state: ^Sim_State, handle: Gear_Ratios_Handle) {
	spanpool.free(&sim_state.gear_ratios_pool, spanpool.Handle(handle))
}

get_gear_ratios :: proc(sim_state: ^Sim_State, handle: Gear_Ratios_Handle) -> []f32 {
	return spanpool.resolve_slice(&sim_state.gear_ratios_pool, spanpool.Handle(handle))
}

update_engine_from_config :: proc(
	sim_state: ^Sim_State,
	engine: Engine_Ptr,
	config: Engine_Config,
) {
	remove_engine_curve(sim_state, engine.rpm_torque_curve)
	engine.rpm_torque_curve = add_engine_curve(sim_state, config.rpm_torque_curve)

	remove_gear_ratios(sim_state, engine.gear_ratios)
	engine.gear_ratios = add_gear_ratios(sim_state, config.gear_ratios)

	engine.lowest_rpm = config.lowest_rpm
	engine.inertia = config.inertia
	engine.internal_friction = config.internal_friction

	engine.axle.final_drive_ratio = config.axle.final_drive_ratio
	engine.axle.wheels[0].wheel = config.axle.wheels[0]
	engine.axle.wheels[1].wheel = config.axle.wheels[1]
	engine.axle.wheel_count = config.axle.wheel_count
	engine.axle.diff_type = config.axle.diff_type
}

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

	sim_state.num_bodies += 1

	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)

		sim_state.num_bodies -= 1
	}
}

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

invalid_engine: Engine

get_engine :: proc(sim_state: ^Sim_State, handle: Engine_Handle) -> Engine_Ptr {
	index := int(handle) - 1
	if index < 0 || index >= len(sim_state.bodies_slice) {
		return &invalid_engine
	}

	return &sim_state.engines[index]
}

add_engine :: proc(sim_state: ^Sim_State, config: Engine_Config) -> Engine_Handle {
	sim_state.num_engines += 1

	engine: Engine
	update_engine_from_config(sim_state, &engine, config)

	engine.alive = true

	if sim_state.first_free_engine_plus_one > 0 {
		index := sim_state.first_free_engine_plus_one
		new_engine := get_engine(sim_state, Engine_Handle(index))
		next_plus_one := new_engine.next_plus_one
		new_engine^ = engine
		sim_state.first_free_engine_plus_one = next_plus_one
		return Engine_Handle(index)
	}

	append(&sim_state.engines, engine)
	index := len(sim_state.engines)
	return Engine_Handle(index)
}

remove_engine :: proc(sim_state: ^Sim_State, handle: Engine_Handle) {
	if is_handle_valid(handle) {
		engine := get_engine(sim_state, handle)

		remove_engine_curve(sim_state, engine.rpm_torque_curve)
		remove_gear_ratios(sim_state, engine.gear_ratios)

		engine.alive = false
		engine.next_plus_one = sim_state.first_free_engine_plus_one
		sim_state.first_free_engine_plus_one = i32(handle)
		sim_state.num_engines -= 1
	}
}


_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)
	delete_soa(sim_state.contact_container.contacts)
	delete_map(sim_state.contact_container.lookup)
	delete(sim_state.engines)
	convex_container_destroy(&sim_state.convex_container)
	spanpool.destroy_spanpool(&sim_state.rpm_torque_curves_pool)
	spanpool.destroy_spanpool(&sim_state.gear_ratios_pool)
}

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