Rides/TrackRides/TrackDefinition.cs
using System;
using System.Collections.Immutable;

namespace HC3.Rides;

#nullable enable

public enum SupportStyle
{
	Column,
	Scaffold
}

/// <summary>
/// Describes how to generate the mesh for a style of track.
/// </summary>
[AssetType( Name = "Track Definition", Extension = "trkdef" )]
public sealed class TrackDefinition : GameResource, IHotloadManaged
{
	private Model? _rails;
	private Model? _ratchet;
	private Model? _endCap;
	private Model? _strut;
	private Model? _stationPlatform;
	private Model? _verticalSupport;
	private Model? _horizontalSupport;
	private float _sourceGauge = 50f;
	private float _sourceLength = 100f;
	private float _supportHeight = 100f;
	private float _supportBaseHeight = 20f;
	private float _gauge = 24f;

	/// <summary>
	/// Rails and any supports that are fully parallel to the track.
	/// </summary>
	[Group( "Source Models" )]
	public Model? Rails
	{
		get => _rails;
		set
		{
			_rails = value;
			InvalidateNormalized();
		}
	}

	/// <summary>
	/// Support structure that connects the rails together.
	/// </summary>
	[Group( "Source Models" )]
	public Model? Strut
	{
		get => _strut;
		set
		{
			_strut = value;
			InvalidateNormalized();
		}
	}

	/// <summary>
	/// Will go at the end of disconnected track.
	/// </summary>
	[Group( "Source Models" )]
	public Model? EndCap
	{
		get => _endCap;
		set
		{
			_endCap = value;
			InvalidateNormalized();
		}
	}

	/// <summary>
	/// Thing fixed to the track for chain lifts.
	/// </summary>
	[Group( "Source Models" )]
	public Model? Ratchet
	{
		get => _ratchet;
		set
		{
			_ratchet = value;
			InvalidateNormalized();
		}
	}

	/// <summary>
	/// Platforms to place next to the track.
	/// </summary>
	[Group( "Source Models" )]
	public Model? StationPlatform
	{
		get => _stationPlatform;
		set
		{
			_stationPlatform = value;
			InvalidateNormalized();
		}
	}

	/// <summary>
	/// Support column, including base.
	/// </summary>
	[Group( "Source Models" )]
	public Model? VerticalSupport
	{
		get => _verticalSupport;
		set
		{
			_verticalSupport = value;
			InvalidateNormalized();
		}
	}

	/// <summary>
	/// Support beam, for connecting track to support columns.
	/// </summary>
	[Group( "Source Models" )]
	public Model? HorizontalSupport
	{
		get => _horizontalSupport;
		set
		{
			_horizontalSupport = value;
			InvalidateNormalized();
		}
	}

	[Group( "Source Models" )]
	public SupportStyle SupportStyle
	{
		get => _supportStyle;
		set
		{
			_supportStyle = value;
			InvalidateNormalized();
		}
	}

	/// <summary>
	/// Distance between rails in the source model.
	/// </summary>
	[Group( "Source Measurements" )]
	public float SourceGauge
	{
		get => _sourceGauge;
		set
		{
			_sourceGauge = value;
			InvalidateNormalized();
		}
	}

	/// <summary>
	/// Length of track in the source model.
	/// </summary>
	[Group( "Source Measurements" )]
	public float SourceLength
	{
		get => _sourceLength;
		set
		{
			_sourceLength = value;
			InvalidateNormalized();
		}
	}

	/// <summary>
	/// Height of <see cref="VerticalSupport"/>'s source model.
	/// </summary>
	[Group( "Source Measurements" )]
	public float SourceSupportHeight
	{
		get => _supportHeight;
		set
		{
			_supportHeight = value;
			InvalidateNormalized();
		}
	}

	/// <summary>
	/// Height of the base of <see cref="VerticalSupport"/>, so it doesn't get stretched.
	/// </summary>
	[Group( "Source Measurements" )]
	public float SourceSupportBaseHeight
	{
		get => _supportBaseHeight;
		set
		{
			_supportBaseHeight = value;
			InvalidateNormalized();
		}
	}

	/// <summary>
	/// Distance between rails in-game.
	/// </summary>
	[Group( "Game Measurements" )]
	public float Gauge
	{
		get => _gauge;
		set
		{
			_gauge = value;
			InvalidateNormalized();
		}
	}

	/// <summary>
	/// How many <see cref="Strut"/>s per in-game tile.
	/// </summary>
	[Group( "Game Measurements" )]
	public int StrutFrequency { get; set; } = 2;

	/// <summary>
	/// Distance from ground to the top of the track, in game units.
	/// </summary>
	[Group( "Game Measurements" )]
	public float RailHeight { get; set; } = 12f;

	/// <summary>
	/// Try to place a support for each section of track this long, in tiles.
	/// </summary>
	[Group( "Game Measurements" )]
	public int SupportSpacing { get; set; } = 4;

