class_name ShipAIController
extends AIController3D

# Training-side bridge between godot_rl_agents and a ship. This is the only
# class that touches plugin types (AIController3D / the Sync node protocol) —
# everything else stays behind the ShipController seam: actions received from
# the trainer are written into an RLShipController, which the ship pulls like
# any other controller.
#
# Action space is ShipAction verbatim: 6 continuous axes (thrust xyz,
# rotation xyz, each -1..1) + binary turbo. ShipAction axes are ship-local
# (body frame), so they need no team mirroring — only observations do
# (see ShipObservations.canon).

# Reward shaping weights. Dense terms accrue per physics tick (60 sim-ticks
# per sim-second); event terms fire once. Exported so tuning needs no code
# edits. Goal rewards are added by TrainingMode, which owns goal events.
@export var ball_touch_reward := 0.1
@export var velocity_to_ball_weight := 0.001
@export var ball_velocity_to_goal_weight := 0.004

var ship: Ship
var rl_controller: RLShipController
var ball: RigidBody3D
var opponent: Ship
var attack_goal_position: Vector3


# Wire up references after the ship is spawned. `attack_goal` is the goal
# this ship scores into (goal.team == opponent's team).
func setup(p_ship: Ship, p_rl_controller: RLShipController, p_ball: RigidBody3D, p_opponent: Ship, p_attack_goal_position: Vector3) -> void:
	ship = p_ship
	rl_controller = p_rl_controller
	ball = p_ball
	opponent = p_opponent
	attack_goal_position = p_attack_goal_position
	init(ship)

	# Contact monitoring for the ball-touch reward (training-only cost;
	# the shipped game leaves contact_monitor off).
	ship.contact_monitor = true
	ship.max_contacts_reported = 8
	ship.body_entered.connect(_on_ship_body_entered)


func get_obs() -> Dictionary:
	return {"obs": ShipObservations.build(ship, opponent, ball, attack_goal_position)}


func get_reward() -> float:
	return reward


func get_action_space() -> Dictionary:
	return {
		"thrust": {"size": 3, "action_type": "continuous"},
		"rotation": {"size": 3, "action_type": "continuous"},
		"turbo": {"size": 2, "action_type": "discrete"},
	}


func set_action(action) -> void:
	var thrust: Array = action["thrust"]
	var rot: Array = action["rotation"]
	rl_controller.action.thrust = Vector3(thrust[0], thrust[1], thrust[2])
	rl_controller.action.rotation = Vector3(rot[0], rot[1], rot[2])
	rl_controller.action.turbo = int(action["turbo"]) == 1


func _physics_process(delta):
	super(delta)
	if not is_instance_valid(ship) or not is_instance_valid(ball):
		return

	# Dense shaping: own velocity toward the ball
	var to_ball := ball.global_position - ship.global_position
	if to_ball.length_squared() > 0.0001:
		var closing_speed := ship.linear_velocity.dot(to_ball.normalized())
		reward += velocity_to_ball_weight * closing_speed / ship.max_speed

	# Dense shaping: ball velocity toward the goal we attack
	var ball_to_goal := attack_goal_position - ball.global_position
	if ball_to_goal.length_squared() > 0.0001:
		var ball_progress := ball.linear_velocity.dot(ball_to_goal.normalized())
		reward += ball_velocity_to_goal_weight * ball_progress / ShipObservations.BALL_SPEED_SCALE


func _on_ship_body_entered(body: Node) -> void:
	if body.is_in_group("ball"):
		reward += ball_touch_reward