using Editor.Prism.Core;
using Editor.Prism.Model;

namespace Editor.Prism.Compiler;

/// <summary>
/// Implemented by a node that cannot run everywhere, or that would rather not.
/// <para>
/// Most nodes are stage-agnostic and never implement this. A node that reads a vertex-only attribute
/// returns <c>StageMask.Vertex</c>; one that needs screen-space derivatives returns
/// <c>StageMask.Pixel</c>. <see cref="PreferredStage"/> is a soft hint used to decide whether a value
/// demanded in the pixel stage should instead be computed per vertex and interpolated.
/// </para>
/// </summary>
public interface IStageConstrained
{
	/// <summary>Stages this node can legally run in.</summary>
	StageMask RequiredStages { get; }

	/// <summary>
	/// The stage this node would rather run in, or <see cref="ShaderStage.None"/> for no preference.
	/// Only honoured when the node and its whole dependency subtree are legal there.
	/// </summary>
	ShaderStage PreferredStage { get; }
}

/// <summary>One value the finished shader has to produce, and the stage that has to produce it.</summary>
public readonly record struct StageRoot( ShaderStage Stage, NodeId Node, PortId Port, string Target )
{
	/// <inheritdoc/>
	public override string ToString() => $"{Stage}:{Target} <- {Node}.{Port}";
}

/// <summary>
/// Where every node runs, and which values cross from the vertex stage to the pixel stage.
/// </summary>
public sealed class StagePlan
{
	internal StagePlan(
		StageMask stages,
		IReadOnlyList<StageRoot> roots,
		IReadOnlyDictionary<NodeId, StageMask> demanded,
		IReadOnlyDictionary<NodeId, StageMask> legal,
		IReadOnlyList<PortRef> varyings,
		bool ok )
	{
		Stages = stages;
		Roots = roots;
		Demanded = demanded;
		Legal = legal;
		VaryingPorts = varyings;
		Ok = ok;

		_varyingSet = new HashSet<PortRef>( varyings );
	}

	readonly HashSet<PortRef> _varyingSet;

	/// <summary>An empty plan.</summary>
	public static StagePlan Empty { get; } = new(
		StageMask.None, Array.Empty<StageRoot>(),
		new Dictionary<NodeId, StageMask>(), new Dictionary<NodeId, StageMask>(),
		Array.Empty<PortRef>(), true );

	/// <summary>Every stage the module has to emit.</summary>
	public StageMask Stages { get; }

	/// <summary>The values the shader has to produce.</summary>
	public IReadOnlyList<StageRoot> Roots { get; }

	/// <summary>Which stages each node is reachable from.</summary>
	public IReadOnlyDictionary<NodeId, StageMask> Demanded { get; }

	/// <summary>Which stages each node may legally run in.</summary>
	public IReadOnlyDictionary<NodeId, StageMask> Legal { get; }

	/// <summary>Producer ports whose value travels through an interpolator.</summary>
	public IReadOnlyList<PortRef> VaryingPorts { get; }

	/// <summary>True when no stage conflict had to be reported.</summary>
	public bool Ok { get; }

	/// <summary>The roots belonging to one stage, in declaration order.</summary>
	public IEnumerable<StageRoot> RootsFor( ShaderStage stage ) => Roots.Where( x => x.Stage == stage );

	/// <summary>Stages a node is demanded in.</summary>
	public StageMask DemandedIn( NodeId node ) =>
		Demanded.TryGetValue( node, out var mask ) ? mask : StageMask.None;

	/// <summary>Stages a node may legally run in.</summary>
	public StageMask LegalIn( NodeId node ) =>
		Legal.TryGetValue( node, out var mask ) ? mask : StageMask.All;

	/// <summary>
	/// The single stage a node runs in, or <see cref="ShaderStage.None"/> when it is emitted in more
	/// than one — which is legal, and means the node is emitted once per stage because the expressions
	/// genuinely differ there.
	/// </summary>
	public ShaderStage AssignedTo( NodeId node )
	{
		var mask = DemandedIn( node );
		var stages = mask.Stages().ToArray();

		return stages.Length == 1 ? stages[0] : ShaderStage.None;
	}

