Terrain/TerrainMesh.cs
using System;

namespace HC3.Terrain;

#nullable enable

public sealed class TerrainMesh : Component, Component.ExecuteInEditor, ITerrainEvent
{
	[RequireComponent]
	public ParkTerrain Terrain { get; private set; } = null!;

	[Property]
	public TerrainPalette? Palette { get; set; }

	[Property]
	public Material? GroundMaterial { get; set; }

	[Property]
	public Material? WallMaterial { get; set; }

	public int ChunkSize => 32;

	private readonly HashSet<Vector2Int> _invalidChunks = new();
	private readonly Dictionary<Vector2Int, TerrainMeshChunk> _chunks = new();

	protected override void OnEnabled()
	{
		Terrain_Changed( Terrain.Bounds );

		// We should set this on each chunk's ModelRenderer instead,
		// but there's no nice way of doing that without waiting for
		// it to create a SceneObject

		Palette?.Apply( Scene.RenderAttributes );
	}

	protected override void OnDestroy()
	{
		foreach ( var chunk in _chunks.Values )
		{
			chunk.Dispose();
		}

		_chunks.Clear();
	}

	private void Terrain_Changed( RectInt tileRange )
	{
		foreach ( var chunkIndex in tileRange.Grow( 1 ).Clamp( Terrain.Bounds ).GetChunkIndices( ChunkSize ) )
		{
			_invalidChunks.Add( chunkIndex );
		}
	}

	protected override void OnUpdate()
	{
		if ( _invalidChunks.Count == 0 ) return;

		ITerrainData terrain = Terrain;

		foreach ( var chunkIndex in _invalidChunks )
		{
			if ( !_chunks.TryGetValue( chunkIndex, out var chunk ) )
			{
				_chunks[chunkIndex] = chunk = new TerrainMeshChunk( this, chunkIndex );
			}

			chunk.Update( terrain );
		}

		_invalidChunks.Clear();
	}

	void ITerrainEvent.Changed( ParkTerrain terrain, RectInt tileRange ) =>
		Terrain_Changed( tileRange );
}

public sealed class TerrainMeshChunk : IDisposable
{
	public TerrainMesh Parent { get; }
	public Vector2Int Index { get; }
	public RectInt TileRange { get; }
	public BBox WorldRange { get; }

	private int Margin => 1;

	private Transform LocalTransform => new( new Vector3( TileRange.Position * Parent.Terrain.TileWidth, 0f ) );

	private readonly TileArraySlice _dataCopy;

	private GameObject? _gameObject;
	private ModelRenderer? _renderer;
	private PhysicsBody? _body;
	private PhysicsShape? _shape;

	private Model? _model;
	private Mesh? _groundMesh;
	private Mesh? _wallMesh;

	private bool _hadGround;
	private bool _hadWalls;

	public TerrainMeshChunk( TerrainMesh parent, Vector2Int index )
	{
		Parent = parent;
		Index = index;
		TileRange = index.GetTileRange( parent.ChunkSize );

		var scale = new Vector3( parent.Terrain.TileWidth, parent.Terrain.TileWidth, parent.Terrain.TileHeight );
		var min = new Vector3( TileRange.Position, 0f ) * scale;
		var max = new Vector3( TileRange.Position + TileRange.Size, ushort.MaxValue + 1 ) * scale;

		WorldRange = new BBox( min, max );

		_dataCopy = new TileArraySlice( parent.ChunkSize + Margin * 2 );
	}

	public void Update( ITerrainData data )
	{
		data.CopyToClamped( TileRange.Grow( Margin ), _dataCopy );

		UpdateRenderMesh( _dataCopy );
		UpdatePhysicsMesh( _dataCopy );
	}

	public void Dispose()
	{
		if ( _body.IsValid() )
		{
			_body.Remove();
		}

		_body = null;
	}

	[field: ThreadStatic]
	private static GroundMeshWriter? GroundWriter { get; set; }

