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

namespace Editor.Prism.Compiler;

/// <summary>
/// The result of running <see cref="TypeSolver"/> over a graph: a concrete <see cref="ShaderType"/>
/// for every port, the conversion each edge performs, and a topological node order the rest of the
/// pipeline can reuse.
/// </summary>
public sealed class TypeSolution
{
	internal TypeSolution(
		IReadOnlyDictionary<PortRef, ShaderType> types,
		IReadOnlyDictionary<EdgeId, ConversionKind> conversions,
		IReadOnlyList<NodeId> order,
		int unresolved,
		bool ok )
	{
		Types = types;
		Conversions = conversions;
		TopologicalOrder = order;
		UnresolvedCount = unresolved;
		Ok = ok;
	}

	/// <summary>An empty solution, used when there is nothing to solve.</summary>
	public static TypeSolution Empty { get; } = new(
		new Dictionary<PortRef, ShaderType>(), new Dictionary<EdgeId, ConversionKind>(),
		Array.Empty<NodeId>(), 0, true );

	/// <summary>True when every port resolved and no unification failed.</summary>
	public bool Ok { get; }

	/// <summary>How many ports had to fall back to a default type.</summary>
	public int UnresolvedCount { get; }

	/// <summary>The solved type of every port in the graph.</summary>
	public IReadOnlyDictionary<PortRef, ShaderType> Types { get; }

	/// <summary>The conversion each edge performs, for wire markers and tooltips.</summary>
	public IReadOnlyDictionary<EdgeId, ConversionKind> Conversions { get; }

	/// <summary>
	/// Producers before consumers. Nodes caught in a cycle are appended at the end in document order,
	/// so this is always a total order even for a malformed graph.
	/// </summary>
	public IReadOnlyList<NodeId> TopologicalOrder { get; }

	/// <summary>The solved type of one port, or void when it is not in the solution.</summary>
	public ShaderType TypeOf( NodeId node, PortId port ) =>
		Types.TryGetValue( new PortRef( node, port ), out var type ) ? type : ShaderType.Void;

	/// <summary>The solved type of one port.</summary>
	public ShaderType TypeOf( Port port ) =>
		port is null ? ShaderType.Void : TypeOf( port.Node?.Id ?? NodeId.None, port.Id );

	/// <summary>The conversion an edge performs, or <see cref="ConversionKind.Identity"/> when unknown.</summary>
	public ConversionKind ConversionOn( EdgeId edge ) =>
		Conversions.TryGetValue( edge, out var kind ) ? kind : ConversionKind.Identity;
}

/// <summary>
/// Hindley–Milner-lite unification over a whole graph.
/// <para>
/// A port's declared type is either a concrete spelling (<c>float3</c>, <c>Texture2D</c>) or a term in
/// a small algebra: <c>T</c> — a variable shared by every port on the node that names it; <c>T.scalar</c>
/// — the component type of <c>T</c>; <c>vecN</c> — a float vector whose width unifies; <c>float{N}</c> —
/// a float vector sharing the width variable <c>N</c>; <c>any</c> — a passthrough that adopts whatever
/// reaches it. Every unrecognised spelling is treated as a fresh variable named after itself, so
/// <c>U</c> and <c>Element</c> work exactly like <c>T</c>.
/// </para>
/// <para>
/// The solver runs forward along the topological order, then backward for anything still open, then
/// defaults what is left to <c>float</c>. Afterwards every <see cref="Port.ResolvedType"/> is concrete
/// and every edge has been classified, which is what lets the IR be built without a single type guess.
/// </para>
/// </summary>
public sealed class TypeSolver
{
	const string PassthroughGroup = "
passthrough";

	readonly IPrismGraph _graph;
	readonly DiagnosticSink _diagnostics;

	readonly List<Slot> _slots = new();
	readonly Dictionary<(NodeId Node, string Name), int> _vars = new();
	readonly Dictionary<PortRef, Term> _terms = new();

	bool _failed;

	/// <summary>Build a solver for one graph.</summary>
	public TypeSolver( IPrismGraph graph, DiagnosticSink diagnostics )
	{
		_graph = graph;
		_diagnostics = diagnostics ?? new DiagnosticSink();
	}

	/// <summary>How many forward/backward sweeps to run before giving up on convergence.</summary>
	public int MaxIterations { get; set; } = 8;

	/// <summary>Write the solved types back onto <see cref="Port.ResolvedType"/>. On by default.</summary>
	public bool ApplyToPorts { get; set; } = true;

	/// <summary>Report warnings for lossy and padded edge conversions. On by default.</summary>
	public bool ReportConversions { get; set; } = true;

