using System;
using Sandbox.Rendering;
using SquishyTerry.Sim;

namespace SquishyTerry;

/// <summary>
/// Runs the <see cref="SkinLayer"/> wrinkle cloth on the GPU, like the original game did. The CPU
/// SkinLayer still builds the topology and rest lengths; every frame this uploads the simulation
/// mesh and records the same passes (attach, constrain, 15 x neighbour/opposite relax, normals) from
/// <c>shaders/squishy_skin_cs.shader</c> into a command list that the camera runs while rendering.
///
/// The skin is drawn from separate "render" buffers that the command list refreshes after the opaque
/// pass, so the depth pre-pass and the colour pass of a frame always see the same surface.
/// </summary>
public sealed class GpuSkin : IDisposable
{
	const int Iterations = 15;

	readonly SkinLayer _layer;
	readonly ComputeShader _shader = new( "shaders/squishy_skin_cs.shader" );
	readonly CommandList _commands = new( "Squishy skin cloth" );
	readonly int _count;

	readonly GpuBuffer<Vector4> _basePositions, _baseNormals, _bary, _attach, _cloth, _scratch, _normals;
	readonly GpuBuffer<Vector4> _renderPositions, _renderNormals;
	readonly GpuBuffer<int> _baseIndices, _skinTriangle, _fixed, _neighbours, _opposites;
	readonly GpuBuffer<float> _neighbourRest, _oppositeRest;

	readonly Vector4[] _baseUpload;
	readonly int[] _fixedUpload;

	CameraComponent _camera;
	bool _initialized;
	int _pinnedTriangle = -1;

	/// <summary>Cloth positions to draw (simulation space), one float4 per skin vertex.</summary>
	public GpuBuffer<Vector4> Positions => _renderPositions;

	/// <summary>Cloth normals to draw (simulation space).</summary>
	public GpuBuffer<Vector4> Normals => _renderNormals;

	public GpuSkin( Face face, SkinLayer layer )
	{
		_layer = layer;
		_count = layer.VertexCount;
		layer.InitializeRest( face );

		int n = _count;
		int baseCount = face.VertexCount;

		_basePositions = new GpuBuffer<Vector4>( baseCount );
		_baseNormals = new GpuBuffer<Vector4>( baseCount );
		_baseIndices = new GpuBuffer<int>( face.Indices.Length );
		_baseIndices.SetData( face.Indices, 0 );
		_baseUpload = new Vector4[baseCount];

		_skinTriangle = new GpuBuffer<int>( n );
		_skinTriangle.SetData( layer.BaseTriangle, 0 );

		var bary = new Vector4[n];
		for ( int i = 0; i < n; i++ )
			bary[i] = new Vector4( layer.BaryW[i], layer.BaryU[i], layer.BaryV[i], layer.Strength[i] );
		_bary = new GpuBuffer<Vector4>( n );
		_bary.SetData( bary, 0 );

		_fixedUpload = new int[n];
		_fixed = new GpuBuffer<int>( n );
		_fixed.SetData( _fixedUpload, 0 );

		_neighbours = new GpuBuffer<int>( n * 6 );
		_neighbours.SetData( layer.Neighbours, 0 );
		_opposites = new GpuBuffer<int>( n * 6 );
		_opposites.SetData( layer.Opposites, 0 );
		_neighbourRest = new GpuBuffer<float>( n * 6 );
		_neighbourRest.SetData( layer.NeighbourRest, 0 );
		_oppositeRest = new GpuBuffer<float>( n * 6 );
		_oppositeRest.SetData( layer.OppositeRest, 0 );

		_attach = new GpuBuffer<Vector4>( n );
		_cloth = new GpuBuffer<Vector4>( n );
		_scratch = new GpuBuffer<Vector4>( n );
		_normals = new GpuBuffer<Vector4>( n );

		// Seed the render buffers with the CPU rest state so the first frames draw a proper surface.
		var seed = new Vector4[n];
		for ( int i = 0; i < n; i++ ) seed[i] = new Vector4( layer.Pos[i].X, layer.Pos[i].Y, layer.Pos[i].Z, 0 );
		_renderPositions = new GpuBuffer<Vector4>( n );
		_renderPositions.SetData( seed, 0 );
		for ( int i = 0; i < n; i++ ) seed[i] = new Vector4( layer.Normals[i].X, layer.Normals[i].Y, layer.Normals[i].Z, 0 );
		_renderNormals = new GpuBuffer<Vector4>( n );
		_renderNormals.SetData( seed, 0 );
	}