	[field: ThreadStatic]
	private static WallMeshWriter? WallWriter { get; set; }

	private void UpdateRenderMesh( TileArraySlice data )
	{
		var groundWriter = GroundWriter ??= new GroundMeshWriter();
		var wallWriter = WallWriter ??= new WallMeshWriter();

		groundWriter.Clear();
		groundWriter.WriteTiles( data, Parent.Terrain.TileWidth, Parent.Terrain.TileHeight );

		var hasGround = !groundWriter.IsEmpty;

		if ( hasGround )
		{
			if ( !_groundMesh.IsValid() )
			{
				_groundMesh = new Mesh( Parent.GroundMaterial );
			}

			groundWriter.CopyTo( _groundMesh );
		}

		wallWriter.Clear();
		wallWriter.WriteTiles( data, Parent.Terrain.TileWidth, Parent.Terrain.TileHeight );

		var hasWalls = !wallWriter.IsEmpty;

		if ( hasWalls )
		{
			if ( !_wallMesh.IsValid() )
			{
				_wallMesh = new Mesh( Parent.WallMaterial );
			}

			wallWriter.CopyTo( _wallMesh );
		}

		if ( !hasWalls && !hasGround )
		{
			if ( _gameObject.IsValid() )
			{
				_gameObject.Enabled = false;
			}

			return;
		}

		if ( !_model.IsValid() || hasGround != _hadGround || hasWalls != _hadWalls )
		{
			var builder = new ModelBuilder();

			if ( hasGround )
			{
				builder.AddMesh( _groundMesh );
			}

			if ( hasWalls )
			{
				builder.AddMesh( _wallMesh );
			}

			_model = builder.Create();

			_hadGround = hasGround;
			_hadWalls = hasWalls;
		}

		if ( !_gameObject.IsValid() )
		{
			_gameObject = new GameObject( Parent.GameObject, false, $"Chunk ({Index})" )
			{
				Flags = GameObjectFlags.NotNetworked | GameObjectFlags.NotSaved,
				LocalTransform = LocalTransform
			};
		}

		if ( !_renderer.IsValid() )
		{
			_renderer = _gameObject.GetOrAddComponent<ModelRenderer>();
		}

		_renderer.Model = _model;
		_gameObject.Enabled = true;
	}

	[field: ThreadStatic]
	private static PhysicsMeshWriter? PhysicsWriter { get; set; }

	private void UpdatePhysicsMesh( TileArraySlice data )
	{
		var physicsWriter = PhysicsWriter ??= new PhysicsMeshWriter();

		physicsWriter.Clear();
		physicsWriter.WriteTiles( data, Parent.Terrain.TileWidth, Parent.Terrain.TileHeight );

		if ( physicsWriter.IsEmpty )
		{
			if ( _body.IsValid() )
			{
				_body.Enabled = false;
			}

			return;
		}

		if ( !_body.IsValid() )
		{
			_body = new PhysicsBody( Parent.Scene.PhysicsWorld ) { BodyType = PhysicsBodyType.Static };
			_body.Component = Parent;
		}

		physicsWriter.UpdateShape( _body, ref _shape );

		_body.Enabled = true;
		_body.Transform = Parent.GameObject.WorldTransform.ToWorld( LocalTransform );

		// We're using A* now, so no need for this
		// Parent.Scene.NavMesh.GenerateTiles( Parent.Scene.PhysicsWorld, WorldRange );
	}

