using System;
using System.Collections.Generic;
using System.Runtime.InteropServices;
using Sandbox;
using Sandbox.Rendering;
using RenderStage = Sandbox.Rendering.Stage;

namespace RedSnail.GrassTool;

/// <summary>
/// Renders a painted grass field. Each frame a compute pass expands the density map into blade
/// instances for whatever is near the camera, then a single indirect draw renders them all.
/// No blade ever exists on the CPU, and nothing is stored per-blade in the scene.
/// </summary>
[Icon( "grass" ), Group( "Grass" ), Title( "Grass Renderer" )]
public sealed class GrassRenderer : Component, Component.ExecuteInEditor, Component.DontExecuteOnServer
{
	[StructLayout( LayoutKind.Sequential )]
	private struct GpuChunk
	{
		public Vector2 Origin;
		public int CellOffset;
		public int Pad;
	}

	[StructLayout( LayoutKind.Sequential )]
	private struct GpuBlade
	{
		public Vector3 Position;
		public float Yaw;
		public Vector3 Normal;
		public float Height;
		public float Width;
		public float Tint;
		public float Phase;
		public float Pad; // keeps the struct at 48 bytes
	}

	// Must match GRASS_BLADE_* in grass_shared.fxc.
	private const int BladeStripVerts = 7;
	private const int BladeVertexCount = (BladeStripVerts - 2) * 3;

	// Byte offset of InstanceCount within IndirectDrawArguments. The struct is sequential
	// { uint VertexCount; uint InstanceCount; uint FirstVertex; uint FirstInstance; }, so this is
	// fixed at 4. Marshal.OffsetOf would express it directly but is outside the sandbox whitelist.
	private const int ArgsInstanceCountOffset = 4;

	[Property, Group( "General" )]
	public GrassDefinition Definition { get; set; }

	/// <summary>
	/// Ceiling on blades alive at once. Each costs 48 bytes of GPU memory, so 500k is ~24 MB.
	/// Raise it if dense fields visibly clip out at the far edge of the render distance.
	/// </summary>
	[Property, Group( "General" ), Range( 50000, 4000000 )]
	public int MaxBlades { get; set; } = 500000;

	/// <summary>Painted coverage. Serialized as a binary blob, not JSON.</summary>
	[Property, Hide]
	public GrassStorage Storage { get; set; } = new();

	private ComputeShader _generateShader;
	private Material _material;
	private CommandList _commandList;

	private GpuBuffer<GpuChunk> _chunkBuffer;
	private GpuBuffer<GrassStorage.Cell> _cellBuffer;
	private GpuBuffer<int> _visibleChunkBuffer;
	private GpuBuffer<Vector4> _frustumPlaneBuffer;
	private GpuBuffer<GpuBlade> _bladeBuffer;
	private GpuBuffer<GpuBuffer.IndirectDrawArguments> _argsBuffer;

	private readonly List<Vector2> _chunkOrigins = [];
	private int[] _visibleScratch = [];
	private readonly Vector4[] _planeScratch = new Vector4[6];

	private CameraComponent _lastCamera;
	private int _chunkCount;
	private int _uploadedVersion = -1;
	private int _uploadedMaxBlades = -1;

	protected override void OnEnabled()
	{
		Storage ??= new GrassStorage();

		_generateShader = new ComputeShader("grass_generate_cs");
		_material = Material.FromShader("grass");
		_commandList = new CommandList("Grass Rendering");

		// Force a full re-upload: the buffers were released on disable, so a matching version
		// number here would otherwise leave the compute pass reading freed resources.
		_uploadedVersion = -1;
		_uploadedMaxBlades = -1;

		RefreshBuffers();
	}

	protected override void OnDisabled()
	{
		_lastCamera?.RemoveCommandList( _commandList );
		_lastCamera = null;

		_commandList?.Reset();
		_commandList = null;

		ReleaseBuffers();

		_generateShader = null;
		_material = null;

		_uploadedVersion = -1;
		_uploadedMaxBlades = -1;
	}

	protected override void OnUpdate()
	{
		// Only a scene camera's command list actually replays, in the editor viewport as well as
		// in game. The editor camera never replays one, so it can't own the list.
		var renderCamera = Scene.Camera;

		if ( renderCamera != _lastCamera )
		{
			_lastCamera?.RemoveCommandList( _commandList );
			renderCamera?.AddCommandList( _commandList, RenderStage.AfterOpaque );
			_lastCamera = renderCamera;
		}

		if ( !renderCamera.IsValid() )
			return;

		// ...but culling has to follow whichever camera the viewport is actually looking through,
		// or nothing generates outside the game camera's frustum while you're editing.
		var cullCamera = GetCullCamera();
		if ( !cullCamera.IsValid() )
			return;

		RefreshBuffers();
		RecordCommandList( cullCamera );
	}

