using Sandbox;
using Sandbox.Rendering;
using System;
using System.Collections.Generic;
using System.Linq;

namespace Saandy.Tilemapper;

public sealed class TilemapRenderObject : SceneCustomObject
{
	public struct TileData
	{
		public Vector4 Position;
		public Vector4 UvRect; // xy = offset, zw = scale
	}

	public struct ChunkCoord : IEquatable<ChunkCoord>
	{
		public int X;
		public int Y;

		public ChunkCoord( int x, int y )
		{
			X = x;
			Y = y;
		}

		public bool Equals( ChunkCoord other ) => X == other.X && Y == other.Y;
		public override bool Equals( object obj ) => obj is ChunkCoord other && Equals( other );
		public override int GetHashCode() => HashCode.Combine( X, Y );
	}

	public sealed class TilemapChunk
	{
		public ChunkCoord Coord;
		public List<TilemapChunkBatch> Batches = new();
		public BBox Bounds;
		public int Revision;
	}

	public sealed class TilemapChunkBatch
	{
		public int LayerIndex;
		public ushort TilesetId;
		public GpuBuffer ArgsBuffer;
		public GpuBuffer<TileData> TileDataBuffer;
		public Material Material;

		public void Dispose()
		{
			TileDataBuffer?.Dispose();
			ArgsBuffer?.Dispose();

			TileDataBuffer = null;
			ArgsBuffer = null;
		}
	}

	private readonly Material _tileMaterial;
	private readonly Model _tileModel;
	private readonly Vector3 _tileModelCenter;
	private readonly float _tileModelScale;
	private readonly Queue<TilemapChunk> _pendingDestroy = new();
	private int _tilemapRevision = -1;
	private TileMapAxis _tilemapAxis = (TileMapAxis)(-1);
	private int _tilemapLayerCount = -1;
	private float _tilemapLayerSpacing = -1.0f;

	public readonly Dictionary<ChunkCoord, TilemapChunk> ActiveChunks = new();
	public TileMap Tilemap { get; private set; }
	public CameraComponent CullingCamera { get; set; }
	public float TileSize { get; set; } = 1.0f;
	public float ChunkSize { get; set; } = 256.0f;
	public int ChunkResolution { get; set; } = 64;
	public int RenderRadius { get; set; } = 6;
	public TileMapAxis Axis { get; set; } = TileMapAxis.XZ;

	private float ChunkWorldSize => Math.Max( ChunkSize, Math.Max( TileSize, 0.0001f ) );

	public TilemapRenderObject( SceneWorld sceneWorld ) : base( sceneWorld )
	{
		_tileMaterial = Material.FromShader( "shaders/tilemap.shader" );
		_tileModel = LoadTileModel();

		BBox tileBounds = _tileModel.Bounds;
		_tileModelCenter = tileBounds.Center;
		float tileExtent = Math.Max( 0.0001f, Math.Max( tileBounds.Size.x, tileBounds.Size.z ) );
		_tileModelScale = 1.0f / tileExtent;

		// Keep the custom object itself from being culled before its streamed chunks exist.
		// Individual chunks still use their own bounds when a culling camera is available.
		Bounds = BBox.FromPositionAndSize( Vector3.Zero, float.MaxValue );
	}

	private Model LoadTileModel()
	{
        return Model.Load( "models/tile/tile.vmdl" );

	}

	public void SetTilemap( TileMap tilemap )
	{
		if ( Tilemap == tilemap )
			return;

		Tilemap = tilemap;
		_tilemapRevision = -1;
	}

	public void UpdateStreaming( Vector3 cameraPos )
	{
		if ( Tilemap == null || !Tilemap.IsValid() )
		{
			DestroyMissingChunks();
			return;
		}

		TileSize = Math.Max( Tilemap.TileSize, 0.0001f );
		Axis = Tilemap.Axis;

		if ( _tilemapRevision != Tilemap.Revision || _tilemapAxis != Tilemap.Axis || _tilemapLayerCount != Tilemap.LayerCount || Math.Abs( _tilemapLayerSpacing - Tilemap.LayerSpacing ) > 0.0001f )
		{
			_tilemapRevision = Tilemap.Revision;
			_tilemapAxis = Tilemap.Axis;
			_tilemapLayerCount = Tilemap.LayerCount;
			_tilemapLayerSpacing = Tilemap.LayerSpacing;

			foreach ( var chunk in ActiveChunks.Values )
			{
				chunk.Revision = -1;
			}
		}

		HashSet<ChunkCoord> needed = new();
		ChunkCoord cameraChunk = WorldToChunk( cameraPos );

		for ( int x = -RenderRadius; x <= RenderRadius; x++ )
		{
			for ( int y = -RenderRadius; y <= RenderRadius; y++ )
			{
				ChunkCoord coord = new( cameraChunk.X + x, cameraChunk.Y + y );
				needed.Add( coord );

				if ( !ActiveChunks.TryGetValue( coord, out TilemapChunk chunk ) )
				{
					chunk = CreateChunk( coord );
					if ( chunk != null )
						ActiveChunks.Add( coord, chunk );
				}

				if ( chunk != null && chunk.Revision != Tilemap.Revision )
				{
					BuildChunk( chunk );
				}
			}
		}

		List<ChunkCoord> remove = new();
		foreach ( var pair in ActiveChunks )
		{
			if ( !needed.Contains( pair.Key ) )
				remove.Add( pair.Key );
		}

		foreach ( var coord in remove )
		{
			_pendingDestroy.Enqueue( ActiveChunks[coord] );
			ActiveChunks.Remove( coord );
		}

		ProcessPendingDestroy();
	}

