Rides/TrackRides/TrackRide.cs
using System;
using System.Collections.Immutable;
using System.Text.Json.Nodes;
using HC3.Terrain;

namespace HC3.Rides;

#nullable enable

/// <summary>
/// This component has an opinion about what track elements are allowed when building <see cref="TrackSection"/>s
/// in descendant objects.
/// </summary>
public interface ITrackConstraints
{
	bool IsElementAllowed( TrackElement element );
}

/// <summary>
/// Base type for rides that contain <see cref="HC3.Rides.TrackSection"/>s and <see cref="TrainCar"/>s.
/// </summary>
public class TrackRide : BasicRide, ITrackConstraints, ITrackEvent
{
	private TrackDefinition? _trackDef;

	/// <summary>
	/// Describes how to generate a model for this ride's track.
	/// </summary>
	[Property, Feature( "Track" )]
	public TrackDefinition? TrackDefinition
	{
		get => _trackDef;
		set
		{
			_trackDef = value;

			foreach ( var trackSection in TrackSections )
			{
				trackSection.TrackDefinition = value;
			}
		}
	}

	/// <summary>
	/// Steepest incline this ride can use when building track.
	/// </summary>
	[Property, Feature( "Track" )]
	public TrackIncline MaxIncline { get; set; } = TrackIncline.SteepUp;

	/// <summary>
	/// Steepest banking this ride can use when building track.
	/// </summary>
	[Property, Feature( "Track" )]
	public TrackBanking MaxBanking { get; set; } = TrackBanking.SteepRight;

	/// <summary>
	/// How high off the ground we can build, in multiples of <see cref="GridManager.HeightStep"/>.
	/// </summary>
	[Property, Feature( "Track" )]
	public int MaxElevation { get; set; } = 256;

	/// <summary>
	/// Special elements this ride is allowed to use.
	/// </summary>
	[Property, Feature( "Track" )]
	public List<TrackElementDefinition> SpecialElements { get; set; } = new();

	/// <summary>
	/// Prefab for the first car in a train. Defaults to <see cref="TrainCarPrefab"/>.
	/// </summary>
	[Property, Feature( "Vehicle" )]
	public GameObject? TrainHeadPrefab { get; set; }

	/// <summary>
	/// Prefab for cars in a train.
	/// </summary>
	[Property, Feature( "Vehicle" )]
	public GameObject? TrainCarPrefab { get; set; }

	/// <summary>
	/// Prefab for the last car in a train. Defaults to <see cref="TrainCarPrefab"/>.
	/// </summary>
	[Property, Feature( "Vehicle" )]
	public GameObject? TrainTailPrefab { get; set; }

	/// <summary>
	/// How much one tile of track costs at ground level.
	/// </summary>
	[Property, Feature( "Construction" )]
	public int BaseTrackCost { get; set; } = 50;

	/// <summary>
	/// Additional cost per height step that a tile of track is elevated.
	/// </summary>
	[Property, Feature( "Construction" )]
	public int ElevationCost { get; set; } = 10;

	/// <summary>
	/// Calculated based on the length of the station.
	/// </summary>
	private int MaxTotalTrainCars => _platforms
		.Select( x => x.Length )
		.DefaultIfEmpty( 0 )
		.Max();

	/// <summary>
	/// How many trains we want, if the station is long enough.
	/// </summary>
	private int _targetTrainCount = 4;

	/// <summary>
	/// How long each train should be, if the station is long enough.
	/// </summary>
	private int _targetTrainLength = 4;

	[Feature( "Vehicle" ), Inspectable, MinMax( 1, 8 )]
	public int TrainCount
	{
		get => MaxTotalTrainCars == 0 ? 0 : Math.Clamp( _targetTrainCount, 1, (int)Math.Round( (float)MaxTotalTrainCars / TrainLength ) );
		set
		{
			_targetTrainCount = value;
			UpdateTrains();
		}
	}

	[Feature( "Vehicle" ), Inspectable( Suffix = " cars" ), MinMax( 1, 10 )]
	public int TrainLength
	{
		get => MaxTotalTrainCars == 0 ? 0 : Math.Clamp( _targetTrainLength, 1, MaxTotalTrainCars );
		set
		{
			_targetTrainLength = value;
			UpdateTrains();
		}
	}