	private CameraComponent GetCullCamera()
	{
		if ( Game.IsPlaying )
			return Scene.Camera;

		var editorCamera = Application.Editor?.Camera;
		return editorCamera.IsValid() ? editorCamera : Scene.Camera;
	}

	/// <summary>
	/// Re-packs painted chunks into GPU buffers. Skipped entirely unless the painted data or the
	/// blade budget actually changed.
	/// </summary>
	private void RefreshBuffers()
	{
		if ( Storage is null || Storage.ChunkCount == 0 )
		{
			if ( _chunkCount != 0 )
			{
				ReleaseBuffers();
				_chunkCount = 0;
			}

			_uploadedVersion = Storage?.Revision ?? -1;
			return;
		}

		if ( _uploadedVersion == Storage.Revision && _uploadedMaxBlades == MaxBlades && _chunkBuffer is not null )
			return;

		ReleaseBuffers();

		_chunkCount = Storage.ChunkCount;
		_uploadedVersion = Storage.Revision;
		_uploadedMaxBlades = MaxBlades;

		var chunks = new GpuChunk[_chunkCount];
		var cells = new GrassStorage.Cell[_chunkCount * GrassStorage.CellsPerChunk];

		_chunkOrigins.Clear();

		var index = 0;
		foreach ( var (coord, cellData) in Storage.Chunks )
		{
			var origin = GrassStorage.ChunkOrigin( coord );
			var offset = index * GrassStorage.CellsPerChunk;

			chunks[index] = new GpuChunk { Origin = origin, CellOffset = offset };
			Array.Copy( cellData, 0, cells, offset, GrassStorage.CellsPerChunk );
			_chunkOrigins.Add( origin );

			index++;
		}

		_chunkBuffer = new GpuBuffer<GpuChunk>( _chunkCount, GpuBuffer.UsageFlags.Structured );
		_chunkBuffer.SetData( chunks );

		_cellBuffer = new GpuBuffer<GrassStorage.Cell>( cells.Length, GpuBuffer.UsageFlags.Structured );
		_cellBuffer.SetData( cells );

		_visibleChunkBuffer = new GpuBuffer<int>( _chunkCount, GpuBuffer.UsageFlags.Structured );
		_visibleScratch = new int[_chunkCount];

		_frustumPlaneBuffer = new GpuBuffer<Vector4>( 6, GpuBuffer.UsageFlags.Structured );

		_bladeBuffer = new GpuBuffer<GpuBlade>( MaxBlades, GpuBuffer.UsageFlags.Structured | GpuBuffer.UsageFlags.Append );

		_argsBuffer = new GpuBuffer<GpuBuffer.IndirectDrawArguments>( 1, GpuBuffer.UsageFlags.IndirectDrawArguments );
		_argsBuffer.SetData( new[]
		{
			new GpuBuffer.IndirectDrawArguments { VertexCount = BladeVertexCount }
		} );
	}

	private void RecordCommandList( CameraComponent camera )
	{
		_commandList.Reset();

		if ( _chunkCount == 0 || _bladeBuffer is null )
			return;

		var definition = Definition;
		var cameraPosition = camera.WorldPosition;

		var visibleCount = CollectVisibleChunks( cameraPosition, definition?.MaxDistance ?? 12000.0f );
		if ( visibleCount == 0 )
			return;

		UploadFrustumPlanes( camera );

		_commandList.Attributes.Set( "GrassChunks", _chunkBuffer );
		_commandList.Attributes.Set( "GrassCells", _cellBuffer );
		_commandList.Attributes.Set( "GrassVisibleChunks", _visibleChunkBuffer );
		_commandList.Attributes.Set( "GrassFrustumPlanes", _frustumPlaneBuffer );
		_commandList.Attributes.Set( "GrassBlades", _bladeBuffer );
		_commandList.Attributes.Set( "GrassVisibleChunkCount", visibleCount );
		_commandList.Attributes.Set( "GrassCameraPos", cameraPosition );
		_commandList.Attributes.Set( "GrassTime", RealTime.Now );

		definition?.ApplyTo( _commandList.Attributes );

		_commandList.ResourceBarrierTransition( _bladeBuffer, ResourceState.UnorderedAccess );
		_commandList.SetCounterValue( _bladeBuffer, 0 );

		_commandList.DispatchCompute( _generateShader, visibleCount * GrassStorage.CellsPerChunk, 1, 1 );

		// The appends must all land before the counter is read into the draw arguments.
		_commandList.UavBarrier( _bladeBuffer );

		_commandList.ResourceBarrierTransition( _argsBuffer, ResourceState.CopyDestination );
		_commandList.CopyStructureCount( _bladeBuffer, _argsBuffer, ArgsInstanceCountOffset );

		_commandList.ResourceBarrierTransition( _bladeBuffer, ResourceState.GenericRead );
		_commandList.ResourceBarrierTransition( _argsBuffer, ResourceState.IndirectArgument );

		_commandList.DrawInstancedIndirect( _material, _argsBuffer );
	}