	private readonly record struct NormalizedMesh(
		ImmutableArray<Vertex> Vertices,
		ImmutableArray<int> Indices )
	{
		public bool IsEmpty => Vertices.IsDefaultOrEmpty || Indices.IsDefaultOrEmpty;

		public static NormalizedMesh Empty { get; } = new(
			ImmutableArray<Vertex>.Empty,
			ImmutableArray<int>.Empty );
	}

	private NormalizedMesh? _normalizedRails;
	private NormalizedMesh? _normalizedStrut;
	private NormalizedMesh? _normalizedRatchet;
	private NormalizedMesh? _normalizedStationPlatform;
	private NormalizedMesh? _normalizedVerticalSupport;
	private NormalizedMesh? _normalizedHorizontalSupport;
	private SupportStyle _supportStyle;

	private void InvalidateNormalized()
	{
		_normalizedRails = null;
		_normalizedStrut = null;
		_normalizedRatchet = null;
		_normalizedStationPlatform = null;
		_normalizedVerticalSupport = null;
		_normalizedHorizontalSupport = null;
	}

	public void WriteRail( MeshWriter<Vertex> writer, in Transform prev, in Transform next, TrackFeature feature )
	{
		WriteStretched( writer, Rails, ref _normalizedRails, prev, next );

		switch ( feature )
		{
			case TrackFeature.ChainLift:
				WriteStretched( writer, Ratchet, ref _normalizedRatchet, prev, next );
				break;

			case TrackFeature.Station:
				WriteStretched( writer, StationPlatform, ref _normalizedStationPlatform, prev, next );
				break;
		}
	}

	private void WriteStretched( MeshWriter<Vertex> writer, Model? model, ref NormalizedMesh? normalized,
		Transform prev, Transform next )
	{
		if ( model is null ) return;

		normalized ??= NormalizeTrackModel( model, true );

		WriteMesh( writer, normalized, vertex =>
		{
			var t = vertex.Position.x;
			return Transform.Lerp( prev, next, t, true );
		} );
	}

	public void WriteStrut( MeshWriter<Vertex> writer, Transform transform )
	{
		if ( Strut is not { } strut ) return;

		_normalizedStrut ??= NormalizeTrackModel( strut, false );

		WriteMesh( writer, _normalizedStrut, _ => transform );
	}

	private void WriteMesh( MeshWriter<Vertex> writer, NormalizedMesh? mesh, Func<Vertex, Transform> transformFunc )
	{
		if ( mesh is not { IsEmpty: false } validMesh ) return;

		var indexOffset = writer.VertexCount;

		foreach ( var vertex in validMesh.Vertices )
		{
			var transform = transformFunc( vertex );

			writer.AddVertex( vertex with
			{
				Position = transform.PointToWorld( vertex.Position ),
				Normal = transform.NormalToWorld( vertex.Normal ),
				Tangent = new Vector4( transform.NormalToWorld( (Vector3)vertex.Tangent ), vertex.Tangent.w )
			} );
		}

		writer.AddIndices( validMesh.Indices, indexOffset );
	}

	public const float UnitSupportHeight = 48f;

	public void WriteSupportColumn( MeshWriter<Vertex> writer, in Transform transform, float connectionOffset, float baseZ )
	{
		if ( VerticalSupport is not { } vertModel || HorizontalSupport is not { } horzModel ) return;

		_normalizedVerticalSupport ??= NormalizeVerticalSupportModel( vertModel );
		_normalizedHorizontalSupport ??= NormalizeHorizontalSupportModel( horzModel );

		const float connectorScale = 1f;

		var rotation = Rotation.FromYaw( transform.Rotation.Yaw() );
		var connectionPosition = transform.Position - transform.Up * connectionOffset;

		var baseTransform = new Transform( connectionPosition.WithZ( baseZ ), rotation );
		var topTransform = new Transform( connectionPosition + Vector3.Up * (-UnitSupportHeight + 4f), rotation );
		var connectorTransform = new Transform( (connectionPosition + transform.Position) * 0.5f,
			Rotation.LookAt( transform.Position - connectionPosition ),
			new Vector3( (connectionOffset + 4f) / UnitSupportHeight, connectorScale, connectorScale ) );

		var baseHeight = (SourceSupportBaseHeight / SourceSupportHeight) * UnitSupportHeight;
		var denom = 1f / (UnitSupportHeight - baseHeight);

		WriteMesh( writer, _normalizedVerticalSupport, vertex =>
		{
			var t = (vertex.Position.z - baseHeight) * denom;
			return Transform.Lerp( baseTransform, topTransform, t, true );
		} );

		WriteMesh( writer, _normalizedHorizontalSupport, _ => connectorTransform );
	}

