Rides/TrackRides/TrackBuilder.cs
using HC3.UI;
using System;
using System.Collections.Immutable;
using HC3.Terrain;

namespace HC3.Rides;

#nullable enable

/// <summary>
/// Are we building at the start or end of a track section?
/// </summary>
public enum TrackBuildKind
{
	Append,
	Prepend
}

/// <summary>
/// Building tool for <see cref="TrackRide"/>.
/// </summary>
public sealed class TrackBuilder : Component
{
	public static TrackBuilder? Instance { get; private set; } = null!;

	public bool IsBuilding => SelectedSection.IsValid();

	public TrackRide? SelectedRide { get; private set; }
	public TrackSection? SelectedSection { get; private set; }
	public TrackBuildKind BuildKind => TrackBuildKind.Append;

	public IEnumerable<TrackElement> PossibleElements =>
		GetPossibleElements( Filter );

	public IEnumerable<TrackElementDefinition> PossibleSpecialElements =>
		GetPossibleElements( Filter with { Feature = TrackFeature.Special } )
			.Select( x => x.Definition )
			.Distinct();

	public bool CanBuild => SelectedSection is { } section
		&& PreviewElement is { } element
		&& CalculateBuildCost( SelectedRide, section.Nodes[^1], element, section.WorldTransform, true ) <= ParkManager.Instance?.Money;

	public TrackNode LastNode => SelectedSection?.Nodes[BuildKind == TrackBuildKind.Append ? ^1 : 0] ?? default;

	public TrackElement? LastElement => SelectedSection is { Elements.Count: > 0 } section
		? section.Elements[BuildKind == TrackBuildKind.Append ? ^1 : 0]
		: null;

	private TrackElementFilter _filter;
	private TrackElementSize _size;
	private TrackElementDefinition? _special;
	private int _lastHash;

	public TrackElementFilter Filter
	{
		get => _filter;
		set
		{
			_filter = value;
			UpdatePreview();
		}
	}

	public TrackElementSize Size
	{
		get => _size;
		set
		{
			_size = value;
			UpdatePreview();
		}
	}

	public TrackElementDefinition? SpecialElement
	{
		get => _special;
		set
		{
			_special = value;
			UpdatePreview();
		}
	}

	private ImmutableDictionary<TrackElementFilter, ImmutableArray<TrackElement>> _allElements =
		ImmutableDictionary<TrackElementFilter, ImmutableArray<TrackElement>>.Empty;

	public IEnumerable<TurnDirection> AllDirections =>
		_allElements.Keys
			.Select( x => x.Direction )
			.Distinct()
			.Order();

	public IEnumerable<TrackBanking> AllBanking =>
		_allElements.Keys
			.Select( x => x.Banking )
			.Distinct()
			.Order();

	public IEnumerable<TrackIncline> AllInclines =>
		_allElements.Keys
			.Select( x => x.Incline )
			.Distinct()
			.Order();

	public IEnumerable<TrackElementSize> AllSizes =>
		PossibleElements
			.Select( x => x.Definition.Size )
			.Distinct()
			.Order();

	private TrackSection? _preview;

	public IEnumerable<TrackElement> AllPossibleElements =>
		_allElements.Values
			.SelectMany( x => x )
			.Where( x => LastNode.Append( x ) is not null );

	public IEnumerable<TrackElement> GetPossibleElements( TrackElementFilter filter ) =>
		_allElements.TryGetValue( filter, out var filtered )
			? filtered.Where( x => LastNode.Append( x ) is not null )
			: Enumerable.Empty<TrackElement>();

	public IEnumerable<TrackElement> GetPossibleElements( Func<TrackElementFilter, bool> filter ) =>
		_allElements.Keys.Where( filter )
			.SelectMany( GetPossibleElements );

	public IEnumerable<TrackElementFilter> PossibleFilters =>
		_allElements.Keys.Where( x => GetPossibleElements( x ).Any() )
			.OrderBy( x => x.GetDistance( Filter ) );