	/// <summary>True when this producer port's value crosses stages through an interpolator.</summary>
	public bool IsVarying( NodeId node, PortId port ) => _varyingSet.Contains( new PortRef( node, port ) );

	/// <summary>True when this producer port's value crosses stages through an interpolator.</summary>
	public bool IsVarying( PortRef port ) => _varyingSet.Contains( port );
}

/// <summary>
/// Decides where each node runs and what has to become an interpolator.
/// <para>
/// The built-in editor supports exactly two stages, runs only the position offset in the vertex stage,
/// and offers no way at all to pass a vertex-computed value to the pixel stage. This planner walks
/// back from every active output, and when it reaches a node that cannot run in the stage that
/// demanded it — a vertex-attribute read demanded by the pixel stage, for instance — it inserts a
/// varying boundary and keeps walking in the stage that <em>can</em> run it. The same mechanism moves
/// expensive stage-invariant work off the pixel shader when a node asks for it.
/// </para>
/// </summary>
public sealed class StagePlanner
{
	readonly IPrismGraph _graph;
	readonly IReadOnlyList<StageRoot> _roots;
	readonly DiagnosticSink _diagnostics;

	readonly Dictionary<NodeId, StageMask> _demanded = new();
	readonly Dictionary<NodeId, StageMask> _legal = new();
	readonly HashSet<(PortRef Port, ShaderStage Stage)> _visited = new();
	readonly HashSet<(NodeId Node, ShaderStage Stage)> _active = new();
	readonly Dictionary<NodeId, bool> _vertexSafe = new();
	readonly List<PortRef> _varyings = new();
	readonly HashSet<PortRef> _varyingSet = new();

	Dictionary<PortRef, List<Edge>> _incoming = new();

	bool _ok = true;

	/// <summary>Build a planner for one graph and one set of roots.</summary>
	public StagePlanner( IPrismGraph graph, IReadOnlyList<StageRoot> roots, DiagnosticSink diagnostics )
	{
		_graph = graph;
		_roots = roots ?? Array.Empty<StageRoot>();
		_diagnostics = diagnostics ?? new DiagnosticSink();
	}

	/// <summary>
	/// Allow values to be moved into the vertex stage and interpolated. Turning this off forces
	/// everything into the stage that demanded it, which is what a compute or post-process graph wants.
	/// </summary>
	public bool AllowVaryingPromotion { get; set; } = true;

	/// <summary>Plan a graph in one call.</summary>
	public static StagePlan Plan( IPrismGraph graph, IReadOnlyList<StageRoot> roots, DiagnosticSink diagnostics ) =>
		new StagePlanner( graph, roots, diagnostics ).Plan();

	/// <summary>Run the planner.</summary>
	public StagePlan Plan()
	{
		if ( _graph?.Nodes is null ) return StagePlan.Empty;

		_incoming = GraphIndex.IncomingEdges( _graph );

		foreach ( var node in _graph.Nodes )
		{
			if ( node is null ) continue;

			_legal[node.Id] = LegalStages( node );
		}

		var stages = StageMask.None;

		foreach ( var root in _roots )
		{
			stages |= root.Stage.ToMask();

			var node = _graph.FindNode( root.Node );
			if ( node is null ) continue;

			Demand( root.Node, root.Stage );

			foreach ( var edge in IncomingEdges( root.Node, root.Port ) )
			{
				Walk( edge.From, root.Stage );
			}
		}

		return new StagePlan( stages, _roots, _demanded, _legal, _varyings, _ok );
	}

	// ---- traversal --------------------------------------------------------

