using System;
using System.Collections.Generic;
using System.Linq;
using System.Text.Json.Nodes;
using System.Text.Json.Serialization;
using Sandbox;

namespace Sandbox;

[Title( "Shader Particle Model Renderer" )]
[Category( "Particles" )]
[Description(
	"Adds the \"Particle Shader\" feature which let's you set shader parameters on the particles model using the particle system !" )]
public class ShaderParticleModelRenderer : ParticleController, Component.ExecuteInEditor
{
	private ParticleModelRenderer _particleModelRenderer { get; set; } = new ParticleModelRenderer();

	[Property]
	public bool Batchable { get; set; } = true;

	/**
	 *  These are the dictionaries used to set the values for the render attributes in the particle system
	 */
	[Property, FeatureEnabled("ParticleShader")]
	public bool ParticleShaderEnabled { get; set; } = true;
	
	
	[Property, Group( "ColorParameters" ), Feature( "ParticleShader" )]
	public Dictionary<String, ParticleGradient> Colors { get; set; }

	[Property, Group( "FloatParameters" ), Feature( "ParticleShader" )]
	public Dictionary<String, ParticleFloat> Floats { get; set; }

	[Property, Group( "Float2Parameters" ), Feature( "ParticleShader" )]
	public Dictionary<String, ParticleVector2> Floats2 { get; set; }

	[Property, Group( "Float3Parameters" ), Feature( "ParticleShader" )]
	public Dictionary<String, ParticleVector3> Floats3 { get; set; }

	[Property, Group( "Float4Parameters" ), Feature( "ParticleShader" )]
	public Dictionary<String, Vector4> Floats4 { get; set; }

	[Property, Group( "Textures" ), Feature( "ParticleShader" )]
	public Dictionary<String, Texture> Textures { get; set; }

	[Property, Group( "DynamicCombos" ), Feature( "ParticleShader" )]
	public Dictionary<String, int> DynamicCombos { get; set; }

	
	
	/// <summary>
	/// Particle rotation is world space: the emitter rotation does not orient the particle at all.
	/// When clear, the particle rotation follows the emitter rotation every frame (local space).
	/// </summary>
	[Property]
	public bool IsWorldRotation { get; set; } = false;

	/// <summary>
	/// World rotation, but seeded with the emitter rotation sampled at spawn: the particle is born
	/// facing the emitter's direction and is never re-oriented afterwards. Keeps rate-over-distance
	/// trails spawning in the right direction while staying in world space.
	/// </summary>
	[Property, ShowIf( nameof(IsWorldRotation), true )]
	public bool InheritSpawnRotation { get; set; } = false;

	[Property, FeatureEnabled( "SheetSequence" )]
	public bool SheetSequence { get; set; } = false;

	[Property, Feature( "SheetSequence" )]
	public float SheetSequenceSpeed { get; set; } = 1f;
	
	

	[RequireComponent] public new ParticleEffect ParticleEffect { get; set; }

	[Property, Order( -100 ), InlineEditor( Label = false ), Group( "Advanced Rendering", StartFolded = true )]
	public RenderOptions RenderOptions => _particleModelRenderer.RenderOptions;
	
	
	protected override void OnStart()
	{
		Floats = Floats == null ? new Dictionary<string, ParticleFloat>() : Floats;
		Floats2 = Floats2 == null ? new Dictionary<string, ParticleVector2>() : Floats2;
		Floats3 = Floats3 == null ? new Dictionary<string, ParticleVector3>() : Floats3;
		Floats4 = Floats4 == null ? new Dictionary<string, Vector4>() : Floats4;
		Textures = Textures == null ? new Dictionary<string, Texture>() : Textures;
		Colors = Colors == null ? new Dictionary<String, ParticleGradient>() : Colors;
		DynamicCombos = DynamicCombos == null ? new Dictionary<string, int>() : DynamicCombos;
		
		if( ParticleShaderEnabled ) ReadAttributes();
	}
	