	public bool CanFilter( TurnDirection direction ) =>
		GetPossibleElements( x => x.Direction == direction ).Any();
	public bool CanFilter( TrackBanking banking ) =>
		GetPossibleElements( x => x.Banking == banking ).Any();
	public bool CanFilter( TrackIncline incline ) =>
		GetPossibleElements( x => x.Incline == incline ).Any();
	public bool CanFilter( TrackFeature feature ) =>
		GetPossibleElements( x => x.Feature == feature ).Any();

	public void SetFilter( TurnDirection direction ) =>
		Filter = PossibleFilters.FirstOrDefault( x => x.Direction == direction );

	public void SetFilter( TrackBanking banking ) =>
		Filter = PossibleFilters.FirstOrDefault( x => x.Banking == banking );

	public void SetFilter( TrackIncline incline ) =>
		Filter = PossibleFilters.FirstOrDefault( x => x.Incline == incline );

	public void SetFilter( TrackFeature feature ) =>
		Filter = PossibleFilters.FirstOrDefault( x => x.Feature == feature );

	public void ToggleFilter( TrackFeature feature ) =>
		SetFilter( Filter.Feature == feature ? TrackFeature.None : feature );

	public TrackElement? PreviewElement =>
		PossibleElements.FirstOrDefault( x => Filter.Feature == TrackFeature.Special ? x.Definition == _special : x.Size == Size ) is { Definition: not null } element
			? element
			: null;

	protected override void OnAwake()
	{
		Instance = this;
	}

	protected override void OnDestroy()
	{
		ClearPreview();

		if ( Instance == this )
		{
			Instance = null!;
		}
	}

	public void ClearSelection()
	{
		ClearPreview();

		SelectedRide = null;
		SelectedSection = null;
	}

	public bool StartBuilding( TrackRide ride )
	{
		if ( ride.TrackSections.FirstOrDefault() is not { } section )
		{
			Log.Warning( "Ride has no track sections!" );
			return false;
		}

		if ( section.Nodes.Count < 2 )
		{
			Log.Warning( "Ride section is invalid!" );
			return false;
		}

		StartBuilding( ride, section, TrackBuildKind.Append );
		return true;
	}

	public void StartBuilding( TrackRide? ride, TrackSection section, TrackBuildKind kind )
	{
		SelectedRide = ride;
		SelectedSection = section;

		if ( ride is not null )
		{
			ride.OpenState = OpenState.Closed;
		}

		var constraints = section.GetComponentInParent<ITrackConstraints>();

		_allElements = FindAllElements( constraints )
			.Where( x => x.Definition.Feature != TrackFeature.Disabled )
			.GroupBy( x => TrackElementFilter.FromElement( x, kind ) )
			.ToImmutableDictionary( x => x.Key, x => x.Order().ToImmutableArray() );

		_filter = LastElement is { } element
			? TrackElementFilter.FromElement( element, kind )
			: _allElements.Keys.Min();

		_size = PossibleElements
			.Select( x => x.Size )
			.DefaultIfEmpty( TrackElementSize.None )
			.Min();

		_special = PossibleElements
			.Where( x => x.Feature == TrackFeature.Special )
			.Select( x => x.Definition )
			.FirstOrDefault();

		UpdatePreview();

		WindowManager.Instance.Open( new TrackBuilderWindow() );
	}

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

		var hash = HashCode.Combine( SelectedSection?.Elements.Count, ParkManager.Instance?.Money );