	void Walk( PortRef producer, ShaderStage stage )
	{
		if ( !_visited.Add( (producer, stage) ) ) return;

		// One managed frame per node in the dependency chain, exactly like NodeEmitter.Demand — and a
		// StackOverflowException is uncatchable, so a deep enough graph would take the editor down here
		// before emission ever ran. Ask the runtime for room rather than guessing a depth. The plan is
		// left incomplete for the part that was not reached; the emitter's own guard then reports it
		// against the node the user can actually see.
		if ( !System.Runtime.CompilerServices.RuntimeHelpers.TryEnsureSufficientExecutionStack() ) return;

		var node = _graph.FindNode( producer.Node );
		if ( node is null ) return;

		var effective = Resolve( node, producer, stage );

		if ( !_active.Add( (node.Id, effective) ) ) return;

		Demand( node.Id, effective );

		foreach ( var input in node.Inputs )
		{
			foreach ( var edge in IncomingEdges( node.Id, input.Id ) )
			{
				Walk( edge.From, effective );
			}
		}

		_active.Remove( (node.Id, effective) );
	}

	/// <summary>
	/// Decide which stage a node actually runs in when <paramref name="stage"/> asked for its value,
	/// inserting a varying boundary when that is what makes the demand satisfiable or cheaper.
	/// </summary>
	ShaderStage Resolve( PrismNode node, PortRef producer, ShaderStage stage )
	{
		var legal = LegalIn( node.Id );

		if ( !legal.Contains( stage ) )
		{
			if ( stage == ShaderStage.Pixel && legal.Contains( ShaderStage.Vertex ) && AllowVaryingPromotion )
			{
				AddVarying( producer );
				return ShaderStage.Vertex;
			}

			_ok = false;

			_diagnostics.Error( DiagnosticCode.StageViolation,
				$"'{Describe( node )}' cannot run in the {stage.DisplayName().ToLowerInvariant()} stage",
				GraphRef.ForNode( node.Id ),
				$"Legal stages: {legal}" );

			return stage;
		}

		if ( stage != ShaderStage.Pixel || !AllowVaryingPromotion ) return stage;
		if ( node is not IStageConstrained constrained ) return stage;
		if ( constrained.PreferredStage != ShaderStage.Vertex ) return stage;
		if ( !IsVertexSafe( node.Id ) ) return stage;

		AddVarying( producer );
		return ShaderStage.Vertex;
	}

	void AddVarying( PortRef producer )
	{
		if ( !_varyingSet.Add( producer ) ) return;

		_varyings.Add( producer );
	}

	void Demand( NodeId node, ShaderStage stage )
	{
		_demanded.TryGetValue( node, out var mask );
		_demanded[node] = mask | stage.ToMask();
	}

	/// <summary>True when this node and everything it depends on can run in the vertex stage.</summary>
	bool IsVertexSafe( NodeId id )
	{
		if ( _vertexSafe.TryGetValue( id, out var cached ) ) return cached;

		// Assume safe while recursing so a cycle cannot make this hang or flip-flop.
		_vertexSafe[id] = true;

		var node = _graph.FindNode( id );
		var safe = node is not null && LegalIn( id ).Contains( ShaderStage.Vertex );

		if ( safe )
		{
			foreach ( var input in node.Inputs )
			{
				foreach ( var edge in IncomingEdges( id, input.Id ) )
				{
					if ( IsVertexSafe( edge.FromNode ) ) continue;

					safe = false;
					break;
				}

				if ( !safe ) break;
			}
		}

		_vertexSafe[id] = safe;
		return safe;
	}

	StageMask LegalIn( NodeId node ) => _legal.TryGetValue( node, out var mask ) ? mask : StageMask.All;

	static StageMask LegalStages( PrismNode node )
	{
		if ( node is IStageConstrained constrained )
		{
			var mask = constrained.RequiredStages;
			if ( mask != StageMask.None ) return mask;
		}

		return StageMask.All;
	}

	IReadOnlyList<Edge> IncomingEdges( NodeId node, PortId port ) =>
		_incoming.TryGetValue( new PortRef( node, port ), out var edges ) ? edges : Array.Empty<Edge>();

	static string Describe( PrismNode node )
	{
		var title = node.Descriptor?.Title;
		return string.IsNullOrEmpty( title ) ? node.GetType().Name : title;
	}
}