	[Button, Feature("ParticleShader")]
	private void ReadAttributes()
	{
		if ( MaterialOverride == null || !FileSystem.Mounted.FileExists( MaterialOverride.Shader.ResourcePath )) return;
		var parsedFloats2 = new Dictionary<string, Vector2>();
		var parsedFloats3 = new Dictionary<string, Vector3>();
		AttributesParser<ParticleFloat, ParticleGradient> parser = new AttributesParser<ParticleFloat, ParticleGradient>(new ParticleAttributeTypeSet());
		parser.Floats = Floats;
		parser.Floats2 = parsedFloats2;
		parser.Floats3 = parsedFloats3;
		parser.Floats4 = Floats4;
		parser.Textures = Textures;
		parser.DynamicCombos = DynamicCombos;
		parser.Colors = Colors;
		parser.ParseAttributes( MaterialOverride.Shader.ResourcePath );

		foreach ( var pair in parsedFloats2 )
		{
			if ( Floats2.ContainsKey( pair.Key ) ) continue;
			Floats2.Add( pair.Key, pair.Value ); // implicit Vector2 -> ParticleVector2
		}

		foreach ( var pair in parsedFloats3 )
		{
			if ( Floats3.ContainsKey( pair.Key ) ) continue;
			Floats3.Add( pair.Key, pair.Value ); // implicit Vector3 -> ParticleVector3
		}
	}

	public sealed class ModelEntry
	{
		private Model _model;

		[KeyProperty]
		public Model Model
		{
			get => _model;
			set
			{
				if ( _model == value )
					return;

				_model = value;

				MaterialGroup = default;
				BodyGroups = _model?.DefaultBodyGroupMask ?? default;
			}
		}

		[Model.MaterialGroup, ShowIf( nameof(HasMaterialGroups), true )]
		public string MaterialGroup { get; set; }

		[Model.BodyGroupMask, ShowIf( nameof(HasBodyGroups), true )]
		public ulong BodyGroups { get; set; }

		[Hide, JsonIgnore] public bool HasMaterialGroups => Model?.MaterialGroupCount > 0;

		[Hide, JsonIgnore] public bool HasBodyGroups => Model?.BodyParts.Sum( x => x.Choices.Count ) > 1;

		public static implicit operator ModelEntry( Model model ) => new() { Model = model };
	}

	[Hide, Obsolete( "Use Choices" )] public List<Model> Models { get; set; } = new();

	[Property] public List<ModelEntry> Choices { get; set; } = new List<ModelEntry> { Model.Cube };

	[Property] public Material MaterialOverride { get; set; }

	[Property, Feature( "ScaleXYZ" )] public ParticleFloat ScaleX { get; set; } = 1;
	[Property, Feature( "ScaleXYZ" )] public ParticleFloat ScaleY { get; set; } = 1;
	[Property, Feature( "ScaleXYZ" )] public ParticleFloat ScaleZ { get; set; } = 1;

	[Property, FeatureEnabled( "ScaleXYZ" )]
	public bool ApplyScaleXYZ { get; set; } = true;

	[Property] public ParticleFloat Scale { get; set; } = 1;
	
	[Property] public bool CastShadows { get; set; } = true;

	[Property] public Allignement Allignement { get; set; }

	protected override void OnParticleCreated( Particle p )
	{
		var particleModel = new CustomParticleModel( this );
		p.AddListener( particleModel, this );
		
	}

	// Scene objects of dead particles, kept hidden and handed to new particles of the same model instead of
	// creating and deleting one per particle. Particles spawn and die on the parallel particle step, so it's locked.
	private readonly Dictionary<Model, Stack<SceneObject>> _pool = new();
	private bool _poolOpen;

	internal SceneObject RentSceneObject( Model model )
	{
		lock ( _pool )
		{
			if ( _poolOpen && _pool.TryGetValue( model, out var stack ) )
			{
				while ( stack.TryPop( out var so ) )
				{
					if ( !so.IsValid() ) continue;

					so.RenderingEnabled = true;
					return so;
				}
			}
		}

		return new SceneObject( Scene.SceneWorld, model );
	}

	internal void ReturnSceneObject( SceneObject so, Model model )
	{
		if ( !so.IsValid() ) return;

		lock ( _pool )
		{
			if ( _poolOpen )
			{
				so.RenderingEnabled = false;

				if ( !_pool.TryGetValue( model, out var stack ) )
					_pool[model] = stack = new Stack<SceneObject>();

				stack.Push( so );
				return;
			}
		}

		so.Delete();
	}

	protected override void OnEnabled()
	{
		lock ( _pool ) _poolOpen = true;

		base.OnEnabled();
	}

	protected override void OnDisabled()
	{
		// Live particles hand their scene objects back while the base disables, so close the pool after
		base.OnDisabled();

		lock ( _pool )
		{
			_poolOpen = false;

			foreach ( var stack in _pool.Values )
			{
				foreach ( var so in stack )
				{
					if ( so.IsValid() ) so.Delete();
				}
			}

			_pool.Clear();
		}
	}

