class_name TrainingMode
extends GameMode

# Headless self-play training mode: two RL-driven ships, no HUD, no camera.
# The scene also contains the godot_rl_agents Sync node, which speaks TCP to
# the Python trainer; this mode owns the environment rules — episodes, goal
# rewards, and randomized episode-start states (the RLGym "state setter"
# lesson: varied starts massively speed up learning versus kickoff-only).
#
# Run: godot --headless --path Game res://scenes/training.tscn
# (started automatically by training/train.py; boots into idle ships with a
# warning if no trainer is listening).
#
# Eval mode (used by training/evaluate.py): pass --eval_model_a=<path> and
# --eval_model_b=<path> (+ optional --eval_episodes=N) and both ships are
# instead driven by those exported policies via AIShipController; each episode
# ends at the first goal (or a draw on timeout), and a final machine-readable
# "EVAL_RESULT {...}" line is printed before quitting.

@export var episode_length_seconds := 30.0
@export var goal_reward := 10.0

# Episode-start state mix; remaining probability = fully random state.
@export_range(0.0, 1.0) var kickoff_state_chance := 0.2
@export_range(0.0, 1.0) var ball_near_goal_chance := 0.2

# Placement bounds, inset from the arena (goals at z ≈ ±15.56).
const FIELD_HALF_X := 10.0
const FIELD_HALF_Z := 13.0
const FIELD_MIN_Y := 1.5
const FIELD_MAX_Y := 8.0
const MAX_RANDOM_BALL_SPEED := 12.0
const MAX_RANDOM_SHIP_SPEED := 8.0

# Sim runs at 60 physics ticks per sim-second regardless of speedup.
const TICKS_PER_SIM_SECOND := 60.0

# The arena has no walls/ceiling yet (see TODO.md), so an untrained policy can
# simply fly away. Leaving this volume ends the episode with a small penalty.
const BOUNDS_HALF_X := 25.0
const BOUNDS_HALF_Z := 25.0
const BOUNDS_MIN_Y := -5.0
const BOUNDS_MAX_Y := 25.0
@export var out_of_bounds_penalty := 1.0

var _agents: Array[ShipAIController] = []

# Eval mode state (see header comment)
var _eval := false
var _eval_models: Array[String] = ["", ""]
var _eval_episodes := 20
var _eval_goals := {0: 0, 1: 0}
var _eval_draws := 0
var _eval_episodes_done := 0
var _episode_ticks := 0


func _start() -> void:
	_parse_eval_args()
	spawn_ball()
	if _eval:
		for team in [0, 1]:
			var bot := AIShipController.new()
			bot.model_path = _eval_models[team]
			spawn_ship(team, 0, bot)
		return
	var ship_team0 := spawn_ship(0, 0, RLShipController.new())
	var ship_team1 := spawn_ship(1, 0, RLShipController.new())
	_attach_agent(ship_team0, ship_team1)
	_attach_agent(ship_team1, ship_team0)


func _parse_eval_args() -> void:
	var args := {}
	for argument in OS.get_cmdline_args():
		if argument.begins_with("--") and argument.find("=") > -1:
			var key_value := argument.lstrip("--").split("=", true, 1)
			args[key_value[0]] = key_value[1]
	if args.has("eval_model_a") and args.has("eval_model_b"):
		_eval = true
		_eval_models[0] = args["eval_model_a"]
		_eval_models[1] = args["eval_model_b"]
		_eval_episodes = int(args.get("eval_episodes", str(_eval_episodes)))


func _attach_agent(ship: Ship, opponent: Ship) -> void:
	var agent := ShipAIController.new()
	agent.name = "ShipAIController"
	agent.reset_after = int(episode_length_seconds * TICKS_PER_SIM_SECOND)
	ship.add_child(agent)
	agent.setup(ship, ship.controller as RLShipController, ball, opponent, _attack_goal_position(ship.team))
	_agents.append(agent)


# The goal this team scores into: the one the opponent defends/concedes.
func _attack_goal_position(team: int) -> Vector3:
	for goal in arena.get_goals():
		if goal.team == 1 - team:
			return goal.global_position
	push_error("TrainingMode: no goal found for team %d to attack" % team)
	return Vector3.ZERO