	/// <summary>Solve a graph in one call.</summary>
	public static TypeSolution Solve( IPrismGraph graph, DiagnosticSink diagnostics ) =>
		new TypeSolver( graph, diagnostics ).Solve();

	/// <summary>Run the solver.</summary>
	public TypeSolution Solve()
	{
		if ( _graph?.Nodes is null || _graph.Nodes.Count == 0 ) return TypeSolution.Empty;

		var order = TopologicalOrder( _graph );

		Seed();

		var edges = ValidEdges().ToArray();

		for ( int pass = 0; pass < Math.Max( 1, MaxIterations ); pass++ )
		{
			var changed = Forward( order, edges );
			changed |= Backward( edges );

			if ( !changed ) break;
		}

		DefaultUnresolved();

		var types = new Dictionary<PortRef, ShaderType>();
		var unresolved = 0;

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

			foreach ( var port in AllPorts( node ) )
			{
				var key = new PortRef( node.Id, port.Id );
				var type = Read( key );

				if ( type.IsVoid )
				{
					type = ShaderType.Float;
					unresolved++;

					if ( !port.Def.IsGeneric )
					{
						// A concrete declaration that came back void means the spelling is unparseable.
						_diagnostics.Warn( DiagnosticCode.UnresolvedType,
							$"Port '{port.DisplayName}' declares an unrecognised type '{port.DeclaredType}'; assuming float",
							GraphRef.ForPort( node.Id, port.Id ) );
					}
				}

				types[key] = type;

				if ( ApplyToPorts ) port.ResolvedType = type;
			}
		}

		var conversions = ClassifyEdges( edges, types );