	private abstract class TileMeshWriter<T> : MeshWriter<T>
		where T : unmanaged
	{
		private bool _first = true;

		protected override void OnClear()
		{
			_first = true;
		}

		public void WriteTiles( TileArraySlice data, float tileWidth, float tileHeight )
		{
			var size = data.Size - 2;

			if ( _first )
			{
				Bounds = new BBox( Vector3.Zero, new Vector3( size, 0f ) * tileWidth );
			}

			for ( var y = 0; y < size.y; ++y )
			{
				for ( var x = 0; x < size.x; ++x )
				{
					var origin = new Vector3( x, y, 0f ) * tileWidth;
					var neighborhood = data.Slice( new RectInt( x, y, 3, 3 ) );
					var tile = neighborhood[new Vector2Int( 1, 1 )];

					OnWriteTile( neighborhood, origin, tileWidth, tileHeight );

					Bounds = Bounds.AddBBox( new BBox( origin, origin + new Vector3( tileWidth, tileWidth, tile.MaxHeight * tileHeight ) ) );
				}
			}
		}

		protected abstract void OnWriteTile( TileArraySlice neighborhood, Vector3 origin, float tileWidth, float tileHeight );

		protected void AddWalls( TileArraySlice neighborhood, Vector3 origin, float tileWidth, float tileHeight )
		{
			var yNeg = neighborhood[new Vector2Int( 1, 0 )];
			var xNeg = neighborhood[new Vector2Int( 0, 1 )];
			var tile = neighborhood[new Vector2Int( 1, 1 )];
			var xPos = neighborhood[new Vector2Int( 2, 1 )];
			var yPos = neighborhood[new Vector2Int( 1, 2 )];

			var min = origin;
			var max = origin + new Vector3( tileWidth, tileWidth, tile.BaseHeight * tileHeight );

			AddWall(
				tile.GetCornerHeight( TileCorner.XMinYMin ), tile.GetCornerHeight( TileCorner.XMinYMax ),
				xNeg.GetCornerHeight( TileCorner.XMaxYMin ), xNeg.GetCornerHeight( TileCorner.XMaxYMax ),
				new Vector3( min.x, min.y, min.z ), new Vector3( 0f, 1f, 0f ), tileWidth, tileHeight );

			AddWall(
				tile.GetCornerHeight( TileCorner.XMinYMax ), tile.GetCornerHeight( TileCorner.XMaxYMax ),
				yPos.GetCornerHeight( TileCorner.XMinYMin ), yPos.GetCornerHeight( TileCorner.XMaxYMin ),
				new Vector3( min.x, max.y, min.z ), new Vector3( 1f, 0f, 0f ), tileWidth, tileHeight );

			AddWall(
				tile.GetCornerHeight( TileCorner.XMaxYMax ), tile.GetCornerHeight( TileCorner.XMaxYMin ),
				xPos.GetCornerHeight( TileCorner.XMinYMax ), xPos.GetCornerHeight( TileCorner.XMinYMin ),
				new Vector3( max.x, max.y, min.z ), new Vector3( 0f, -1f, 0f ), tileWidth, tileHeight );

			AddWall(
				tile.GetCornerHeight( TileCorner.XMaxYMin ), tile.GetCornerHeight( TileCorner.XMinYMin ),
				yNeg.GetCornerHeight( TileCorner.XMaxYMax ), yNeg.GetCornerHeight( TileCorner.XMinYMax ),
				new Vector3( max.x, min.y, min.z ), new Vector3( -1f, 0f, 0f ), tileWidth, tileHeight );
		}

		private void AddWall( int thisHeight0, int thisHeight1, int thatHeight0, int thatHeight1, Vector3 origin, Vector3 tangent, float tileWidth, float tileHeight )
		{
			if ( thisHeight0 <= thatHeight0 && thisHeight1 <= thatHeight1 ) return;

			if ( thisHeight0 < thatHeight0 || thisHeight1 < thatHeight1 )
			{
				thatHeight0 = thatHeight1 = Math.Min( thatHeight0, thatHeight1 );
			}

			OnAddWall( thisHeight0, thisHeight1, thatHeight0, thatHeight1, origin, tangent, tileWidth, tileHeight );
		}

		protected virtual void OnAddWall( int thisHeight0, int thisHeight1, int thatHeight0, int thatHeight1, Vector3 origin, Vector3 tangent, float tileWidth, float tileHeight )
		{
			throw new NotImplementedException();
		}

		[SkipHotload]

		private static TileCorner[] DefaultTriangulation { get; } =
		[
			TileCorner.XMinYMin, TileCorner.XMaxYMin, TileCorner.XMaxYMax,
			TileCorner.XMinYMin, TileCorner.XMaxYMax, TileCorner.XMinYMax
		];
		[SkipHotload]

		private static TileCorner[] FlippedTriangulation { get; } =
		[
			TileCorner.XMaxYMin, TileCorner.XMaxYMax, TileCorner.XMinYMax,
			TileCorner.XMaxYMin, TileCorner.XMinYMax, TileCorner.XMinYMin
		];

		protected static TileCorner[] GetGroundTriangulation( TerrainTile tile )
		{
			return tile.Slope.Triangulation is TileTriangulation.Default ? DefaultTriangulation : FlippedTriangulation;
		}
	}