	public void ProcessPendingDestroy()
	{
		while ( _pendingDestroy.Count > 0 )
		{
			var chunk = _pendingDestroy.Dequeue();

			foreach ( var batch in chunk.Batches )
			{
				batch.Dispose();
			}

			chunk.Batches.Clear();
		}
	}

	public void Disable()
	{
		foreach ( var pair in ActiveChunks )
		{
			_pendingDestroy.Enqueue( pair.Value );
		}

		ActiveChunks.Clear();
		ProcessPendingDestroy();
	}

	public override void RenderSceneObject()
	{
		base.RenderSceneObject();

		if ( Tilemap == null || !Tilemap.IsValid() )
			return;

		if ( CullingCamera == null )
			return;

		foreach ( var pair in ActiveChunks )
		{
			RenderChunk( pair.Value );
		}
	}

	private TilemapChunk CreateChunk( ChunkCoord coord )
	{
		TilemapChunk chunk = new()
		{
			Coord = coord,
			Revision = -1
		};

		chunk.Bounds = BuildChunkBounds( coord );
		return chunk;
	}

	private void BuildChunk( TilemapChunk chunk )
	{
		if ( Tilemap == null || !Tilemap.IsValid() )
			return;

		foreach ( var batch in chunk.Batches )
		{
			batch.Dispose();
		}

		chunk.Batches.Clear();

		int baseTileX = MathX.FloorToInt( (chunk.Coord.X * ChunkSize) / TileSize );
		int baseTileY = MathX.FloorToInt( (chunk.Coord.Y * ChunkSize) / TileSize );
		int maxTileX = baseTileX + ChunkResolution - 1;
		int maxTileY = baseTileY + ChunkResolution - 1;

		Dictionary<(int LayerIndex, ushort TilesetId), List<TileData>> tilesByLayerAndTileset = new();

		// Bottom layers are farther from the camera. Topmost layer 0 is closest.
		// Each layer has its own tile data; layers never affect each other's masks.
		for ( int layerIndex = Tilemap.LayerCount - 1; layerIndex >= 0; layerIndex-- )
		{
			var layer = Tilemap.GetLayer( layerIndex );
			if ( layer == null || !layer.IsVisible )
				continue;

			float layerOffset = Tilemap.GetLayerRenderNormalOffset( layerIndex );

			foreach ( var cell in Tilemap.GetFilledCells( layerIndex ) )
			{
				if ( cell.x < baseTileX || cell.x > maxTileX || cell.y < baseTileY || cell.y > maxTileY )
					continue;

				if ( !Tilemap.TryGetTile( layerIndex, cell.x, cell.y, out var tile ) )
					continue;

				if ( !Tilemap.IsRenderableTile( tile ) )
					continue;

				var tileset = Tilemap.GetTileset( tile.TilesetId );

				if ( tileset == null || !tileset.IsValid() )
					continue;

				var key = (layerIndex, tile.TilesetId);

				if ( !tilesByLayerAndTileset.TryGetValue( key, out var list ) )
				{
					list = new List<TileData>();
					tilesByLayerAndTileset[key] = list;
				}

				Vector3 worldPosition = Tilemap.CellToWorld( cell.x, cell.y, TileSize, layerOffset );

				list.Add( new TileData
				{
					Position = new Vector4(
						worldPosition.x,
						worldPosition.y,
						worldPosition.z,
						1.0f
					),

					UvRect = tileset.GetUvRect( tile.TileId )
				} );
			}
		}

		foreach ( var pair in tilesByLayerAndTileset )
		{
			int layerIndex = pair.Key.LayerIndex;
			ushort tilesetId = pair.Key.TilesetId;
			var data = pair.Value;

			if ( data.Count == 0 )
				continue;

			var tileset = Tilemap.GetTileset( tilesetId );

			if ( tileset == null || !tileset.IsValid() )
				continue;

			var batch = new TilemapChunkBatch
			{
				LayerIndex = layerIndex,
				TilesetId = tilesetId,
				Material = _tileMaterial.CreateCopy(),
				TileDataBuffer = new GpuBuffer<TileData>( data.Count ),
				ArgsBuffer = new GpuBuffer( 1, 5 * sizeof( uint ), GpuBuffer.UsageFlags.ByteAddress )
			};

			batch.TileDataBuffer.SetData( data.ToArray() );
			batch.ArgsBuffer.SetData( BuildArgs( _tileModel, (uint)data.Count ) );

			if ( !string.IsNullOrWhiteSpace( tileset.FilePath ) )
			{
				var texture = Texture.Load( tileset.FilePath );
				batch.Material.Attributes.Set( "_TileTexture", texture );
			}

			chunk.Batches.Add( batch );
		}

		// Draw farthest layers first, closest layers last. The normal offset/depth still
		// does the real work, but this keeps the order deterministic.
		chunk.Batches.Sort( ( a, b ) => b.LayerIndex.CompareTo( a.LayerIndex ) );

		chunk.Revision = Tilemap.Revision;
		chunk.Bounds = BuildChunkBounds( chunk.Coord );
	}