	/// <summary>
	/// Narrows the dispatch to chunks near the camera. Per-blade frustum culling happens on the
	/// GPU; this only has to be conservative.
	/// </summary>
	private int CollectVisibleChunks( Vector3 cameraPosition, float maxDistance )
	{
		// A chunk's near corner can be in range while its centre isn't, hence the circumradius.
		var cullRadius = maxDistance + GrassStorage.ChunkSize * 0.7072f;
		var cullRadiusSq = cullRadius * cullRadius;

		var halfChunk = GrassStorage.ChunkSize * 0.5f;
		var count = 0;

		for ( var i = 0; i < _chunkOrigins.Count; i++ )
		{
			var origin = _chunkOrigins[i];
			var dx = origin.x + halfChunk - cameraPosition.x;
			var dy = origin.y + halfChunk - cameraPosition.y;

			if ( dx * dx + dy * dy > cullRadiusSq )
				continue;

			_visibleScratch[count++] = i;
		}

		if ( count > 0 )
			_visibleChunkBuffer.SetData( _visibleScratch.AsSpan( 0, count ) );

		return count;
	}

	private void UploadFrustumPlanes( CameraComponent camera )
	{
		var frustum = camera.GetFrustum();

		// Plane.GetDistance is dot( point, Normal ) - Distance, so w is negated to let the shader
		// use a plain dot( xyz, p ) + w.
		_planeScratch[0] = ToVector4( frustum.LeftPlane );
		_planeScratch[1] = ToVector4( frustum.RightPlane );
		_planeScratch[2] = ToVector4( frustum.TopPlane );
		_planeScratch[3] = ToVector4( frustum.BottomPlane );
		_planeScratch[4] = ToVector4( frustum.NearPlane );
		_planeScratch[5] = ToVector4( frustum.FarPlane );

		_frustumPlaneBuffer.SetData( _planeScratch );

		static Vector4 ToVector4( Plane plane ) =>
			new( plane.Normal.x, plane.Normal.y, plane.Normal.z, -plane.Distance );
	}

	private void ReleaseBuffers()
	{
		_chunkBuffer?.Dispose();
		_chunkBuffer = null;

		_cellBuffer?.Dispose();
		_cellBuffer = null;

		_visibleChunkBuffer?.Dispose();
		_visibleChunkBuffer = null;

		_frustumPlaneBuffer?.Dispose();
		_frustumPlaneBuffer = null;

		_bladeBuffer?.Dispose();
		_bladeBuffer = null;

		_argsBuffer?.Dispose();
		_argsBuffer = null;

		_chunkOrigins.Clear();
		_chunkCount = 0;
	}

	/// <summary>
	/// Called by the editor tool after painting, so the next frame re-uploads the density map.
	/// </summary>
	public void MarkDirty() => _uploadedVersion = -1;

	protected override void DrawGizmos()
	{
		if ( !Gizmo.IsSelected || Storage is null || Storage.ChunkCount == 0 )
			return;

		Gizmo.Draw.Color = Color.Green.WithAlpha( 0.35f );

		foreach ( var (coord, _) in Storage.Chunks )
		{
			var origin = GrassStorage.ChunkOrigin( coord );
			var mins = new Vector3( origin.x, origin.y, 0 );
			var maxs = new Vector3( origin.x + GrassStorage.ChunkSize, origin.y + GrassStorage.ChunkSize, 0 );

			Gizmo.Draw.LineBBox( new BBox( WorldTransform.PointToLocal( mins ), WorldTransform.PointToLocal( maxs ) ) );
		}
	}
}