	public void WriteSupportScaffold( MeshWriter<Vertex> writer, Transform prev, Transform next, float baseZ )
	{
		if ( VerticalSupport is not { } vertModel ) return;

		_normalizedVerticalSupport ??= NormalizeScaffoldModel( vertModel );

		const float halfGauge = 12f;

		var topZ = Math.Min(
			Math.Min( prev.PointToWorld( Vector3.Left * halfGauge ).z, prev.PointToWorld( Vector3.Right * halfGauge ).z ),
			Math.Min( next.PointToWorld( Vector3.Left * halfGauge ).z, next.PointToWorld( Vector3.Right * halfGauge ).z ) ) - 8f;

		topZ = Math.Max( baseZ, topZ );

		var prevBase = prev with { Position = prev.Position.WithZ( baseZ ), Rotation = Rotation.FromYaw( prev.Rotation.Yaw() ) };
		var nextBase = next with { Position = next.Position.WithZ( baseZ ), Rotation = Rotation.FromYaw( next.Rotation.Yaw() ) };

		for ( var z = baseZ; z < topZ; z += UnitSupportHeight )
		{
			WriteMesh( writer, _normalizedVerticalSupport, vertex =>
			{
				var u = vertex.Position.x;
				var v = vertex.Position.z;

				var maxZ = Math.Min( topZ, z + UnitSupportHeight );

				var vertexZ = MathX.Lerp( z, maxZ, v );

				var a = prevBase with { Position = prevBase.Position.WithZ( vertexZ ) };
				var b = nextBase with { Position = nextBase.Position.WithZ( vertexZ ) };

				return Transform.Lerp( a, b, u, true );
			} );
		}

		if ( prev.Position.z <= topZ + 2f && next.Position.z <= topZ + 2f ) return;

		prevBase = prevBase with { Position = prevBase.Position.WithZ( topZ ) };
		nextBase = nextBase with { Position = nextBase.Position.WithZ( topZ ) };

		WriteMesh( writer, _normalizedVerticalSupport, vertex =>
		{
			var u = vertex.Position.x;
			var v = vertex.Position.z;

			var a = Transform.Lerp( prevBase, prev, v, true );
			var b = Transform.Lerp( nextBase, next, v, true );

			return Transform.Lerp( a, b, u, true );
		} );
	}

	private NormalizedMesh NormalizeTrackModel( Model model, bool unitLength )
	{
		var scaleYz = Gauge / SourceGauge;
		var scaleX = unitLength ? 1f / SourceLength : scaleYz;
		var transform = new Transform( new Vector3( unitLength ? 0.5f : 0f, 0f, 0f ), Rotation.Identity, new Vector3( scaleX, scaleYz, scaleYz ) );

		// TODO: transform normal / tangent?

		return NormalizeModel( model, x => x with
		{
			Position = transform.PointToWorld( x.Position )
		} );
	}

	private NormalizedMesh NormalizeVerticalSupportModel( Model model )
	{
		var scale = UnitSupportHeight / SourceSupportHeight;

		return NormalizeModel( model, x => x with
		{
			Position = x.Position * scale
		} );
	}

	private NormalizedMesh NormalizeHorizontalSupportModel( Model model )
	{
		var scale = UnitSupportHeight / SourceSupportHeight;

		return NormalizeModel( model, x => x with
		{
			Position = x.Position * scale
		} );
	}

	private NormalizedMesh NormalizeScaffoldModel( Model model )
	{
		var scaleXZ = 1f / SourceSupportHeight;
		var transform = new Transform( new Vector3( 0.5f, 0f, 0f ), Rotation.Identity, new Vector3( scaleXZ, UnitSupportHeight * scaleXZ, scaleXZ ) );

		// TODO: transform normal / tangent?

		return NormalizeModel( model, x => x with
		{
			Position = transform.PointToWorld( x.Position )
		} );
	}

	private NormalizedMesh NormalizeModel( Model model, Func<Vertex, Vertex> func )
	{
		var vertices = model.GetVertices();
		var indices = model.GetIndices();

		var baseVertex = model.GetBaseVertex( 0 );
		var indexStart = model.GetIndexStart( 0 );
		var indexCount = model.GetIndexCount( 0 );

		if ( indexCount == 0 || vertices.Length == 0 )
		{
			return NormalizedMesh.Empty;
		}

		var indicesSlice = indices
			.Skip( indexStart )
			.Take( indexCount )
			.Select( x => (int)x )
			.ToImmutableArray();

		var vertexCount = indicesSlice.Max() + 1;

		var verticesSlice = vertices
			.Skip( baseVertex )
			.Take( vertexCount )
			.Select( func )
			.ToImmutableArray();

		return new NormalizedMesh( verticesSlice, indicesSlice );
	}

	protected override void PostReload()
	{
		InvalidateNormalized();
	}

	void IHotloadManaged.Created( IReadOnlyDictionary<string, object?> state )
	{
		InvalidateNormalized();
	}
}