		if ( hash != _lastHash )
		{
			_lastHash = hash;
			UpdatePreview();
		}
	}

	private static IEnumerable<TrackElement> FindAllElements( ITrackConstraints constraints )
	{
		return
			from elemDef in ResourceLibrary.GetAll<TrackElementDefinition>()
			from element in elemDef.Elements
			where constraints?.IsElementAllowed( element ) ?? true
			select element;
	}

	private void ClearPreview()
	{
		_preview?.DestroyGameObject();
	}

	private void UpdatePreview()
	{
		if ( SelectedSection is not { } source || !IsBuilding || source.IsCycle )
		{
			ClearPreview();
			return;
		}

		if ( !PossibleElements.Any() )
		{
			_filter = PossibleFilters.FirstOrDefault();
		}

		if ( _filter.Feature == TrackFeature.Special )
		{
			if ( !PossibleSpecialElements.Contains( _special ) )
			{
				_special = PossibleSpecialElements.FirstOrDefault();
			}
		}
		else if ( PossibleElements.All( x => x.Size != Size ) )
		{
			_size = PossibleElements.Select( x => x.Size )
				.DefaultIfEmpty( TrackElementSize.None )
				.Min();
		}

		if ( PreviewElement is not { } element )
		{
			ClearPreview();
			return;
		}

		if ( !_preview.IsValid() )
		{
			var previewObj = new GameObject( true, "Track Preview" )
			{
				WorldTransform = source.WorldTransform,
				Flags = GameObjectFlags.NotNetworked | GameObjectFlags.NotSaved
			};

			_preview = previewObj.AddComponent<TrackSection>();
			previewObj.AddComponent<TrackMesh>();
		}

		_preview.TrackDefinition = source.TrackDefinition;
		_preview.WorldTransform = source.WorldTransform;

		if ( _preview.GetComponent<ModelRenderer>()?.SceneObject is { } sceneObj )
		{
			sceneObj.Attributes.Set( "Ghost", CanBuild ? 1 : 2 );
			sceneObj.Batchable = false;
		}

		_preview.SetElements( LastNode, [element] );
	}

	public void Build()
	{
		if ( SelectedSection is not { } section ) return;
		if ( PreviewElement is not { } element ) return;

		BuildCore( SelectedRide, section, element, TrackBuildKind.Append );
	}

	[Rpc.Host]
	private void BuildCore( TrackRide? ride, TrackSection section, TrackElement.RpcSafe element, TrackBuildKind kind )
	{
		if ( kind != TrackBuildKind.Append ) throw new NotImplementedException();

		if ( CalculateBuildCost( ride, section.Nodes[^1], element, section.WorldTransform, true ) is not { } cost ) return;
		if ( !ParkManager.Instance?.TakeMoney( cost, "Rollercoasters" ) ?? false ) return;
		if ( !section.Build( element, kind ) ) return;

		BuildEffect( section.Nodes[^2], element, section.WorldTransform, cost );
	}

	public void Demolish()
	{
		if ( SelectedSection is not { Elements.Count: > 1 } section ) return;

		DemolishCore( SelectedRide, section, BuildKind );
	}

	[Rpc.Host]
	private void DemolishCore( TrackRide? ride, TrackSection section, TrackBuildKind kind )
	{
		if ( kind != TrackBuildKind.Append ) throw new NotImplementedException();

		var node = section.Nodes[^2];
		var element = section.Elements[^1];

		if ( CalculateBuildCost( ride, node, element, section.WorldTransform, false ) is not { } cost ) return;

		var refund = Math.Min( (cost * 3 / 4).SnapToGrid( 10 ), cost );

		if ( !section.Demolish( kind ) ) return;

		BuildEffect( node, element, section.WorldTransform, -refund );

		ParkManager.Instance?.GiveMoney( refund, "Rollercoasters" );
	}

	private int? CalculateBuildCost( TrackRide? ride, TrackNode prev, TrackElement element, Transform worldTransform, bool includeElevation )
	{
		if ( prev.Append( element ) is not { } next ) return null;

		var spline = new Spline();

		if ( !spline.AddTrack( prev, [element] ) ) return null;

		var length = (int)MathF.Round( spline.Length / GridManager.GridSize );
		var elevation = includeElevation ? GetElevation( spline, worldTransform ) : default; // TODO: use minElement, count

		// Can't build underground yet

		if ( elevation.Min < 0 ) return null;
		if ( elevation.Max > ride?.MaxElevation ) return null;

		// Check for obstructions

		if ( GridManager.Instance is { } grid && HasGridObstructions( grid, spline, worldTransform ) ) return null;

		// Track needs special handling

		var nodes = new HashSet<TrackNode> { prev, next };

		foreach ( var track in Scene.GetAllComponents<TrackSection>() )
		{
			if ( track == _preview ) continue;

			if ( HasTrackObstructions( new OverlappingTrack( track, worldTransform.ToLocal( track.WorldTransform ) ), spline, nodes ) )
			{
				return null;
			}
		}

		return length * (ride?.BaseTrackCost ?? 0) + elevation.Total * (ride?.ElevationCost ?? 0);
	}

	/// <summary>
	/// For each tile-sized segment of track, sum how high from the ground it is.
	/// </summary>
	private static (int Min, int Max, int Total) GetElevation( Spline spline, Transform worldTransform )
	{
		if ( GridManager.Instance?.Terrain is not { IsValid: true } terrain ) return default;

		var tiles = Math.Max( 1, (int)MathF.Round( spline.Length / GridManager.GridSize ) );
		var tileSize = spline.Length / tiles;

		var min = int.MaxValue;
		var max = int.MinValue;
		var total = 0;

		for ( var i = 0; i <= tiles; ++i )
		{
			var t = i * tileSize;
			var sample = spline.SampleAtDistance( t );

			var worldPos = worldTransform.PointToWorld( sample.Position );
			var gridPos = terrain.WorldToGrid( worldPos );
			var tile = terrain[terrain.GetTileIndex( (Vector2)gridPos )];

			var elevation = (int)MathF.Floor( gridPos.z - (tile.MinHeight + tile.MaxHeight) * 0.5f );

			min = Math.Min( min, elevation );
			max = Math.Max( max, elevation );
			total += elevation;
		}

		return (min, max, total);
	}

	private static bool HasGridObstructions( GridManager grid, Spline spline, Transform worldTransform )
	{
		const float sampleResolution = 16f;

		// TODO: allow building through scenery if we have enough money

		for ( var t = sampleResolution * 0.5f; t <= spline.Length; t += sampleResolution )
		{
			var pos = worldTransform.PointToWorld( spline.SampleAtDistance( t ).Position );

			if ( grid.GetCell( GridManager.WorldToGridPosition( pos ) ) is not { } cell ) continue;

			foreach ( var obj in cell )
			{
				if ( !obj.BlocksConstruction ) continue;

				if ( obj.WorldPosition.z >= pos.z + 40f ) continue;
				if ( obj.WorldPosition.z + obj.Height * GridManager.HeightStep <= pos.z ) continue;

				return true;
			}
		}

		return false;
	}

	private bool HasTrackObstructions( OverlappingTrack overlap, Spline spline, HashSet<TrackNode>? nodes )
	{
		const float sampleResolution = 16f;

		var ignoreNodes = overlap.Transform == global::Transform.Zero ? nodes : null;

		for ( var t = sampleResolution * 0.5f; t <= spline.Length; t += sampleResolution )
		{
			var sample = spline.SampleAtDistance( t );

			var pos = overlap.Transform.PointToLocal( sample.Position + sample.Up * 8f );
			var gridPos = (pos / TrackSection.GridSize).FloorToInt();

			if ( overlap.Track.HasObstruction( gridPos, ignoreNodes ) ) return true;
		}

		return false;
	}

	[Rpc.Broadcast( NetFlags.HostOnly )]
	private void BuildEffect( TrackNode prev, TrackElement.RpcSafe element, Transform worldTransform, int cost )
	{
		if ( cost == 0 ) return;

		var spline = new Spline();

		if ( !spline.AddTrack( prev, [element] ) ) return;

		var length = spline.Length;

		var centerSample = spline.SampleAtDistance( length * 0.5f );
		var pos = worldTransform.PointToWorld( centerSample.Position + centerSample.Up * 32f );

		Sound.Play( "sounds/gameplay/building_placed.sound", pos );

		for ( var t = 32f; t < length; t += 64f )
		{
			var sample = spline.SampleAtDistance( t );

			GameObject.Clone( "prefabs/particles/place_dustcloud.prefab",
				worldTransform.ToWorld( new Transform( sample.Position, Rotation.LookAt( sample.Tangent ) ) ) );
		}

		if ( cost > 0 )
		{
			MoneyEffect.Broadcast( pos + Vector3.Up * 32f, $"-${cost}", Color.Red );
		}
		else
		{
			MoneyEffect.Broadcast( pos + Vector3.Up * 32f, $"${-cost}", Color.Green );
		}
	}
}