	private void RenderChunk( TilemapChunk chunk )
	{
		if ( chunk == null )
			return;

		if ( CullingCamera != null && !CullingCamera.GetFrustum().IsInside( chunk.Bounds, true ) )
			return;

		foreach ( var batch in chunk.Batches )
		{
			if ( batch.ArgsBuffer == null || batch.TileDataBuffer == null || batch.Material == null )
				continue;

			batch.Material.Attributes.Set( "_TileDataBuffer", batch.TileDataBuffer );
			batch.Material.Attributes.Set( "_TileSize", TileSize );
			batch.Material.Attributes.Set( "_TileModelCenter", _tileModelCenter );
			batch.Material.Attributes.Set( "_TileModelScale", _tileModelScale );
			batch.Material.Attributes.Set( "_TilePlaneAxisU", Tilemap.GetPlaneAxisU() );
			batch.Material.Attributes.Set( "_TilePlaneAxisV", Tilemap.GetPlaneAxisV() );
			batch.Material.Attributes.Set( "_TilePlaneNormal", Tilemap.GetPlaneNormal() );

			Graphics.DrawModelInstancedIndirect(
				_tileModel,
				batch.ArgsBuffer,
				0,
				batch.Material.Attributes
			);
		}
	}

	private BBox BuildChunkBounds( ChunkCoord coord )
	{
		float chunkWorldSize = ChunkWorldSize;
		float centerX = coord.X * chunkWorldSize + chunkWorldSize * 0.5f;
		float centerY = coord.Y * chunkWorldSize + chunkWorldSize * 0.5f;

		if ( Tilemap == null || !Tilemap.IsValid() )
			return BBox.FromPositionAndSize( Vector3.Zero, chunkWorldSize );

		return Tilemap.GetMapRectBounds( centerX, centerY, chunkWorldSize, chunkWorldSize, 10.0f + Tilemap.GetLayerRenderNormalThickness() );
	}

	private ChunkCoord WorldToChunk( Vector3 worldPos )
	{
		float chunkWorldSize = ChunkWorldSize;

		if ( Tilemap == null || !Tilemap.IsValid() )
			return new ChunkCoord( 0, 0 );

		Vector2 mapPosition = Tilemap.WorldToMap( worldPos );

		return new ChunkCoord(
			MathX.FloorToInt( mapPosition.x / chunkWorldSize ),
			MathX.FloorToInt( mapPosition.y / chunkWorldSize )
		);
	}

	private uint[] BuildArgs( Model model, uint instanceCount )
	{
		return new uint[]
		{
			(uint)model.GetIndexCount( 0 ),
			instanceCount,
			(uint)model.GetIndexStart( 0 ),
			(uint)model.GetBaseVertex( 0 ),
			0
		};
	}

	private void DestroyMissingChunks()
	{
		List<ChunkCoord> remove = new();
		foreach ( var pair in ActiveChunks )
			remove.Add( pair.Key );

		foreach ( var coord in remove )
		{
			_pendingDestroy.Enqueue( ActiveChunks[coord] );
			ActiveChunks.Remove( coord );
		}

		ProcessPendingDestroy();
	}
}