func _physics_process(_delta):
	if _eval:
		_episode_ticks += 1
		for ship in ships:
			if is_instance_valid(ship) and _out_of_bounds(ship.global_position):
				_place_body(ship, _ship_spawn_transforms[ship], Vector3.ZERO, Vector3.ZERO)
		var ball_lost := is_instance_valid(ball) and _out_of_bounds(ball.global_position)
		if _episode_ticks > int(episode_length_seconds * TICKS_PER_SIM_SECOND) or ball_lost:
			_eval_draws += 1
			_end_eval_episode()
		return

	# Both trainer-requested resets and truncation (reset_after ticks elapsed)
	# surface as needs_reset. Only truncation is an episode end the trainer
	# must be told about via done — a trainer-requested reset already knows.
	var needs_reset := false
	var truncated := false
	for agent in _agents:
		needs_reset = needs_reset or agent.needs_reset
		truncated = truncated or agent.n_steps > agent.reset_after
	if needs_reset:
		if truncated:
			for agent in _agents:
				agent.done = true
		_reset_episode()
		return

	# A ship leaving the play volume is penalized and respawned at its kickoff
	# spawn (the episode continues); a lost ball ends the episode for both.
	for agent in _agents:
		if _out_of_bounds(agent.ship.global_position):
			agent.reward -= out_of_bounds_penalty
			_place_body(agent.ship, _ship_spawn_transforms[agent.ship], Vector3.ZERO, Vector3.ZERO)
	if is_instance_valid(ball) and _out_of_bounds(ball.global_position) and not _agents.is_empty():
		for agent in _agents:
			agent.done = true
		_reset_episode()


func _out_of_bounds(position: Vector3) -> bool:
	return absf(position.x) > BOUNDS_HALF_X \
		or absf(position.z) > BOUNDS_HALF_Z \
		or position.y < BOUNDS_MIN_Y \
		or position.y > BOUNDS_MAX_Y


func _on_goal_scored(conceding_team: int) -> void:
	if _eval:
		_eval_goals[1 - conceding_team] += 1
		_end_eval_episode()
		return
	for agent in _agents:
		agent.reward += goal_reward if agent.ship.team != conceding_team else -goal_reward
		agent.done = true
	_reset_episode()


func _end_eval_episode() -> void:
	_eval_episodes_done += 1
	_episode_ticks = 0
	if _eval_episodes_done >= _eval_episodes:
		print("EVAL_RESULT " + JSON.stringify({
			"model_a": _eval_models[0],
			"model_b": _eval_models[1],
			"episodes": _eval_episodes_done,
			"goals_a": _eval_goals[0],
			"goals_b": _eval_goals[1],
			"draws": _eval_draws,
		}))
		get_tree().quit()
		return
	# Randomized states (not kickoff): deterministic policies would otherwise
	# replay the identical episode every time.
	_reset_episode()


func _reset_episode() -> void:
	for agent in _agents:
		agent.reset()

	var roll := randf()
	if roll < kickoff_state_chance:
		reset_ball()
		reset_ships()
	elif roll < kickoff_state_chance + ball_near_goal_chance:
		_place_ships_random()
		_place_ball_near_goal()
	else:
		_place_ships_random()
		_place_ball_random()


func _place_ball_random() -> void:
	var velocity := _random_direction() * randf_range(0.0, MAX_RANDOM_BALL_SPEED)
	_place_body(ball, Transform3D(Basis.IDENTITY, _random_position()), velocity, Vector3.ZERO)


# Attacking/defending drill states: ball close to a goal, moving toward it.
func _place_ball_near_goal() -> void:
	var goals := arena.get_goals()
	var goal: Goal = goals[randi() % goals.size()]
	var toward_centre := -signf(goal.global_position.z)
	var position := Vector3(
		randf_range(-4.0, 4.0),
		randf_range(FIELD_MIN_Y, 4.0),
		goal.global_position.z + toward_centre * randf_range(3.0, 6.0)
	)
	var to_goal := (goal.global_position - position).normalized()
	var velocity := (to_goal + _random_direction() * 0.3).normalized() * randf_range(2.0, MAX_RANDOM_BALL_SPEED)
	_place_body(ball, Transform3D(Basis.IDENTITY, position), velocity, Vector3.ZERO)


func _place_ships_random() -> void:
	for ship in ships:
		var orientation := Basis.from_euler(Vector3(
			randf_range(-0.4, 0.4),
			randf_range(-PI, PI),
			randf_range(-0.4, 0.4)
		))
		var velocity := _random_direction() * randf_range(0.0, MAX_RANDOM_SHIP_SPEED)
		_place_body(ship, Transform3D(orientation, _random_position()), velocity, Vector3.ZERO)


func _random_position() -> Vector3:
	return Vector3(
		randf_range(-FIELD_HALF_X, FIELD_HALF_X),
		randf_range(FIELD_MIN_Y, FIELD_MAX_Y),
		randf_range(-FIELD_HALF_Z, FIELD_HALF_Z)
	)


func _random_direction() -> Vector3:
	var direction := Vector3(randf_range(-1, 1), randf_range(-1, 1), randf_range(-1, 1))
	return direction.normalized() if direction.length_squared() > 0.001 else Vector3.FORWARD


func _place_body(body: RigidBody3D, to: Transform3D, linear_velocity: Vector3, angular_velocity: Vector3) -> void:
	# Deferred: a RigidBody3D transform can't be set mid-physics-step
	body.set_deferred("global_transform", to)
	body.set_deferred("linear_velocity", linear_velocity)
	body.set_deferred("angular_velocity", angular_velocity)