	private abstract class RenderMeshWriter : TileMeshWriter<Vertex>;

	private sealed class GroundMeshWriter : RenderMeshWriter
	{
		protected override void OnWriteTile( TileArraySlice neighborhood, Vector3 origin, float width, float height )
		{
			var tile = neighborhood[new Vector2Int( 1, 1 )];

			var min = origin;
			var max = origin + new Vector3( width, width, tile.BaseHeight * height );

			AddGround( tile, min, max, height );
		}

		private int AddVertex( Vector3 pos, Vector2 uv, Vector3 normal, Vector4 tangent, TilePaint paint )
		{
			var vertex = new Vertex( position: pos, texCoord0: uv, normal: normal, tangent: tangent )
			{
				Color = new Color32(
					(byte)paint.GetMaterialIndex( TileCorner.XMinYMin ),
					(byte)paint.GetMaterialIndex( TileCorner.XMaxYMin ),
					(byte)paint.GetMaterialIndex( TileCorner.XMinYMax ),
					(byte)paint.GetMaterialIndex( TileCorner.XMaxYMax ) )
			};

			return AddVertex( vertex );
		}

		private void AddGround( TerrainTile tile, Vector2 min, Vector2 max, float tileHeight )
		{
			var triangulation = GetGroundTriangulation( tile );

			for ( var i = 0; i < triangulation.Length; i += 3 )
			{
				var c0 = triangulation[i + 0];
				var c1 = triangulation[i + 1];
				var c2 = triangulation[i + 2];

				var uv0 = c0.GetHorizontalOffset();
				var uv1 = c1.GetHorizontalOffset();
				var uv2 = c2.GetHorizontalOffset();

				var p0 = new Vector3( min + uv0 * (max - min), tile.GetCornerHeight( c0 ) * tileHeight );
				var p1 = new Vector3( min + uv1 * (max - min), tile.GetCornerHeight( c1 ) * tileHeight );
				var p2 = new Vector3( min + uv2 * (max - min), tile.GetCornerHeight( c2 ) * tileHeight );

				var normal = Vector3.Cross( p1 - p0, p2 - p0 ).Normal;
				var tangent = new Vector4( Vector3.Cross( p1 - p0, normal ).Normal, 1f );

				var v0 = AddVertex( p0, uv0, normal, tangent, tile.Paint );
				var v1 = AddVertex( p1, uv1, normal, tangent, tile.Paint );
				var v2 = AddVertex( p2, uv2, normal, tangent, tile.Paint );

				AddTriangle( v0, v1, v2 );
			}
		}
	}

	private sealed class WallMeshWriter : RenderMeshWriter
	{
		protected override void OnWriteTile( TileArraySlice neighborhood, Vector3 origin, float tileWidth, float tileHeight )
		{
			AddWalls( neighborhood, origin, tileWidth, tileHeight );
		}