	// Read by every particle on the parallel step; sampled once per step here instead of once per particle
	internal Vector3 StepWorldScale { get; private set; } = Vector3.One;
	internal Rotation StepWorldRotation { get; private set; } = Rotation.Identity;
	internal Vector3? StepCameraPosition { get; private set; }

	protected override void OnBeforeStep( float delta )
	{
		base.OnBeforeStep( delta );

		StepWorldScale = WorldScale;
		StepWorldRotation = WorldRotation;
		var camera = Scene.Camera;
		StepCameraPosition = camera.IsValid() ? camera.WorldPosition : null;

		if ( ParticleShaderEnabled )
			AnimatedAttributes.Rebuild( Floats, Floats2, Floats3, Colors );
	}

	/// <summary>
	/// Shader attributes that change over a particle's life, shared by this renderer's particles.
	/// </summary>
	internal AnimatedAttributes AnimatedAttributes { get; } = new();

	public override int ComponentVersion => 2;

	[JsonUpgrader( typeof(ParticleModelRenderer), 1 )]
	static void Upgrader_v1( JsonObject obj )
	{
		if ( obj.TryGetPropertyValue( "Models", out var node ) )
		{
			var choices = new JsonArray();

			foreach ( var model in node.AsArray() )
			{
				if ( model is null )
					continue;

				choices.Add( new JsonObject { ["Model"] = model.ToString() } );
			}

			obj["Choices"] = choices;
			obj.Remove( "Models" );
		}
	}

	// Floats2 used to be a Dictionary<string, Vector2> (each value serialized as an
	// "x,y" string). It is now a Dictionary<string, ParticleVector2>; convert each
	// entry into a ParticleVector2 with constant X/Y components.
	[JsonUpgrader( typeof(ShaderParticleModelRenderer), 2 )]
	static void Upgrader_v2( JsonObject obj )
	{
		if ( !obj.TryGetPropertyValue( "Floats2", out var node ) || node is not JsonObject floats2 )
			return;

		foreach ( var kv in floats2.ToList() )
		{
			if ( kv.Value is not JsonValue val || !val.TryGetValue<string>( out var str ) )
				continue;

			var parts = str.Split( ',' );
			if ( parts.Length < 2 )
				continue;

			if ( !float.TryParse( parts[0].Trim(), System.Globalization.NumberStyles.Float, System.Globalization.CultureInfo.InvariantCulture, out var x ) ||
			     !float.TryParse( parts[1].Trim(), System.Globalization.NumberStyles.Float, System.Globalization.CultureInfo.InvariantCulture, out var y ) )
				continue;

			floats2[kv.Key] = new JsonObject { ["X"] = x, ["Y"] = y };
		}
	}
	
	
}

public enum Allignement
{
	SimulationSpace,
	FaceCamera,
	FaceVelocity,
	
}





public class CustomParticleModel : Particle.BaseListener
{
	public ShaderParticleModelRenderer Renderer;

	public SceneObject so;

	/// <summary>
	/// Emitter rotation sampled when the particle spawned. Used as the fixed base rotation when
	/// <see cref="ShaderParticleModelRenderer.InheritSpawnRotation"/> is set, so a particle inherits
	/// the emitter's orientation at birth but is not re-oriented afterwards.
	/// </summary>
	private Rotation _spawnRotation = Rotation.Identity;

	private ParticleAttributesSetter _particleAttributesSetter;

	// Last values pushed to the scene object; each set is a native call plus a change notification
	// Nullable so the first frame always sets them: a reused scene object still holds the last particle's values
	private Color? _tint;
	private bool? _castShadows;
	private int _renderOptionsMask = -1;
	private Material _material;

	// Model the scene object was rented for, to hand it back to the right pool
	private Model _model;

	public CustomParticleModel( ShaderParticleModelRenderer renderer )
	{
		Renderer = renderer;
	}

	public override void OnEnabled( Particle p )
	{
		_spawnRotation = Renderer.StepWorldRotation;

		var entry = Random.Shared.FromList( Renderer.Choices );
		var model = entry?.Model;
		_model = model ?? Model.Cube;
		so = Renderer.RentSceneObject( _model );
		so.Batchable = false;
		if ( model is not null )
		{
			so.MeshGroupMask = entry.BodyGroups;
			so.SetMaterialGroup( entry.MaterialGroup );
		}

		if ( Renderer.SheetSequence )
		{
			so.Attributes.Set( "AnimSpeed", Renderer.SheetSequenceSpeed / p.LifeTimeRemaining );
		}

		if ( !Renderer.ParticleShaderEnabled ) return;
		_particleAttributesSetter = new ParticleAttributesSetter( so.Attributes, p );
		SetRenderAttributes();
	}