	public bool IsElementAllowed( TrackElement element )
	{
		var isSpecial = element.Feature == TrackFeature.Special;

		if ( isSpecial && !SpecialElements.Contains( element.Definition ) ) return false;

		var start = new TrackOrientation( Incline: element.StartIncline, Banking: element.StartBanking, Inverted: element.StartInverted );
		var end = new TrackOrientation( Heading: element.EndHeading, Incline: element.EndIncline, Banking: element.EndBanking, Inverted: element.StartInverted );

		return IsOrientationAllowed( start, isSpecial ) && IsOrientationAllowed( end, isSpecial );
	}

	public bool IsOrientationAllowed( TrackOrientation orientation, bool allowInverted )
	{
		if ( Math.Abs( (int)orientation.Incline ) > Math.Abs( (int)MaxIncline ) ) return false;
		if ( Math.Abs( (int)orientation.Banking ) > Math.Abs( (int)MaxBanking ) ) return false;

		return !orientation.Inverted || allowInverted;
	}

	public IEnumerable<TrackSection> TrackSections =>
		GetComponentsInChildren<TrackSection>();

	protected override JsonObject GetPersistentMetadata()
	{
		var data = base.GetPersistentMetadata();
		var trackSections = TrackSections
			.Select( x => x.Serialized )
			.ToImmutableArray();

		data[nameof( TrackSections )] = Json.ToNode( trackSections );

		var trains = Trains
			.Select( x => x.Serialized )
			.ToImmutableArray();

		data[nameof( TrainCount )] = _targetTrainCount;
		data[nameof( TrainLength )] = _targetTrainLength;
		data[nameof( Trains )] = Json.ToNode( trains );

		return data;
	}

	protected override void SetPersistentMetadata( JsonObject obj )
	{
		base.SetPersistentMetadata( obj );

		if ( Json.FromNode<ImmutableArray<TrackSection.Model>?>( obj[nameof( TrackSections )] ) is not { } trackSections ) return;

		foreach ( var trackSection in TrackSections.ToImmutableArray() )
		{
			trackSection.DestroyGameObject();
		}

		foreach ( var model in trackSections )
		{
			var sectionObj = new GameObject( GameObject, true, "Track Section" );
			var trackSection = sectionObj.AddComponent<TrackSection>();

			sectionObj.AddComponent<TrackMesh>();

			trackSection.TrackDefinition = TrackDefinition;
			trackSection.Serialized = model;
		}

		UpdatePlatforms();

		_targetTrainCount = obj[nameof( TrainCount )]?.GetValue<int>() ?? _targetTrainCount;
		_targetTrainLength = obj[nameof( TrainLength )]?.GetValue<int>() ?? _targetTrainLength;

		UpdateTrains();

		if ( Json.FromNode<ImmutableArray<Train.Model>?>( obj[nameof( Trains )] ) is not { } trains ) return;

		for ( var i = 0; i < _trains.Count && i < trains.Length; ++i )
		{
			_trains[i].Serialized = trains[i];
		}
	}

	protected override void GenerateDecorations()
	{
		// nuffin
	}

	private sealed class PlatformInfo
	{
		public TrackSection Track { get; }
		public TileEdge ForwardEdge { get; }
		public int MinElementIndex { get; }
		public Vector3Int TailGridPos { get; set; }

		public int MaxElementIndex { get; private set; }
		public Vector3Int HeadGridPos { get; private set; }

		public int Length { get; private set; }

		public PlatformInfo( TrackSection track, TileEdge forwardEdge, int elementIndex, Vector3Int gridPos )
		{
			Track = track;
			ForwardEdge = forwardEdge;
			MinElementIndex = MaxElementIndex = elementIndex;
			HeadGridPos = TailGridPos = gridPos;
			Length = 1;
		}

		public void AddElement( int elementIndex, Vector3Int gridPos )
		{
			MaxElementIndex = elementIndex;
			HeadGridPos = gridPos;
			Length += 1;
		}
	}

