Rides/TrackRides/Train.cs
using System;

namespace HC3.Rides;

#nullable enable

/// <summary>
/// Describes the state of a <see cref="Train"/>.
/// </summary>
public enum TrainState
{
	/// <summary>
	/// Train has left the station and is doing the main ride loop.
	/// Will transition to <see cref="ArrivingAtStation"/> when a platform is encountered.
	/// </summary>
	MainCycle = 0,

	/// <summary>
	/// Train is taxiing down at a platform. Will wait behind other trains in front of it.
	/// When at the head of the platform, will stop and transition to <see cref="UnloadingGuests"/>.
	/// </summary>
	ArrivingAtStation = 1,

	/// <summary>
	/// Stopped at a station and waiting for guests to get out. Will transition to <see cref="LoadingGuests"/>
	/// when empty.
	/// </summary>
	UnloadingGuests = 2,

	/// <summary>
	/// Stopped at a station and waiting for guests to board. Will transition to <see cref="LeavingStation"/>
	/// when full, or if <see cref="BasicRide.MaximumLoadtime"/> elapses.
	/// </summary>
	LoadingGuests = 3,

	/// <summary>
	/// Train is taxiing to leave a station. Will transition to <see cref="MainCycle"/> when fully out of the
	/// station.
	/// </summary>
	LeavingStation = 4
}

/// <summary>
/// A vehicle attached to a <see cref="Rides.TrackSection"/>, containing one or more <see cref="TrainCar"/>s.
/// </summary>
public sealed class Train : Component
{
	private float _trackPosition;
	private float _length;
	private float _mass;

	private TrainState _state;

	private bool _consistInvalid = true;
	private readonly List<TrainCar> _consist = new();
	private readonly List<SlotMarker> _slotMarkers = new();

	[Property, Sync( SyncFlags.FromHost )]
	public TrackSection? TrackSection { get; set; }

	/// <summary>
	/// Train immediately ahead of us.
	/// </summary>
	[Property]
	public Train? NextTrain { get; set; }

	[Property, Sync( SyncFlags.FromHost | SyncFlags.Interpolate )]
	public float TrackPosition
	{
		get => _trackPosition;
		set
		{
			_trackPosition = value;
			UpdateCarTransforms();
		}
	}

	[Property, Sync( SyncFlags.FromHost | SyncFlags.Interpolate )]
	public float Velocity { get; set; }

	public float Momentum => Velocity * Mass;

	public float? NextTrainDistance { get; set; }

	[Property, Sync( SyncFlags.FromHost )]
	public TrainState State
	{
		get => _state;
		set
		{
			if ( _state == value ) return;

			_state = value;
			StateTime = 0f;
		}
	}

	/// <summary>
	/// How long has the train been in the current <see cref="State"/>, in seconds?
	/// </summary>
	public float StateTime { get; private set; }

	public IReadOnlyList<TrainCar> Consist
	{
		get
		{
			UpdateConsist();
			return _consist;
		}
	}

	public IReadOnlyList<SlotMarker> SlotMarkers
	{
		get
		{
			UpdateConsist();
			return _slotMarkers;
		}
	}

	public float Length
	{
		get
		{
			UpdateConsist();
			return _length;
		}
	}

	public float Mass
	{
		get
		{
			UpdateConsist();
			return _mass;
		}
	}

	protected override void OnEnabled()
	{
		InvalidateConsist();
	}

	public void InvalidateConsist()
	{
		_consistInvalid = true;
	}

	private void UpdateConsist()
	{
		if ( !_consistInvalid ) return;

		foreach ( var trainCar in _consist )
		{
			if ( !trainCar.IsValid() ) continue;

			if ( trainCar.Train == this )
			{
				trainCar.Train = null;
			}
		}

		_consistInvalid = false;

		_consist.Clear();
		_consist.AddRange( GetComponentsInChildren<TrainCar>() );

		_slotMarkers.Clear();
		_slotMarkers.AddRange( GetComponentsInChildren<SlotMarker>() );

		_length = _consist.Sum( x => x.Length );
		_mass = _consist.Sum( x => x.Mass );
	}

	private void UpdateCarTransforms()
	{
		var position = TrackPosition;

		foreach ( var car in Consist )
		{
			car.Train = this;

			position -= car.Length * 0.5f;
			car.TrackPosition = position;
			position -= car.Length * 0.5f;
		}
	}

	private float SimulateGravity()
	{
		if ( TrackSection is not ITrackSection { Elements.Count: > 0 } trackSection ) return 0f;

		var force = 0f;

		foreach ( var car in Consist )
		{
			var rotation = trackSection.SampleAtDistance( car.TrackPosition ).Rotation;
			var gForce = Gravity * -rotation.Forward.z * car.Mass;

			force += gForce;
		}

		return force;
	}