	public override void OnDisabled( Particle p )
	{
		if ( !so.IsValid() ) return;
		Renderer.ReturnSceneObject( so, _model );
		so = null;
	}

	public override void OnUpdate( Particle p, float dt )
	{
		if ( !so.IsValid() ) return;


		var angles = ComputeRotation( p );
		
		var scale = p.Size * Renderer.StepWorldScale;
		if ( Renderer.ApplyScaleXYZ )
		{
			scale *= EvaluateScale( p );
		}

		so.Transform = new Transform( p.Position, angles, scale * Renderer.Scale.Evaluate( p, 6211 ) );
		var tint = p.Color.WithAlphaMultiplied( p.Alpha );
		if ( tint != _tint )
		{
			so.ColorTint = tint;
			_tint = tint;
		}

		if ( _castShadows != Renderer.CastShadows )
		{
			so.Flags.CastShadows = Renderer.CastShadows;
			_castShadows = Renderer.CastShadows;
		}

		if ( _material != Renderer.MaterialOverride )
		{
			so.SetMaterialOverride( Renderer.MaterialOverride );
			_material = Renderer.MaterialOverride;
		}

		// Constant attributes were set at spawn and stay on the scene object. While editing the scene
		// everything is re-set every frame, so inspector changes show on live particles straight away.
		if ( Renderer.ParticleShaderEnabled )
		{
			if ( Renderer.Scene.IsEditor ) _particleAttributesSetter.SetAttributes();
			else _particleAttributesSetter.SetAnimatedAttributes();
		}
		

		// Apply sets four native flags; only redo it when the options change
		var renderOptions = Renderer.RenderOptions;
		if ( renderOptions != null )
		{
			var mask = (renderOptions.Game ? 1 : 0) | (renderOptions.Overlay ? 2 : 0) | (renderOptions.Bloom ? 4 : 0) | (renderOptions.AfterUI ? 8 : 0);
			if ( mask != _renderOptionsMask )
			{
				renderOptions.Apply( so );
				_renderOptionsMask = mask;
			}
		}
	}

	


	private Vector3 EvaluateScale( Particle p )
	{
		
		var scaleX = Renderer.ScaleX.Evaluate( p, 6211 );
		var scaleY = Renderer.ScaleY.Evaluate( p, 6211 );
		var scaleZ = Renderer.ScaleZ.Evaluate( p, 6211 );
		return new Vector3( scaleX, scaleY, scaleZ );
	}


	private Angles ComputeRotation(Particle p)
	{
		var angles = new Rotation();
		switch ( Renderer.Allignement )
		{
			case Allignement.FaceCamera :
				if ( Renderer.StepCameraPosition is not { } cameraPosition ) break;
				var dir = cameraPosition - p.Position;
				angles = Rotation.LookAt( dir, Vector3.Up ) * p.Angles.ToRotation();
				break;
			case Allignement.FaceVelocity :
				angles = Rotation.LookAt( p.Velocity.Normal, Vector3.Up ) * p.Angles.ToRotation();
				break;
			case Allignement.SimulationSpace :
				angles = BaseRotation().Angles();
				angles *= p.Angles;
				break;
		}
		return angles;
	}

	/// <summary>
	/// Base rotation the particle's own <see cref="Particle.Angles"/> are applied on top of, for
	/// <see cref="Allignement.SimulationSpace"/>.
	/// </summary>
	private Rotation BaseRotation()
	{
		if ( !Renderer.IsWorldRotation )
			return Renderer.StepWorldRotation; // follows the emitter every frame

		return Renderer.InheritSpawnRotation ? _spawnRotation : Rotation.Identity;
	}

	private void SetRenderAttributes()
	{
		_particleAttributesSetter.Floats = Renderer.Floats;
		_particleAttributesSetter.Floats2 = Renderer.Floats2;
		_particleAttributesSetter.Floats3 = Renderer.Floats3;
		_particleAttributesSetter.Floats4 = Renderer.Floats4;
		_particleAttributesSetter.Textures = Renderer.Textures;
		_particleAttributesSetter.Colors = Renderer.Colors;
		_particleAttributesSetter.DynamicCombos = Renderer.DynamicCombos;
		_particleAttributesSetter.Animated = Renderer.AnimatedAttributes;
		_particleAttributesSetter.SetAttributes();
	}
}