	private readonly List<PlatformInfo> _platforms = new();
	private readonly List<BoxCollider> _platformColliders = new();

	public override IEnumerable<(Vector3Int GridPos, TileEdge Edge)> ValidEntranceExitPositions
	{
		get
		{
			foreach ( var platform in _platforms )
			{
				var left = platform.ForwardEdge.Rotate( -1 );
				var right = platform.ForwardEdge.Rotate( 1 );

				var forward = platform.ForwardEdge.GetDirection();
				var pos = platform.HeadGridPos;

				for ( var i = 0; i < platform.Length; ++i )
				{
					yield return (pos + left.GetDirection(), left);
					yield return (pos + right.GetDirection(), right);

					pos -= forward;
				}
			}
		}
	}

	protected override IEnumerable<SlotMarker> GetLoadableSlots()
	{
		return Trains
			.Where( x => x.State == TrainState.LoadingGuests )
			.SelectMany( x => x.SlotMarkers );
	}

	void ITrackEvent.Changed( TrackSection track )
	{
		if ( TrackSections.Contains( track ) )
		{
			UpdatePlatforms();
		}
	}

	private void UpdatePlatforms()
	{
		if ( GridManager.Instance?.Terrain is not { } terrain ) return;

		_platforms.Clear();

		foreach ( var section in TrackSections )
		{
			FindPlatforms( section, _platforms );
		}

		// TODO: support obstructions at arbitrary heights on a subset of tiles

		// Find bounds for GridObject

		var worldBounds = new RectInt( _platforms[0].HeadGridPos.x, _platforms[0].HeadGridPos.y, 1, 1 );

		foreach ( var platform in _platforms )
		{
			worldBounds.Add( new RectInt( platform.HeadGridPos.x, platform.HeadGridPos.y, 1, 1 ) );
			worldBounds.Add( new RectInt( platform.TailGridPos.x, platform.TailGridPos.y, 1, 1 ) );
		}

		// bounds is already in world space, need to un-rotate to be local again

		GridObject.LocalBounds = GridObject.BoundsToLocal( worldBounds, WorldTransform );
		GridObject.UpdatePlacement();

		// Place colliders at stations so we can click on them

		while ( _platformColliders.Count > _platforms.Count )
		{
			_platformColliders[^1].DestroyGameObject();
			_platformColliders.RemoveAt( _platformColliders.Count - 1 );
		}

		while ( _platformColliders.Count < _platforms.Count )
		{
			var obj = new GameObject( GameObject, true, "Collider" );
			var collider = obj.AddComponent<BoxCollider>();

			_platformColliders.Add( collider );
		}

		for ( var i = 0; i < _platformColliders.Count; ++i )
		{
			var platform = _platforms[i];
			var collider = _platformColliders[i];

			var start = terrain.GridToWorld( platform.HeadGridPos + new Vector3( 0.5f, 0.5f, 0f ) );
			var end = terrain.GridToWorld( platform.TailGridPos + new Vector3( 0.5f, 0.5f, 0f ) );

			collider.Scale = (Vector3.Max( start - end, end - start ) + 64f) with { z = 16f };
			collider.Center = new Vector3( 0f, 0f, 8f );
			collider.WorldRotation = Rotation.Identity;
			collider.WorldPosition = (start + end) * 0.5f;
		}

		UpdateEntranceExit();
	}

	private void FindPlatforms( TrackSection track, List<PlatformInfo> platforms )
	{
		if ( GridManager.Instance?.Terrain is not { } terrain ) return;

		var startIndex = 0;

		if ( track.Elements[^1].Feature == TrackFeature.Station )
		{
			startIndex = track.Elements.Count - 1;

			while ( startIndex > 0 && track.Elements[startIndex].Feature == TrackFeature.Station )
			{
				--startIndex;
			}
		}

		PlatformInfo? platform = null;

		for ( var i = 0; i < track.Elements.Count; ++i )
		{
			var index = (startIndex + i) % track.Elements.Count;
			var element = track.Elements[index];

			if ( element.Feature != TrackFeature.Station )
			{
				platform = null;
				continue;
			}

			var prev = track.Nodes[index];
			var next = track.Nodes[index + 1];
			var trackTileIndex = (prev.Position + next.Position) / 2;
			var worldPos = track.TrackGridToWorld( trackTileIndex );
			var worldTileIndex = terrain.GetTileIndex( terrain.WorldToGrid( worldPos ) );

			if ( platform is null )
			{
				var forwardEdge = WorldTransform.RotationToWorld( next.Orientation ).ToTileEdge();

				platform = new PlatformInfo( track, forwardEdge, index, worldTileIndex );
				platforms.Add( platform );
			}
			else
			{
				platform.AddElement( index, worldTileIndex );
			}
		}
	}

