State machine doesn't work as expected

Hi, total beginner here! I just dived into state machines and how they work and have crafted a nice little state machine. Except there are some bugs which didn’t appear with the normal imperative way of how you would code a player script. The main issues I am facing are: Jump count doesn’t reset properly even after resetting it when state is idle or on the ground, the player sometimes floats away when doing a double jump and maybe there are some more issues. Also, is this a preferred way of doing a state machine or can improvements be made? Thanks in advance!

class_name Player extends CharacterBody2D


@export var FRICTION = 2000.0
@export var SPEED = 1500.0
@export var STEP_DISTANCE = 250.0
@export var JUMP_VELOCITY = -600.0
@export var MAX_JUMPS = 2
# Coyote time is the time the user has to react if they fall off the platform
@export var COYOTE_TIME_WINDOW = 0.1

enum States {IDLE, MOVING, JUMPING, FALLING}
var state: States = States.IDLE: set = set_state

var jump_count := 0
var time_in_air := 0.0
var direction_sign := 0
var target_distance := 0.0

func set_state(new_state: int) -> void:
	state = new_state
	
	if state == States.MOVING:
		target_distance = STEP_DISTANCE
	
	## TODO: add animations for the state

func _physics_process(delta: float) -> void:
	var is_initiating_jump := Input.is_action_just_pressed("jump") and (jump_count == MAX_JUMPS - 1 or (is_on_floor() or time_in_air < COYOTE_TIME_WINDOW))
	var is_initiating_burst := Input.is_action_just_pressed("move_left") or Input.is_action_just_pressed("move_right")
	
	if is_initiating_jump:
		state = States.JUMPING
		jump_count += 1
		velocity.y = JUMP_VELOCITY
	elif state == States.JUMPING and velocity.y > 0.0:
		state = States.FALLING
	elif is_initiating_burst:
		state = States.MOVING
	
	if state in [States.JUMPING, States.FALLING]:
		velocity += get_gravity() * delta
		time_in_air += delta
	else:
		var direction := Input.get_axis("move_left", "move_right")
		direction_sign = sign(direction)
		if state == States.MOVING:
			velocity.x = direction_sign * SPEED 
		else:
			velocity.x = move_toward(velocity.x, 0, FRICTION * delta)
			
		if is_on_floor():
			time_in_air = 0
			jump_count = 0
		
		if target_distance > 0:
			target_distance -= abs(velocity.x * delta)
		elif target_distance <= 0 and is_on_floor():
			velocity.x = 0
			state = States.IDLE
	
	move_and_slide()

jump_count is only getting reset when the state is neither JUMPING nor FALLING, but the state doesn’t change upon reaching the ground.
And pressing either “move_left” or “move_right” will always change the state to MOVING, even if the player is in the air. But gravity is only applied in JUMPING and FALLING state.

In general, the benefit of good state machines is to have states with distinct behaviors and clear ways to transition between them. Your current code doesn’t seem to achieve any of that. In my opinion, this kind of state machine mainly makes the code less comprehensible.

But would you just remove the state machine fully if you were me or would you rewrite it to look better?

If it’s just about the current code I would remove it. Since IDLE+MOVING (and JUMPING+FALLING respectively) behave exactly the same, you basically have just two different states (and even between those two aren’t that many differences). This doesn’t require a state machine.
You can keep the States enum though and use it for animations later. The movement shouldn’t depend on the current state then, but only keep it updated.

However, if you already know you will be expanding the movement later and maybe have more distinct states to add in the future you could think about rewriting it instead. In this case, for anything to be considered as a ‘state’ it should have distinguishable behavior and not just a different animation to play (e. g. if the horizontal movement shall be different in the air then on the ground, it would make more sense to use a state machine).