using Editor.Prism.Compiler;
using Editor.Prism.Compiler.Ir;
using Editor.Prism.Core;
using Editor.Prism.Model;
using System.Text;

namespace Editor.Prism.Nodes;

/// <summary>Which shader stages a custom-code body is allowed to run in.</summary>
public enum PrismCodeStage
{
	/// <summary>Wherever the graph needs it.</summary>
	Any,
	/// <summary>Vertex program only. Values are interpolated to the pixel program automatically.</summary>
	VertexOnly,
	/// <summary>Pixel program only. Required for derivatives, implicit-LOD sampling and discard.</summary>
	PixelOnly
}

/// <summary>
/// A node whose body the user writes.
/// <para>
/// The body is lowered to a <see cref="HelperFunction"/> rather than pasted inline, which is what
/// makes it behave like everything else in the graph: it is deduplicated per module, calls to it are
/// hash-consed so two identical invocations collapse into one, its lines land in the source map so a
/// compiler error can point back at the exact line inside <em>this node's</em> editor, and it inherits
/// the same stage and shader-model gates as the built-in nodes.
/// </para>
/// <para>
/// A node with more than one output emits one helper per output, each returning that output. Every
/// helper carries the same body, so the shader compiler eliminates whatever the returned value does
/// not depend on; the alternative — a single helper with <c>out</c> parameters — cannot be expressed
/// as an expression and would defeat common-subexpression elimination entirely.
/// </para>
/// </summary>
[NodeInfo( Id = CustomCodeNode.TypeId, Title = "Custom Code", Category = "Utility",
	Icon = "code", Tier = NodeTier.Advanced,
	Keywords = new[] { "custom", "code", "hlsl", "slang", "function", "snippet", "expression" },
	Description = "Write HLSL — and optionally a separate Slang body — against a port table you " +
		"declare yourself." )]
[NodeVersion( 1 )]
public sealed class CustomCodeNode : PrismNode, IStageConstrained
{
	/// <summary>The stable type id.</summary>
	public const string TypeId = "prism.custom.code";

	/// <summary>The body a new node starts with.</summary>
	public const string DefaultBody = "// Assign each declared output. Inputs are in scope by name.\r\nOut = In;";

	/// <summary>What a node with no port table declares.</summary>
	static readonly SubgraphPortInfo[] s_defaultInputs =
	[
		new() { Name = "In", Type = "float" }
	];

	/// <summary>What a node with no port table returns.</summary>
	static readonly SubgraphPortInfo[] s_defaultOutputs =
	[
		new() { Name = "Out", Type = "float" }
	];

	/// <summary>The function name the generated helper is built around. Sanitized before use.</summary>
	public string FunctionName { get; set; } = "CustomCode";

	/// <summary>The HLSL body. Statements only — the signature is generated.</summary>
	public string Hlsl { get; set; } = DefaultBody;

	/// <summary>
	/// A separate Slang body, used only when <see cref="UseSlangBody"/> is set. Leave it off unless
	/// the two languages genuinely need to differ.
	/// </summary>
	public string Slang { get; set; }

	/// <summary>True when <see cref="Slang"/> replaces <see cref="Hlsl"/> on the Slang backend.</summary>
	public bool UseSlangBody { get; set; }

	/// <summary>Headers the body needs, added to the module's include list.</summary>
	public List<string> Includes { get; set; }

	/// <summary>Which stages the body is allowed to run in.</summary>
	public PrismCodeStage Stage { get; set; } = PrismCodeStage.Any;

	/// <summary>The node's declared inputs. Each row becomes a port and a function parameter.</summary>
	public List<SubgraphPortInfo> InputPorts { get; set; }

	/// <summary>The node's declared outputs. Each row becomes a port and a generated helper.</summary>
	public List<SubgraphPortInfo> OutputPorts { get; set; }

	// ---- port table -------------------------------------------------------

	/// <summary>The input rows actually in effect, falling back to a single float in.</summary>
	public IReadOnlyList<SubgraphPortInfo> EffectiveInputs =>
		InputPorts is { Count: > 0 } ? InputPorts : s_defaultInputs;

	/// <summary>The output rows actually in effect, falling back to a single float out.</summary>
	public IReadOnlyList<SubgraphPortInfo> EffectiveOutputs =>
		OutputPorts is { Count: > 0 } ? OutputPorts : s_defaultOutputs;

