Struggle with Swim controls and logic 2D

Hello!
I’ve tried to make swim mechanic in my game but there’s one bug that if you jump into water from surface gravity will drag you very fast. That also works when you go to water by jumping from the side but you’ll go up fast
How to solve it?
That’s move controls:

extends CharacterBody2D


const SPEED = 230.0
const JUMP_VELOCITY = -450.0
const SWIM_SPEED = 50
const SWIM_VELOCITY = -100
var gravity = ProjectSettings.get_setting("physics/2d/default_gravity")

@onready var jump_buffer_timer: = $JumpBufferTimer
@onready var coyote_timer_jump: = $CoyoteTimerJump

var jump_buffer: = false
var can_jump = true
var in_water = false
func _physics_process(delta):
	if in_water == true:
		swim_movement(delta)
	else:
		normal_movement(delta)

	
	# Handle jump.
	jump()
	
	
	
		#FLipping animation
	var direction = Input.get_axis("move_left", "move_right")
	if direction > 0:
		$AnimatedSprite2D.flip_h = false
	elif direction < 0:
		$AnimatedSprite2D.flip_h = true
		#Animation
	if is_on_floor():
		if direction == 0:
			$AnimatedSprite2D.play("idle")
		else: 
			$AnimatedSprite2D.play("walk")
	else:
		$AnimatedSprite2D.play("jump")
		
	
		
		#Moving character
	if direction:
		velocity.x = direction * SPEED
	else:
		velocity.x = move_toward(velocity.x, 0, SPEED)
		
	var was_on_floor = is_on_floor()
	move_and_slide()
	var just_left_ledge = was_on_floor and not is_on_floor() and velocity.y >= 0
	if just_left_ledge:
		coyote_timer_jump.start()
		
	if is_on_floor() and jump_buffer:
		velocity.y = JUMP_VELOCITY
		jump_buffer = false
		jump_buffer_timer.stop()
		can_jump = false
		
	if is_on_floor():
		can_jump = true
func jump():
	if (is_on_floor() or coyote_timer_jump.time_left > 0.0 or jump_buffer == true) and can_jump:
		if Input.is_action_just_pressed("jump"):
			velocity.y = JUMP_VELOCITY
			jump_buffer = false
			jump_buffer_timer.stop()
			can_jump = false
	elif Input.is_action_just_released("jump") and velocity.y < JUMP_VELOCITY / 2:
			velocity.y = JUMP_VELOCITY / 2
			can_jump = false
	elif Input.is_action_just_pressed("jump"):
			jump_buffer = true
			jump_buffer_timer.start()
			
func normal_movement(delta):
		#gravity
	if not is_on_floor():
		velocity.y += gravity * delta


func _on_jump_buffer_timer_timeout():
	jump_buffer = false

func swim_movement(delta):
	
	if Input.is_action_just_pressed("jump"):
		velocity.y = SWIM_VELOCITY
	else:
		velocity.y += gravity * 0.1 * delta #slower fall
	
	if Input.is_action_just_pressed("move_left"):
		velocity.x = -SWIM_SPEED
	elif Input.is_action_just_pressed("move_right"):
		velocity.x = SWIM_SPEED
	else:
		velocity.x = 0
		
	move_and_slide()

And that’s water code:

extends Area2D


func _on_body_entered(_body):
	_body.in_water = true
	print("in water!")# Replace with function body.



func _on_body_exited(_body):
	_body.in_water = false
	print("exited!") # Replace with function body.