		return new TypeSolution( types, conversions, order, unresolved, !_failed );
	}

	/// <summary>
	/// Producers before consumers, cycles appended in document order. Kahn's algorithm, so a cyclic
	/// graph degrades into a stable-but-arbitrary order instead of hanging.
	/// </summary>
	public static IReadOnlyList<NodeId> TopologicalOrder( IPrismGraph graph )
	{
		if ( graph?.Nodes is null ) return Array.Empty<NodeId>();

		var indegree = new Dictionary<NodeId, int>();
		var successors = new Dictionary<NodeId, List<NodeId>>();

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

			indegree.TryAdd( node.Id, 0 );
			successors.TryAdd( node.Id, new List<NodeId>() );
		}

		foreach ( var edge in graph.Edges ?? Array.Empty<Edge>() )
		{
			if ( edge is null || !edge.IsValid ) continue;
			if ( !indegree.ContainsKey( edge.FromNode ) || !indegree.ContainsKey( edge.ToNode ) ) continue;
			if ( edge.FromNode == edge.ToNode ) continue;

			successors[edge.FromNode].Add( edge.ToNode );
			indegree[edge.ToNode] = indegree[edge.ToNode] + 1;
		}

		// Seed in document order so the result is deterministic run to run.
		var ready = new List<NodeId>();

		foreach ( var node in graph.Nodes )
		{
			if ( node is null ) continue;
			if ( indegree[node.Id] == 0 ) ready.Add( node.Id );
		}

		var order = new List<NodeId>( indegree.Count );
		var cursor = 0;

		while ( cursor < ready.Count )
		{
			var id = ready[cursor++];
			order.Add( id );

			foreach ( var next in successors[id] )
			{
				var remaining = indegree[next] - 1;
				indegree[next] = remaining;

				if ( remaining == 0 ) ready.Add( next );
			}
		}

		if ( order.Count < indegree.Count )
		{
			var seen = new HashSet<NodeId>( order );

			foreach ( var node in graph.Nodes )
			{
				if ( node is null || seen.Contains( node.Id ) ) continue;

				order.Add( node.Id );
			}
		}

		return order;
	}

	// ---- seeding ----------------------------------------------------------

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

			foreach ( var port in AllPorts( node ) )
			{
				_terms[new PortRef( node.Id, port.Id )] = MakeTerm( node.Id, port );
			}
		}
	}

	Term MakeTerm( NodeId node, Port port )
	{
		var declared = port.DeclaredType;

		if ( ( port.Flags & PortFlags.Passthrough ) != 0 )
		{
			return Term.Variable( VarSlot( node, PassthroughGroup ), null );
		}

		if ( !TypeRules.IsTypeVariable( declared ) && ShaderType.TryParse( declared, out var concrete ) )
		{
			return Term.Fixed( concrete );
		}

		var text = ( declared ?? string.Empty ).Trim();

		if ( text.Length == 0 || text == TypeRules.TypeVarAny )
		{
			return Term.Variable( VarSlot( node, PassthroughGroup ), null );
		}

		if ( text == TypeRules.TypeVarVecN )
		{
			return Term.Variable( VarSlot( node, TypeRules.TypeVarVecN ), ScalarKind.Float );
		}

		// "T.scalar" — the component type of another variable on the same node.
		var dot = text.IndexOf( '.' );

		if ( dot > 0 && text[( dot + 1 )..].Equals( "scalar", StringComparison.OrdinalIgnoreCase ) )
		{
			return Term.ScalarOf( VarSlot( node, text[..dot] ) );
		}

		// "float{N}" — a vector of the shared width variable N, with the component kind pinned.
		var open = text.IndexOf( '{' );
		var close = text.IndexOf( '}' );

		if ( open > 0 && close > open + 1 )
		{
			var prefix = text[..open];
			var width = text[( open + 1 )..close];
			var scalar = ShaderType.TryParse( prefix, out var prefixType ) && prefixType.IsNumeric
				? prefixType.Scalar
				: ScalarKind.Float;

			return Term.Variable( VarSlot( node, width ), scalar );
		}

		return Term.Variable( VarSlot( node, text ), null );
	}

	int VarSlot( NodeId node, string name )
	{
		var key = (node, name ?? string.Empty);

		if ( _vars.TryGetValue( key, out var index ) ) return index;

		index = _slots.Count;
		_slots.Add( new Slot() );
		_vars[key] = index;

		return index;
	}

	// ---- propagation ------------------------------------------------------

	bool Forward( IReadOnlyList<NodeId> order, IReadOnlyList<Edge> edges )
	{
		var incoming = new Dictionary<PortRef, List<Edge>>();

		foreach ( var edge in edges )
		{
			var key = edge.To;

			if ( !incoming.TryGetValue( key, out var list ) )
			{
				list = new List<Edge>();
				incoming[key] = list;
			}

			list.Add( edge );
		}

		var changed = false;

		foreach ( var id in order )
		{
			var node = _graph.FindNode( id );
			if ( node is null ) continue;

			foreach ( var input in node.Inputs )
			{
				var key = new PortRef( id, input.Id );

				if ( !incoming.TryGetValue( key, out var sources ) ) continue;

				foreach ( var edge in sources )
				{
					var produced = Read( edge.From );
					if ( produced.IsVoid ) continue;

					changed |= Constrain( key, produced, GraphRef.ForPort( id, input.Id ), edge );
				}
			}
		}

		return changed;
	}

	bool Backward( IReadOnlyList<Edge> edges )
	{
		var changed = false;

		for ( int i = edges.Count - 1; i >= 0; i-- )
		{
			var edge = edges[i];
			var consumed = Read( edge.To );

			if ( consumed.IsVoid ) continue;
			if ( !Read( edge.From ).IsVoid ) continue;

			changed |= Constrain( edge.From, consumed, GraphRef.ForPort( edge.FromNode, edge.FromPort ), edge );
		}

		return changed;
	}

	void DefaultUnresolved()
	{
		for ( int i = 0; i < _slots.Count; i++ )
		{
			var root = _slots[i];

			if ( root.Type.IsVoid ) root.Type = ShaderType.Float;
		}
	}

	// ---- term access ------------------------------------------------------

	ShaderType Read( PortRef port )
	{
		if ( !_terms.TryGetValue( port, out var term ) ) return ShaderType.Void;

		switch ( term.Kind )
		{
			case TermKind.Fixed:
				return term.Concrete;

			case TermKind.Variable:
			{
				var type = _slots[term.Slot].Type;

				if ( type.IsVoid ) return ShaderType.Void;

				return term.Force.HasValue && type.IsNumeric ? type.WithScalar( term.Force.Value ) : type;
			}

			case TermKind.ScalarOf:
			{
				var type = _slots[term.Slot].Type;
				return type.IsVoid ? ShaderType.Void : type.ScalarType;
			}

			default:
				return ShaderType.Void;
		}
	}

	bool Constrain( PortRef port, ShaderType incoming, GraphRef where, Edge edge )
	{
		if ( incoming.IsVoid ) return false;
		if ( !_terms.TryGetValue( port, out var term ) ) return false;

		switch ( term.Kind )
		{
			case TermKind.Fixed:
				return false;

			case TermKind.Variable:
			{
				var wanted = term.Force.HasValue && incoming.IsNumeric
					? incoming.WithScalar( term.Force.Value )
					: incoming;

				return Bind( term.Slot, wanted, where, edge );
			}

			case TermKind.ScalarOf:
			{
				var slot = _slots[term.Slot];
				var wanted = slot.Type.IsVoid
					? incoming.ScalarType
					: slot.Type.WithScalar( TypeRules.PromoteScalar( slot.Type.Scalar, incoming.Scalar ) );

				return Bind( term.Slot, wanted, where, edge );
			}

			default:
				return false;
		}
	}

	bool Bind( int slotIndex, ShaderType incoming, GraphRef where, Edge edge )
	{
		if ( incoming.IsVoid ) return false;

		var slot = _slots[slotIndex];

		if ( slot.Type.IsVoid )
		{
			slot.Type = incoming;
			return true;
		}

		if ( slot.Type == incoming ) return false;

		if ( !TypeRules.Unify( slot.Type, incoming, out var unified ) )
		{
			if ( slot.Failed ) return false;

			slot.Failed = true;
			_failed = true;

			var detail = edge is null ? null : $"Connection {edge}";

			_diagnostics.Error( DiagnosticCode.UnificationFailure,
				$"Cannot reconcile {slot.Type.Hlsl} and {incoming.Hlsl} on the same generic port group",
				where, detail );

			return false;
		}

		if ( unified == slot.Type ) return false;

		slot.Type = unified;
		return true;
	}

	// ---- edges ------------------------------------------------------------

	IEnumerable<Edge> ValidEdges()
	{
		foreach ( var edge in _graph.Edges ?? Array.Empty<Edge>() )
		{
			if ( edge is null || !edge.IsValid ) continue;
			if ( !_terms.ContainsKey( edge.From ) || !_terms.ContainsKey( edge.To ) ) continue;

			yield return edge;
		}
	}

	Dictionary<EdgeId, ConversionKind> ClassifyEdges( IReadOnlyList<Edge> edges,
		IReadOnlyDictionary<PortRef, ShaderType> types )
	{
		var conversions = new Dictionary<EdgeId, ConversionKind>();

		foreach ( var edge in edges )
		{
			if ( !types.TryGetValue( edge.From, out var from ) ) continue;
			if ( !types.TryGetValue( edge.To, out var to ) ) continue;

			var kind = TypeRules.Classify( from, to );
			conversions[edge.Id] = kind;

			if ( !ReportConversions ) continue;

			var where = new GraphRef( edge.ToNode, edge.ToPort, edge.Id );

			switch ( kind )
			{
				case ConversionKind.Illegal:
					_failed = true;
					_diagnostics.Error( DiagnosticCode.IllegalConversion,
						$"{from.Hlsl} cannot connect to {to.Hlsl}", where,
						TypeRules.Describe( from, to, kind ) );
					break;

				case ConversionKind.Truncate:
					_diagnostics.Warn( DiagnosticCode.LossyConversion,
						$"{from.Hlsl} narrows to {to.Hlsl}", where,
						TypeRules.Describe( from, to, kind ) );
					break;

				case ConversionKind.Pad:
					var fill = edge.Fill ?? TypeRules.DefaultFill( from, to, to.Components - 1 );
					_diagnostics.Warn( DiagnosticCode.PaddedConversion,
						$"{from.Hlsl} widens to {to.Hlsl}, filling with {fill}", where,
						TypeRules.Describe( from, to, kind ) );
					break;
			}
		}

		return conversions;
	}

	static IEnumerable<Port> AllPorts( PrismNode node )
	{
		foreach ( var input in node.Inputs ) yield return input;
		foreach ( var output in node.Outputs ) yield return output;
	}

	enum TermKind
	{
		Fixed,
		Variable,
		ScalarOf
	}

	readonly struct Term
	{
		Term( TermKind kind, ShaderType concrete, int slot, ScalarKind? force )
		{
			Kind = kind;
			Concrete = concrete;
			Slot = slot;
			Force = force;
		}

		public TermKind Kind { get; }
		public ShaderType Concrete { get; }
		public int Slot { get; }
		public ScalarKind? Force { get; }

		public static Term Fixed( ShaderType type ) => new( TermKind.Fixed, type, -1, null );
		public static Term Variable( int slot, ScalarKind? force ) => new( TermKind.Variable, ShaderType.Void, slot, force );
		public static Term ScalarOf( int slot ) => new( TermKind.ScalarOf, ShaderType.Void, slot, null );
	}

	/// <summary>
	/// One type variable's current binding.
	/// <para>
	/// There is deliberately no union-find here. Prism's type algebra is per-(node, variable name):
	/// <c>VarSlot</c> mints one slot for each and nothing ever merges two, because a constraint that
	/// spans nodes is expressed by propagating a concrete type along an edge rather than by equating two
	/// variables. The class used to carry a <c>Parent</c> field and a path-compressing <c>Find</c> that
	/// could only ever return its own argument — it read as Hindley-Milner and behaved as a lookup, and
	/// a later pass adding a real cross-node constraint would have assumed the merging worked. If one is
	/// ever needed, add <c>Union</c> and <c>Find</c> together.
	/// </para>
	/// </summary>
	sealed class Slot
	{
		public ShaderType Type;
		public bool Failed;
	}
}