	/// <summary>Add an input row and rebuild the port set.</summary>
	public SubgraphPortInfo AddInput( string name = null, string type = "float" )
	{
		// A node still on its defaults has no table yet. Materialise the defaults before appending, or
		// adding a second input would silently delete the first.
		InputPorts ??= EffectiveInputs.Select( x => x.Clone() ).ToList();

		var row = new SubgraphPortInfo
		{
			Name = Unique( InputPorts, name ?? "In" ),
			Type = type,
			Order = InputPorts.Count
		};

		InputPorts.Add( row );
		RebuildPorts();

		return row;
	}

	/// <summary>Add an output row and rebuild the port set.</summary>
	public SubgraphPortInfo AddOutput( string name = null, string type = "float" )
	{
		OutputPorts ??= EffectiveOutputs.Select( x => x.Clone() ).ToList();

		var row = new SubgraphPortInfo
		{
			Name = Unique( OutputPorts, name ?? "Out" ),
			Type = type,
			Order = OutputPorts.Count
		};

		OutputPorts.Add( row );
		RebuildPorts();

		return row;
	}

	/// <summary>Remove an input row by name. Edges to the port it declared become visible ghosts.</summary>
	public bool RemoveInput( string name ) => Remove( InputPorts, name );

	/// <summary>Remove an output row by name. Edges from the port it declared become visible ghosts.</summary>
	public bool RemoveOutput( string name ) => Remove( OutputPorts, name );

	/// <summary>
	/// Rename an input row. The port id follows the name, so any edge attached to the old id becomes a
	/// ghost rather than silently moving to a port that may now mean something else.
	/// </summary>
	public bool RenameInput( string from, string to ) => Rename( InputPorts, from, to );

	/// <summary>Rename an output row. See <see cref="RenameInput"/> for what happens to its edges.</summary>
	public bool RenameOutput( string from, string to ) => Rename( OutputPorts, from, to );

	/// <summary>
	/// Rebuild the port set after the port table was edited in place — by an inspector widget, say.
	/// Ports whose ids survive keep their connections, their resolved types and their literals.
	/// </summary>
	public void PortsChanged() => RebuildPorts();

	bool Remove( List<SubgraphPortInfo> rows, string name )
	{
		if ( rows is null || string.IsNullOrWhiteSpace( name ) ) return false;

		var id = SubgraphLibrary.Identifier( name );
		var removed = rows.RemoveAll( x => x is not null && SubgraphLibrary.Identifier( x.Name ) == id );

		if ( removed == 0 ) return false;

		RebuildPorts();
		return true;
	}

	bool Rename( List<SubgraphPortInfo> rows, string from, string to )
	{
		if ( rows is null || string.IsNullOrWhiteSpace( from ) || string.IsNullOrWhiteSpace( to ) ) return false;

		var id = SubgraphLibrary.Identifier( from );
		var row = rows.FirstOrDefault( x => x is not null && SubgraphLibrary.Identifier( x.Name ) == id );

		if ( row is null ) return false;

		row.Name = Unique( rows, to, row );
		RebuildPorts();

		return true;
	}

	static string Unique( List<SubgraphPortInfo> rows, string name, SubgraphPortInfo ignore = null )
	{
		var baseName = string.IsNullOrWhiteSpace( name ) ? "Value" : name.Trim();
		var candidate = baseName;
		var index = 2;

		while ( rows.Any( x => x is not null && !ReferenceEquals( x, ignore ) &&
			string.Equals( SubgraphLibrary.Identifier( x.Name ), SubgraphLibrary.Identifier( candidate ),
				StringComparison.Ordinal ) ) )
		{
			candidate = $"{baseName}{index++}";
		}

		return candidate;
	}

	/// <inheritdoc/>
	protected override void OnDefinePorts( PortBuilder b )
	{
		if ( b is null ) return;

		var order = 0;

		foreach ( var row in EffectiveInputs.OrderBy( x => x?.Order ?? 0 ) )
		{
			if ( row is null || string.IsNullOrWhiteSpace( row.Name ) ) continue;
			if ( b.Has( row.PortId ) ) continue;

			b.Input( row.PortId, row.ResolvedType.ToString(), row.Name,
				tooltip: row.Description, order: order++,
				flags: row.Required ? PortFlags.Required : PortFlags.None );
		}

		order = 0;

		foreach ( var row in EffectiveOutputs.OrderBy( x => x?.Order ?? 0 ) )
		{
			if ( row is null || string.IsNullOrWhiteSpace( row.Name ) ) continue;
			if ( b.Has( row.PortId ) ) continue;

			b.Output( row.PortId, row.ResolvedType.ToString(), row.Name,
				tooltip: row.Description, order: order++ );
		}
	}