		private int AddVertex( Vector3 pos, Vector2 uv, Vector3 normal, Vector4 tangent ) =>
			AddVertex( new Vertex( position: pos, texCoord0: uv, normal: normal, tangent: tangent ) );

		protected override void OnAddWall( int thisHeight0, int thisHeight1, int thatHeight0, int thatHeight1, Vector3 origin, Vector3 tangent, float tileWidth, float tileHeight )
		{
			var normal = Vector3.Cross( Vector3.Up, tangent );
			var vUnit = tileHeight / tileWidth;

			var zUnit = Vector3.Up * tileHeight;

			var tangent4 = new Vector4( tangent, 1f );

			var v0 = AddVertex( origin + thatHeight0 * zUnit, new Vector2( 0f, thatHeight0 * vUnit ), normal, tangent4 );
			var v1 = AddVertex( origin + tangent * tileWidth + thatHeight1 * zUnit, new Vector2( 1f, thatHeight1 * vUnit ), normal, tangent4 );
			var v2 = AddVertex( origin + thisHeight0 * zUnit, new Vector2( 0f, thisHeight0 * vUnit ), normal, tangent4 );
			var v3 = AddVertex( origin + tangent * tileWidth + thisHeight1 * zUnit, new Vector2( 1f, thisHeight1 * vUnit ), normal, tangent4 );

			AddTriangle( v0, v3, v1 );
			AddTriangle( v0, v2, v3 );
		}
	}

	private sealed class PhysicsMeshWriter : TileMeshWriter<Vector3>
	{
		public void UpdateShape( PhysicsBody body, ref PhysicsShape? shape )
		{
			if ( shape.IsValid() )
			{
				shape.UpdateMesh( Vertices, Indices );
			}
			else
			{
				shape = body.AddMeshShape( Vertices, Indices );

				// TODO: copy from game object?
				shape.Tags.Add( "ground" );
			}
		}

		protected override void OnWriteTile( TileArraySlice neighborhood, Vector3 origin, float tileWidth, float tileHeight )
		{
			var tile = neighborhood[new Vector2Int( 1, 1 )];
			var triangulation = GetGroundTriangulation( tile );

			var min = new Vector2( origin.x, origin.y );
			var max = min + tileWidth;

			for ( var i = 0; i < triangulation.Length; i += 3 )
			{
				var c0 = triangulation[i + 0];
				var c1 = triangulation[i + 1];
				var c2 = triangulation[i + 2];

				var uv0 = c0.GetHorizontalOffset();
				var uv1 = c1.GetHorizontalOffset();
				var uv2 = c2.GetHorizontalOffset();

				var p0 = new Vector3( min + uv0 * (max - min), tile.GetCornerHeight( c0 ) * tileHeight );
				var p1 = new Vector3( min + uv1 * (max - min), tile.GetCornerHeight( c1 ) * tileHeight );
				var p2 = new Vector3( min + uv2 * (max - min), tile.GetCornerHeight( c2 ) * tileHeight );

				var v0 = AddVertex( p0 );
				var v1 = AddVertex( p1 );
				var v2 = AddVertex( p2 );

				AddTriangle( v0, v1, v2 );
			}

			AddWalls( neighborhood, origin, tileWidth, tileHeight );
		}

		protected override void OnAddWall( int thisHeight0, int thisHeight1, int thatHeight0, int thatHeight1, Vector3 origin, Vector3 tangent, float tileWidth, float tileHeight )
		{
			var zUnit = Vector3.Up * tileHeight;

			var v0 = AddVertex( origin + thatHeight0 * zUnit );
			var v1 = AddVertex( origin + tangent * tileWidth + thatHeight1 * zUnit );
			var v2 = AddVertex( origin + thisHeight0 * zUnit );
			var v3 = AddVertex( origin + tangent * tileWidth + thisHeight1 * zUnit );

			AddTriangle( v0, v3, v1 );
			AddTriangle( v0, v2, v3 );
		}
	}
}