	private readonly List<Train> _trains = new();

	public IEnumerable<Train> Trains => _trains.Where( x => x.IsValid );

	protected override void OpenStateChanged()
	{
		UpdateTrains();

		base.OpenStateChanged();
	}

	private void UpdateTrains()
	{
		if ( !Networking.IsHost ) return;

		if ( _platforms.Count == 0 )
		{
			ClearTrains();
			return;
		}

		// TODO: multiple stations?
		// TODO: properly support disconnected track?
		// TODO: check actual length of cars?

		var targetCount = OpenState == OpenState.Closed ? 0 : TrainCount;
		var targetLength = TrainLength;

		if ( _trains.Count == targetCount && _trains.All( x => x.Consist.Count == targetLength ) ) return;

		ClearTrains();

		if ( targetCount <= 0 ) return;
		if ( TrainCarPrefab is null ) return;

		SetTrains( targetCount, targetLength );
	}

	private void ClearTrains()
	{
		if ( !Networking.IsHost ) return;

		foreach ( var train in _trains )
		{
			train.DestroyGameObject();
		}

		_trains.Clear();
	}

	private void SetTrains( int count, int carsPerTrain )
	{
		if ( !Networking.IsHost ) return;

		if ( count <= 0 )
		{
			ClearTrains();
			return;
		}

		if ( _trains.Count == count && _trains[0].Consist.Count == carsPerTrain )
		{
			return;
		}

		ClearTrains();

		var platform = _platforms.MaxBy( x => x.Length )!;
		var spawnPos = platform.Track.GetDistanceAtNode( platform.MaxElementIndex + 1 ) - 16f;

		for ( var i = 0; i < count; ++i )
		{
			var trainObj = new GameObject( platform.Track.GameObject, true, "Train" );
			var train = trainObj.AddComponent<Train>();

			train.TrackSection = platform.Track;

			_trains.Add( train );

			for ( var j = 0; j < carsPerTrain; ++j )
			{
				var prefab = (j == 0 ? TrainHeadPrefab : j == carsPerTrain - 1 ? TrainTailPrefab : null) ?? TrainCarPrefab;
				if ( !prefab.IsValid() ) continue;

				prefab.Clone( global::Transform.Zero, trainObj );
			}

			train.TrackPosition = spawnPos;
			spawnPos -= train.Length + 32f;

			train.GameObject.NetworkSpawn();
		}
	}

	protected override void OnFixedUpdate()
	{
		base.OnFixedUpdate();

		if ( !Networking.IsHost ) return;

		foreach ( var train in _trains )
		{
			if ( !train.IsValid ) continue;

			switch ( train.State )
			{
				case TrainState.UnloadingGuests:

					foreach ( var slot in train.SlotMarkers.Where( x => x.Contents.IsValid() ) )
					{
						OnUse( slot.Contents );
						Unload( slot.Contents );
					}

					train.State = TrainState.LoadingGuests;
					break;

				case TrainState.LoadingGuests:

					if ( train.StateTime >= MinimumLoadTime && train.SlotMarkers.All( x => !x.IsFree() ) )
					{
						train.State = TrainState.LeavingStation;
					}
					else if ( train.StateTime >= MaximumLoadTime )
					{
						train.State = TrainState.LeavingStation;
					}
					else if ( OpenState == OpenState.Testing && train.StateTime >= 2f )
					{
						train.State = TrainState.LeavingStation;
					}

					break;
			}
		}
	}
}