	private float SimulateFriction( float restVelocity, float coefficient )
	{
		return coefficient * (restVelocity - Velocity);
	}

	private readonly List<PlacedTrackElement> _elements = new();

	public void PrePhysicsUpdate( float dt )
	{
		if ( TrackSection is not ITrackSection { Elements.Count: > 0 } track ) return;

		var force = SimulateGravity() + SimulateFriction( 0f, DefaultFriction );

		_elements.Clear();

		var next = track.GetElements( TrackPosition, Length, _elements );

		if ( _elements.Count > 0 )
		{
			force += State switch
			{
				TrainState.MainCycle => State_MainCycle( track, _elements, next ),
				TrainState.ArrivingAtStation => State_ArrivingAtStation( track, _elements, next ),
				TrainState.LeavingStation => State_LeavingStation( track, _elements, next ),
				TrainState.UnloadingGuests or TrainState.LoadingGuests => State_StoppedAtStation( track, _elements, next ),
				_ => 0f
			};
		}

		Velocity += force * dt / Mass;

		if ( _elements.Any( static x => x.Feature == TrackFeature.ChainLift ) )
		{
			Velocity = Math.Max( ChainLiftVelocity, Velocity );
		}
	}

	public void PostPhysicsUpdate( float dt )
	{
		TrackPosition += Velocity * dt;
		StateTime += dt;
	}

	private const float Gravity = 500f;
	private const float DefaultFriction = 25f;
	private const float ChainLiftVelocity = 40f;
	private const float StationVelocity = 32f;
	private const float StationFriction = 4_000f;

	private const float StoppedSpeed = 0.001f;

	private float State_MainCycle( ITrackSection track, IReadOnlyList<PlacedTrackElement> elements, PlacedTrackElement? next )
	{
		if ( elements[0].Feature is TrackFeature.Station )
		{
			State = TrainState.ArrivingAtStation;
		}

		return 0f;
	}

	private float State_ArrivingAtStation( ITrackSection track, IReadOnlyList<PlacedTrackElement> elements, PlacedTrackElement? next )
	{
		var targetVelocity = StationVelocity;

		if ( elements[0].Feature is not TrackFeature.Station )
		{
			// We went too fast and missed the station!

			State = TrainState.LeavingStation;
		}
		else if ( next is { Feature: not TrackFeature.Station, StartDistance: var stationEnd } )
		{
			// Slow down as we reach the end of the station

			var dist = track.NormalizeDelta( stationEnd - TrackPosition );

			targetVelocity = dist < 16f ? 0f : StationVelocity * 0.5f;

			if ( MathF.Abs( Velocity ) <= StoppedSpeed )
			{
				// Start unloading guests when we stop

				State = TrainState.UnloadingGuests;
			}
		}

		if ( NextTrain is { } nextTrain )
		{
			// Slow down for the train ahead of us

			var dist = track.NormalizeDelta( nextTrain.TrackPosition - nextTrain.Length - TrackPosition );

			targetVelocity = Math.Min( targetVelocity, dist < 16f ? nextTrain.Velocity : MathX.Lerp( nextTrain.Velocity, StationVelocity, (dist - 16f) / 48f ) );
		}

		return SimulateFriction( targetVelocity, StationFriction );
	}

	private float State_LeavingStation( ITrackSection track, IReadOnlyList<PlacedTrackElement> elements, PlacedTrackElement? next )
	{
		if ( elements.All( static x => x.Feature is not TrackFeature.Station ) )
		{
			State = TrainState.MainCycle;
		}

		return Velocity < StationVelocity ? SimulateFriction( StationVelocity, StationFriction ) : 0f;
	}

	private float State_StoppedAtStation( ITrackSection track, IReadOnlyList<PlacedTrackElement> elements, PlacedTrackElement? next )
	{
		if ( elements[0].Feature is not TrackFeature.Station )
		{
			// We went too fast and missed the station!

			State = TrainState.LeavingStation;
		}
		else if ( MathF.Abs( Velocity ) > StoppedSpeed )
		{
			State = TrainState.ArrivingAtStation;
		}

		return SimulateFriction( 0f, StationFriction );
	}

	public sealed record Model( TrainState State, float Position, float Velocity );

	public Model Serialized
	{
		get => new( State, TrackPosition, Velocity );
		set
		{
			State = value.State;
			TrackPosition = value.Position;
			Velocity = value.Velocity;
		}
	}
}