	// ---- stages -----------------------------------------------------------

	/// <inheritdoc/>
	public StageMask RequiredStages => Stage switch
	{
		PrismCodeStage.VertexOnly => StageMask.Vertex,
		PrismCodeStage.PixelOnly => StageMask.Pixel,
		_ => StageMask.All
	};

	/// <inheritdoc/>
	public ShaderStage PreferredStage =>
		Stage == PrismCodeStage.VertexOnly ? ShaderStage.Vertex : ShaderStage.None;

	// ---- naming -----------------------------------------------------------

	/// <summary>The sanitized stem every generated helper name is built from.</summary>
	public string SafeName
	{
		get
		{
			var text = SubgraphLibrary.Identifier(
				string.IsNullOrWhiteSpace( FunctionName ) ? "CustomCode" : FunctionName );

			return text.Length == 0 ? "CustomCode" : text;
		}
	}

	/// <summary>
	/// The name of the helper generated for one output.
	/// <para>
	/// The body's fingerprint is part of the name on purpose. Two nodes that share a name and a body
	/// deduplicate into one function; two that share a name but not a body get different functions
	/// instead of a hard collision error, which is the behaviour a user editing two copies of a snippet
	/// expects.
	/// </para>
	/// </summary>
	public string HelperName( SubgraphPortInfo output ) =>
		$"Prism_{SafeName}_{( output is null ? "Out" : output.PortId )}_{Fingerprint}";

	/// <summary>A stable 8-character fingerprint of everything that affects the generated bodies.</summary>
	public string Fingerprint
	{
		get
		{
			var text = new StringBuilder();

			text.Append( SafeName ).Append( '|' );
			text.Append( Hlsl ?? string.Empty ).Append( '|' );
			text.Append( UseSlangBody ? Slang ?? string.Empty : string.Empty ).Append( '|' );

			foreach ( var row in EffectiveInputs )
			{
				text.Append( row?.PortId ).Append( ':' ).Append( row?.ResolvedType.Hlsl ).Append( ',' );
			}

			text.Append( '|' );

			foreach ( var row in EffectiveOutputs )
			{
				text.Append( row?.PortId ).Append( ':' ).Append( row?.ResolvedType.Hlsl ).Append( ',' );
			}

			return Hash( text.ToString() );
		}
	}

	/// <summary>FNV-1a, rendered as eight lower-case base-36 characters.</summary>
	static string Hash( string text )
	{
		const ulong offset = 14695981039346656037UL;
		const ulong prime = 1099511628211UL;

		var hash = offset;

		foreach ( var c in text ?? string.Empty )
		{
			hash ^= c;
			hash *= prime;
		}

		const string alphabet = "0123456789abcdefghijklmnopqrstuvwxyz";

		var chars = new char[8];

		for ( int i = 7; i >= 0; i-- )
		{
			chars[i] = alphabet[(int)( hash % 36 )];
			hash /= 36;
		}

		return new string( chars );
	}

	// ---- emission ---------------------------------------------------------

	/// <summary>Build the helper that returns one of the declared outputs.</summary>
	public HelperFunction BuildHelper( SubgraphPortInfo output )
	{
		if ( output is null ) return null;

		var parameters = EffectiveInputs
			.Where( x => x is not null && !string.IsNullOrWhiteSpace( x.Name ) )
			.OrderBy( x => x.Order )
			.Select( x => new HelperParam( x.PortId, x.ResolvedType ) )
			.ToArray();

		var name = HelperName( output );

		var hlsl = Compose( name, output, Hlsl );
		var slang = UseSlangBody && !string.IsNullOrWhiteSpace( Slang ) ? Compose( name, output, Slang ) : null;

		return new HelperFunction( name, output.ResolvedType, parameters )
		{
			Hlsl = hlsl,
			Slang = slang,
			// Never null: HelperFunction defaults these to an empty list and the backends read them
			// without a null check.
			Includes = Includes is { Count: > 0 }
				? Includes.Where( x => !string.IsNullOrWhiteSpace( x ) ).ToArray()
				: Array.Empty<string>(),
			Stages = RequiredStages,
			Pure = true
		};
	}