	/// <summary>Hooks the cloth passes into a camera's rendering (after the opaque pass).</summary>
	public void AttachTo( CameraComponent camera )
	{
		if ( camera == _camera ) return;
		_camera?.RemoveCommandList( _commands );
		_camera = camera;
		_camera?.AddCommandList( _commands, Stage.AfterOpaque, 0 );
	}

	/// <summary>Uploads this frame's simulation mesh and records the cloth passes (g.prototype.update in the original).</summary>
	public void Update( Face face, bool dragging, int dragTriangle )
	{
		var pos = face.Pos;
		for ( int i = 0; i < pos.Length; i++ ) _baseUpload[i] = new Vector4( pos[i].X, pos[i].Y, pos[i].Z, 0 );
		_basePositions.SetData( _baseUpload, 0 );
		var nrm = face.Normals;
		for ( int i = 0; i < nrm.Length; i++ ) _baseUpload[i] = new Vector4( nrm[i].X, nrm[i].Y, nrm[i].Z, 0 );
		_baseNormals.SetData( _baseUpload, 0 );

		// Like the original, the pinned set is only rewritten while dragging.
		if ( dragging && dragTriangle != _pinnedTriangle )
		{
			Array.Clear( _fixedUpload );
			foreach ( int v in _layer.VerticesOfBaseTriangle( dragTriangle ) ) _fixedUpload[v] = 1;
			_fixed.SetData( _fixedUpload, 0 );
			_pinnedTriangle = dragTriangle;
		}

		Record();
	}

	void Record()
	{
		var c = _commands;
		c.Reset();

		// Shared inputs.
		c.Set( "Count", _count );
		c.Set( "BasePositions", _basePositions );
		c.Set( "BaseNormals", _baseNormals );
		c.Set( "BaseIndices", _baseIndices );
		c.Set( "SkinTriangle", _skinTriangle );
		c.Set( "SkinBary", _bary );
		c.Set( "Fixed", _fixed );
		c.Set( "Attach", _attach );

		Dispatch( 0 );
		c.UavBarrier( _attach );

		if ( !_initialized )
		{
			c.Set( "Output", _cloth );
			Dispatch( 1 );
			c.UavBarrier( _cloth );
			_initialized = true;
		}

		c.Set( "Output", _cloth );
		Dispatch( 2 );
		c.UavBarrier( _cloth );

		// Ping-pong: neighbours cloth -> scratch, opposites scratch -> cloth; each iteration ends in _cloth.
		for ( int it = 0; it < Iterations; it++ )
		{
			c.Set( "RespectFixed", it < Iterations - 1 ? 1 : 0 );

			c.Set( "Links", _neighbours );
			c.Set( "Rest", _neighbourRest );
			c.Set( "Input", _cloth );
			c.Set( "Output", _scratch );
			Dispatch( 3 );
			c.UavBarrier( _scratch );

			c.Set( "Links", _opposites );
			c.Set( "Rest", _oppositeRest );
			c.Set( "Input", _scratch );
			c.Set( "Output", _cloth );
			Dispatch( 3 );
			c.UavBarrier( _cloth );
		}

		c.Set( "Links", _neighbours );
		c.Set( "Input", _cloth );
		c.Set( "Output", _normals );
		Dispatch( 4 );
		c.UavBarrier( _normals );

		// Publish for the next frame's draw.
		c.Set( "Input", _cloth );
		c.Set( "Output", _renderPositions );
		Dispatch( 5 );
		c.Set( "Input", _normals );
		c.Set( "Output", _renderNormals );
		Dispatch( 5 );
		c.UavBarrier( _renderPositions );
		c.UavBarrier( _renderNormals );
	}

	void Dispatch( int pass )
	{
		_commands.SetCombo( "D_PASS", pass );
		_commands.DispatchCompute( _shader, _count, 1, 1 );
	}

	public void Dispose()
	{
		_camera?.RemoveCommandList( _commands );
		_camera = null;

		_basePositions.Dispose();
		_baseNormals.Dispose();
		_bary.Dispose();
		_attach.Dispose();
		_cloth.Dispose();
		_scratch.Dispose();
		_normals.Dispose();
		_renderPositions.Dispose();
		_renderNormals.Dispose();
		_baseIndices.Dispose();
		_skinTriangle.Dispose();
		_fixed.Dispose();
		_neighbours.Dispose();
		_opposites.Dispose();
		_neighbourRest.Dispose();
		_oppositeRest.Dispose();
	}
}