	/// <summary>Wrap an authored body in a signature that returns one declared output.</summary>
	string Compose( string name, SubgraphPortInfo output, string body )
	{
		var text = new StringBuilder();

		text.Append( output.ResolvedType.Hlsl ).Append( ' ' ).Append( name ).Append( "( " );

		var first = true;

		foreach ( var row in EffectiveInputs.OrderBy( x => x?.Order ?? 0 ) )
		{
			if ( row is null || string.IsNullOrWhiteSpace( row.Name ) ) continue;

			if ( !first ) text.Append( ", " );

			text.Append( row.ResolvedType.Hlsl ).Append( ' ' ).Append( row.PortId );
			first = false;
		}

		if ( first ) text.Append( "void" );

		text.Append( " )\r\n{\r\n" );

		foreach ( var row in EffectiveOutputs.OrderBy( x => x?.Order ?? 0 ) )
		{
			if ( row is null || string.IsNullOrWhiteSpace( row.Name ) ) continue;

			text.Append( '\t' ).Append( row.ResolvedType.Hlsl ).Append( ' ' ).Append( row.PortId )
				.Append( " = (" ).Append( row.ResolvedType.Hlsl ).Append( ")0;\r\n" );
		}

		text.Append( "\r\n" );

		foreach ( var line in ( body ?? string.Empty ).Replace( "\r\n", "\n" ).Split( '\n' ) )
		{
			if ( line.Length == 0 )
			{
				text.Append( "\r\n" );
				continue;
			}

			text.Append( '\t' ).Append( line.TrimEnd() ).Append( "\r\n" );
		}

		text.Append( "\r\n\treturn " ).Append( output.PortId ).Append( ";\r\n}" );

		return text.ToString();
	}

	/// <inheritdoc/>
	public override void OnValidate( ValidationContext ctx )
	{
		if ( ctx is null ) return;

		if ( string.IsNullOrWhiteSpace( Hlsl ) )
		{
			ctx.Warn( "This node has no body, so every output reads as zero", null,
				DiagnosticCode.NodeEmitFailed );
		}

		if ( UseSlangBody && string.IsNullOrWhiteSpace( Slang ) )
		{
			ctx.Warn( "A separate Slang body is enabled but empty; the HLSL body will be used instead",
				null, DiagnosticCode.BackendUnsupported );
		}

		Check( ctx, EffectiveInputs, "input" );
		Check( ctx, EffectiveOutputs, "output" );

		if ( EffectiveOutputs.Count != 0 ) return;

		ctx.Error( "This node declares no outputs", null, DiagnosticCode.NoOutput );
	}

	static void Check( ValidationContext ctx, IReadOnlyList<SubgraphPortInfo> rows, string what )
	{
		var seen = new HashSet<string>( StringComparer.Ordinal );

		foreach ( var row in rows )
		{
			if ( row is null ) continue;

			if ( string.IsNullOrWhiteSpace( row.Name ) )
			{
				ctx.Error( $"An {what} row has no name", null, DiagnosticCode.NodeReadFailed );
				continue;
			}

			if ( !seen.Add( row.PortId ) )
			{
				ctx.Error( $"Two {what} rows both resolve to '{row.PortId}'", null,
					DiagnosticCode.NodeReadFailed );
			}

			if ( ShaderType.TryParse( row.Type, out var type ) && !type.IsVoid ) continue;

			ctx.Error( $"'{row.Type}' is not a type this {what} can have", row.PortId,
				DiagnosticCode.UnresolvedType );
		}
	}

	/// <inheritdoc/>
	public override void Emit( EmitContext ctx )
	{
		if ( ctx is null ) return;

		var arguments = new List<IrValue>();
		var ok = true;

		foreach ( var row in EffectiveInputs.OrderBy( x => x?.Order ?? 0 ) )
		{
			if ( row is null || string.IsNullOrWhiteSpace( row.Name ) ) continue;

			var type = row.ResolvedType;

			var value = ctx.TryIn( row.PortId, out var read )
				? read
				: ctx.Const( type, row.DefaultValue );

			if ( !value.IsValid )
			{
				ctx.Error( $"Input '{row.Name}' produced no value", row.PortId );
				ok = false;
				break;
			}

			arguments.Add( value );
		}

		foreach ( var include in Includes ?? new List<string>() )
		{
			if ( string.IsNullOrWhiteSpace( include ) ) continue;

			ctx.Include( include );
		}

		foreach ( var row in EffectiveOutputs.OrderBy( x => x?.Order ?? 0 ) )
		{
			if ( row is null || string.IsNullOrWhiteSpace( row.Name ) ) continue;

			if ( !ok )
			{
				ctx.Out( row.PortId, IrValue.Invalid );
				continue;
			}

			var helper = BuildHelper( row );

			ctx.Out( row.PortId, helper is null
				? IrValue.Invalid
				: ctx.Helper( helper, arguments.ToArray() ) );
		}
	}
}