using Editor.Prism.Core;
using System.Text;

namespace Editor.Prism.Compiler.Ir;

/// <summary>
/// Builds the IR for one shader stage.
/// <para>
/// Two things make this more than a factory. First, <b>hash-consing</b>: every expression handed out
/// is interned in a structural-equality pool, so two structurally identical trees — even when they
/// come from two unrelated nodes — are the same object. Second, <b>common-subexpression binding</b>:
/// <see cref="Bind"/> looks the interned expression up in a scoped temp table and reuses the temp that
/// already holds it instead of emitting a second declaration. Together those give real CSE across the
/// whole graph rather than the per-connection memo a naive emitter manages.
/// </para>
/// <para>
/// The block stack keeps temps honest: a temp bound inside an <c>if</c> is dropped from the CSE table
/// when that block closes, so it can never be referenced from a sibling branch.
/// </para>
/// </summary>
public sealed class IrBuilder
{
	readonly Dictionary<IrExpr, IrExpr> _pool = new( StructuralComparer.Instance );
	readonly List<Frame> _frames = new();
	readonly HashSet<string> _names = new( StringComparer.Ordinal );

	int _tempSerial;

	/// <summary>Start a builder for one stage.</summary>
	public IrBuilder( ShaderStage stage, bool debugSymbols = false )
	{
		Stage = stage;
		DebugSymbols = debugSymbols;
		Root = new IrBlock();

		_frames.Add( new Frame( Root ) );
	}

	/// <summary>The stage whose code this builder is assembling.</summary>
	public ShaderStage Stage { get; }

	/// <summary>When set, temps get descriptive names and per-node comments are worth emitting.</summary>
	public bool DebugSymbols { get; set; }

	/// <summary>The outermost block. Statements land here unless a nested block is open.</summary>
	public IrBlock Root { get; }

	/// <summary>The block statements are currently being appended to.</summary>
	public IrBlock Current => _frames[^1].Block;

	/// <summary>How many blocks are open above the root.</summary>
	public int Depth => _frames.Count - 1;

	/// <summary>Temps declared so far.</summary>
	public int TempCount => _tempSerial;

	/// <summary>Distinct expressions interned so far.</summary>
	public int InternedCount => _pool.Count;

	/// <summary>Number of times a binding request was answered by an existing temp.</summary>
	public int CseHits { get; private set; }

	/// <summary>Number of times an interning request found an identical expression already in the pool.</summary>
	public int InternHits { get; private set; }

	// ---- interning --------------------------------------------------------

	/// <summary>
	/// Return the canonical instance of an expression, adding it to the pool when it is new. The
	/// returned object is reference-comparable against any structurally identical expression built
	/// through this builder.
	/// </summary>
	public IrExpr Intern( IrExpr expr )
	{
		if ( expr is null ) return null;

		if ( _pool.TryGetValue( expr, out var existing ) )
		{
			InternHits++;
			return existing;
		}

		var hashed = expr.Hash != 0 ? expr : expr with { Hash = StructuralHash( expr ) };

		_pool[hashed] = hashed;
		return hashed;
	}

	IrValue Wrap( IrExpr expr, ShaderType type ) => IrValue.Of( Intern( expr ), type );

	// ---- expression construction ------------------------------------------

	/// <summary>A literal.</summary>
	public IrValue Const( ShaderType type, ConstValue value ) => Wrap( new IrConst( type, value ), type );

	/// <summary>A reference to a named local.</summary>
	public IrValue Var( ShaderType type, string name ) => Wrap( new IrVar( type, name ), type );

	/// <summary>A reference to a module-level declaration.</summary>
	public IrValue GlobalRef( GlobalDecl decl )
	{
		if ( decl is null ) return IrValue.Invalid;

		return Wrap( new IrGlobalRef( decl.Type, decl ), decl.Type );
	}

	/// <summary>A reference to an environment-provided value.</summary>
	public IrValue BuiltinRef( Builtin id )
	{
		var type = Builtins.TypeOf( id );
		if ( type.IsVoid ) return IrValue.Invalid;

		return Wrap( new IrBuiltinRef( type, id ), type );
	}

	/// <summary>A call to a canonical intrinsic.</summary>
	public IrValue Call( ShaderType type, Intrinsic id, params IrValue[] args )
	{
		var operands = Unwrap( args );
		if ( operands is null ) return IrValue.Invalid;

		var pure = IntrinsicCatalog.Get( id ).Pure && AllPure( operands );

		return Wrap( new IrCall( type, id, operands ) { Pure = pure }, type );
	}

	/// <summary>A call to a shared helper function.</summary>
	public IrValue HelperCall( HelperFunction fn, params IrValue[] args )
	{
		if ( fn is null ) return IrValue.Invalid;

		var operands = Unwrap( args );
		if ( operands is null ) return IrValue.Invalid;

		var pure = fn.Pure && AllPure( operands );

		return Wrap( new IrHelperCall( fn.ReturnType, fn, operands ) { Pure = pure }, fn.ReturnType );
	}

	/// <summary>A binary operation. Operand promotion is the caller's job.</summary>
	public IrValue Binary( ShaderType type, BinaryOp op, IrValue l, IrValue r )
	{
		if ( !l.IsValid || !r.IsValid ) return IrValue.Invalid;

		var pure = l.Expr.Pure && r.Expr.Pure;

		return Wrap( new IrBinary( type, op, l.Expr, r.Expr ) { Pure = pure }, type );
	}

	/// <summary>A unary operation.</summary>
	public IrValue Unary( ShaderType type, UnaryOp op, IrValue v )
	{
		if ( !v.IsValid ) return IrValue.Invalid;

		return Wrap( new IrUnary( type, op, v.Expr ) { Pure = v.Expr.Pure }, type );
	}

	/// <summary>A swizzle. The mask must already be normalised to the <c>xyzw</c> alphabet.</summary>
	public IrValue Swizzle( ShaderType type, IrValue v, string mask )
	{
		if ( !v.IsValid || string.IsNullOrEmpty( mask ) ) return IrValue.Invalid;

		return Wrap( new IrSwizzle( type, v.Expr, mask ) { Pure = v.Expr.Pure }, type );
	}

	/// <summary>Construction of a vector, matrix or struct from parts.</summary>
	public IrValue Construct( ShaderType type, params IrValue[] parts )
	{
		var operands = Unwrap( parts );
		if ( operands is null ) return IrValue.Invalid;

		return Wrap( new IrConstruct( type, operands ) { Pure = AllPure( operands ) }, type );
	}

	/// <summary>A conversion.</summary>
	public IrValue Cast( ShaderType type, IrValue v, CastKind kind, float fill = 0f )
	{
		if ( !v.IsValid ) return IrValue.Invalid;
		if ( v.Type == type ) return v;

		return Wrap( new IrCast( type, v.Expr, kind ) { Fill = fill, Pure = v.Expr.Pure }, type );
	}

	/// <summary>A component-wise select.</summary>
	public IrValue Select( ShaderType type, IrValue c, IrValue a, IrValue b )
	{
		if ( !c.IsValid || !a.IsValid || !b.IsValid ) return IrValue.Invalid;

		var pure = c.Expr.Pure && a.Expr.Pure && b.Expr.Pure;

		return Wrap( new IrSelect( type, c.Expr, a.Expr, b.Expr ) { Pure = pure }, type );
	}

	/// <summary>An indexing expression.</summary>
	public IrValue Index( ShaderType type, IrValue v, IrValue i )
	{
		if ( !v.IsValid || !i.IsValid ) return IrValue.Invalid;

		return Wrap( new IrIndex( type, v.Expr, i.Expr ) { Pure = v.Expr.Pure && i.Expr.Pure }, type );
	}

	/// <summary>A struct member access.</summary>
	public IrValue Member( ShaderType type, IrValue v, string field )
	{
		if ( !v.IsValid || string.IsNullOrEmpty( field ) ) return IrValue.Invalid;

		return Wrap( new IrMember( type, v.Expr, field ) { Pure = v.Expr.Pure }, type );
	}

	// ---- binding ----------------------------------------------------------

	/// <summary>
	/// True when a value is worth pulling into its own temp. Leaves and one-hop views of leaves read
	/// better inline; everything else becomes a declaration so it is computed once and so the source
	/// map has a line to point at.
	/// </summary>
	public bool ShouldBind( IrValue value )
	{
		if ( !value.IsValid ) return false;
		if ( value.Type.IsVoid || value.Type.IsObject ) return false;

		return value.Expr switch
		{
			IrConst => false,
			IrVar => false,
			IrGlobalRef => false,
			IrBuiltinRef => false,
			IrSwizzle s => s.V is not ( IrVar or IrGlobalRef or IrBuiltinRef or IrConst ),
			IrMember m => m.V is not ( IrVar or IrGlobalRef or IrBuiltinRef ),
			IrCast c => c.V is not ( IrVar or IrGlobalRef or IrBuiltinRef or IrConst ),
			_ => true
		};
	}

	/// <summary>
	/// Bind a value to a temp when it is worth one, reusing the existing temp for a structurally
	/// identical pure expression that is still in scope. Values not worth a temp are returned as-is.
	/// </summary>
	public IrValue Bind( NodeId origin, string hint, IrValue value )
	{
		if ( !value.IsValid ) return value;
		if ( !ShouldBind( value ) ) return value;

		return BindForced( origin, hint, value );
	}

	/// <summary>
	/// Bind a value to a temp whether or not it is worth one. This is what <c>EmitContext.Let</c> does:
	/// the node is asserting the value deserves a readable name.
	/// </summary>
	public IrValue BindForced( NodeId origin, string hint, IrValue value )
	{
		if ( !value.IsValid ) return value;
		if ( value.Type.IsVoid ) return value;

		var expr = Intern( value.Expr );

		if ( expr.Pure && TryFindBinding( expr, out var existing ) )
		{
			CseHits++;
			return existing;
		}

		var name = NewTempName( hint );
		var declared = value.Type.IsVoid ? expr.Type : value.Type;

		Current.Add( new IrDecl( origin, name, declared, expr ) );

		var temp = Var( declared, name );

		if ( expr.Pure ) _frames[^1].Bindings[expr] = temp;

		return temp;
	}

	/// <summary>Forget every cached binding without dropping the interning pool. Used between functions.</summary>
	public void ClearBindings()
	{
		foreach ( var frame in _frames ) frame.Bindings.Clear();
	}

	bool TryFindBinding( IrExpr expr, out IrValue value )
	{
		for ( int i = _frames.Count - 1; i >= 0; i-- )
		{
			if ( _frames[i].Bindings.TryGetValue( expr, out value ) ) return true;
		}

		value = IrValue.Invalid;
		return false;
	}

	/// <summary>Mint a unique local name. Descriptive in debug-symbol mode, terse otherwise.</summary>
	public string NewTempName( string hint )
	{
		string name;

		if ( DebugSymbols && !string.IsNullOrWhiteSpace( hint ) )
		{
			name = $"{Sanitize( hint )}_{_tempSerial++}";
		}
		else
		{
			name = $"t{_tempSerial++}";
		}

		while ( !_names.Add( name ) )
		{
			name = $"{name}_{_tempSerial++}";
		}

		return name;
	}

	// ---- statements -------------------------------------------------------

	/// <summary>Append a statement to the current block.</summary>
	public IrBuilder Add( IrStmt statement )
	{
		if ( statement is not null ) Current.Add( statement );
		return this;
	}

	/// <summary>
	/// Declare a local with an explicit name, bypassing the CSE table.
	/// <para>
	/// The name is <em>requested</em>, not guaranteed: if it is already taken in this stage the local is
	/// given a fresh one and the returned value names that instead. Emitting two declarations of one
	/// local is a redefinition error from the shader compiler on a generated line, which is exactly the
	/// failure the source map exists to prevent — and this is public IR-builder API, so a custom node
	/// can reach it with any name it likes. Always read the name back off the returned value.
	/// </para>
	/// </summary>
	public IrValue Declare( NodeId origin, string name, ShaderType type, IrValue init )
	{
		if ( string.IsNullOrEmpty( name ) ) return IrValue.Invalid;

		if ( !_names.Add( name ) ) name = NewTempName( name );

		Current.Add( new IrDecl( origin, name, type, init.IsValid ? Intern( init.Expr ) : null ) );

		return Var( type, name );
	}

	/// <summary>Assign to an existing location.</summary>
	public IrBuilder Assign( NodeId origin, IrValue target, IrValue value )
	{
		if ( !target.IsValid || !value.IsValid ) return this;

		Current.Add( new IrAssign( origin, Intern( target.Expr ), Intern( value.Expr ) ) );
		return this;
	}

	/// <summary>Evaluate an expression for its side effects.</summary>
	public IrBuilder ExprStatement( NodeId origin, IrValue value )
	{
		if ( !value.IsValid ) return this;

		Current.Add( new IrExprStmt( origin, Intern( value.Expr ) ) );
		return this;
	}

	/// <summary>Emit a comment. Callers decide whether comments are wanted at all.</summary>
	public IrBuilder Comment( NodeId origin, string text )
	{
		if ( string.IsNullOrWhiteSpace( text ) ) return this;

		Current.Add( new IrComment( origin, text ) );
		return this;
	}

	/// <summary>Return a value from the enclosing function.</summary>
	public IrBuilder Return( NodeId origin, IrValue value )
	{
		Current.Add( new IrReturn( origin, value.IsValid ? Intern( value.Expr ) : null ) );
		return this;
	}

	/// <summary>Kill the fragment unconditionally.</summary>
	public IrBuilder Discard( NodeId origin )
	{
		Current.Add( new IrDiscard( origin ) );
		return this;
	}

	/// <summary>Leave the innermost loop.</summary>
	public IrBuilder Break( NodeId origin )
	{
		Current.Add( new IrBreak( origin ) );
		return this;
	}

	/// <summary>Skip to the next iteration of the innermost loop.</summary>
	public IrBuilder Continue( NodeId origin )
	{
		Current.Add( new IrContinue( origin ) );
		return this;
	}

	// ---- blocks -----------------------------------------------------------

	/// <summary>Append statements to an arbitrary block, e.g. a function body. Dispose to pop.</summary>
	public IDisposable Block( IrBlock block )
	{
		if ( block is null ) return NullScope.Instance;

		_frames.Add( new Frame( block ) );
		return new PopScope( this, null );
	}

	/// <summary>Open a nested scope. Dispose to close it.</summary>
	public IDisposable Scope( NodeId origin )
	{
		var body = new IrBlock();
		_frames.Add( new Frame( body ) );

		return new PopScope( this, _ => Current.Add( new IrScope( origin, body ) ) );
	}

	/// <summary>Open an <c>if</c>. Dispose to close it; <see cref="Else"/> may follow immediately after.</summary>
	public IDisposable If( NodeId origin, IrValue condition )
	{
		if ( !condition.IsValid ) return NullScope.Instance;

		var cond = Intern( condition.Expr );
		var body = new IrBlock();

		_frames.Add( new Frame( body ) );

		return new PopScope( this, parent =>
		{
			parent.Block.Add( new IrIf( origin, cond, body, null ) );
			parent.LastIfIndex = parent.Block.Statements.Count - 1;
		} );
	}

	/// <summary>
	/// Open the <c>else</c> arm of the <c>if</c> that closed most recently in this block. Dispose to
	/// close it. Returns a no-op scope when there is no such <c>if</c>, so a mis-sequenced node cannot
	/// corrupt the block.
	/// </summary>
	public IDisposable Else()
	{
		var frame = _frames[^1];

		if ( frame.LastIfIndex < 0 || frame.LastIfIndex >= frame.Block.Statements.Count ) return NullScope.Instance;
		if ( frame.Block.Statements[frame.LastIfIndex] is not IrIf branch ) return NullScope.Instance;

		var index = frame.LastIfIndex;
		var body = new IrBlock();

		_frames.Add( new Frame( body ) );

		return new PopScope( this, parent =>
		{
			parent.Block.Statements[index] = branch with { Else = body };
			parent.LastIfIndex = -1;
		} );
	}

	/// <summary>Open a counted loop over <c>[0, count)</c>. Dispose to close it.</summary>
	public IDisposable For( NodeId origin, string varName, IrValue count, out IrValue index )
	{
		var name = string.IsNullOrWhiteSpace( varName ) ? NewTempName( "i" ) : Sanitize( varName );
		_names.Add( name );

		index = Var( ShaderType.Int, name );

		if ( !count.IsValid ) return NullScope.Instance;

		var limit = Intern( count.Expr );
		var body = new IrBlock();

		_frames.Add( new Frame( body ) );

		return new PopScope( this, parent => parent.Block.Add( new IrFor( origin, name, limit, body ) ) );
	}

	/// <summary>Open a conditional loop. Dispose to close it.</summary>
	public IDisposable While( NodeId origin, IrValue condition )
	{
		if ( !condition.IsValid ) return NullScope.Instance;

		var cond = Intern( condition.Expr );
		var body = new IrBlock();

		_frames.Add( new Frame( body ) );

		return new PopScope( this, parent => parent.Block.Add( new IrWhile( origin, cond, body ) ) );
	}

	void Pop( Action<Frame> onClose )
	{
		if ( _frames.Count <= 1 ) return;

		_frames.RemoveAt( _frames.Count - 1 );
		onClose?.Invoke( _frames[^1] );
	}

	// ---- helpers ----------------------------------------------------------

	static bool AllPure( IrExpr[] items )
	{
		foreach ( var item in items )
		{
			if ( item is null || !item.Pure ) return false;
		}

		return true;
	}

	IrExpr[] Unwrap( IrValue[] values )
	{
		if ( values is null ) return Array.Empty<IrExpr>();

		var result = new IrExpr[values.Length];

		for ( int i = 0; i < values.Length; i++ )
		{
			if ( !values[i].IsValid ) return null;

			result[i] = Intern( values[i].Expr );
		}

		return result;
	}

	/// <summary>Turn a free-form hint into a legal HLSL identifier fragment.</summary>
	public static string Sanitize( string hint )
	{
		if ( string.IsNullOrWhiteSpace( hint ) ) return "t";

		var sb = new StringBuilder( hint.Length );

		foreach ( var c in hint )
		{
			if ( char.IsLetterOrDigit( c ) ) sb.Append( c );
			else if ( c is '_' ) sb.Append( '_' );
		}

		if ( sb.Length == 0 ) return "t";
		if ( char.IsDigit( sb[0] ) ) sb.Insert( 0, '_' );

		return sb.ToString();
	}

	/// <summary>
	/// Structural hash of one expression node, folding in the already-computed hashes of its interned
	/// children so hashing stays linear in the number of children rather than in subtree size.
	/// </summary>
	public static int StructuralHash( IrExpr expr )
	{
		if ( expr is null ) return 0;

		var hash = expr switch
		{
			IrConst c => HashCode.Combine( 1, c.Value ),
			IrVar v => HashCode.Combine( 2, v.Name ),
			IrGlobalRef g => HashCode.Combine( 3, g.Decl?.Name ),
			IrBuiltinRef b => HashCode.Combine( 4, (int)b.Id ),
			IrCall call => HashCode.Combine( 5, (int)call.Id ),
			IrHelperCall helper => HashCode.Combine( 6, helper.Fn?.Name ),
			IrBinary bin => HashCode.Combine( 7, (int)bin.Op ),
			IrUnary un => HashCode.Combine( 8, (int)un.Op ),
			IrSwizzle sw => HashCode.Combine( 9, sw.Mask ),
			IrConstruct => 10,
			IrCast cast => HashCode.Combine( 11, (int)cast.Kind, cast.Fill ),
			IrSelect => 12,
			IrIndex => 13,
			IrMember member => HashCode.Combine( 14, member.Field ),
			_ => 15
		};

		hash = HashCode.Combine( hash, expr.Type );

		foreach ( var child in expr.Children )
		{
			hash = HashCode.Combine( hash, child is null ? 0 : child.Hash != 0 ? child.Hash : child.GetHashCode() );
		}

		return hash == 0 ? 1 : hash;
	}

	sealed class Frame
	{
		public Frame( IrBlock block ) => Block = block;

		public IrBlock Block { get; }
		public Dictionary<IrExpr, IrValue> Bindings { get; } = new( StructuralComparer.Instance );
		public int LastIfIndex { get; set; } = -1;
	}

	sealed class PopScope : IDisposable
	{
		readonly IrBuilder _builder;
		readonly Action<Frame> _onClose;
		bool _closed;

		public PopScope( IrBuilder builder, Action<Frame> onClose )
		{
			_builder = builder;
			_onClose = onClose;
		}

		public void Dispose()
		{
			if ( _closed ) return;

			_closed = true;
			_builder.Pop( _onClose );
		}
	}

	sealed class NullScope : IDisposable
	{
		public static readonly NullScope Instance = new();

		public void Dispose() { }
	}

	/// <summary>
	/// Structural comparer that consults the cached <see cref="IrExpr.Hash"/> before falling back to the
	/// record's own deep equality, so unequal expressions almost always reject on the hash alone.
	/// </summary>
	sealed class StructuralComparer : IEqualityComparer<IrExpr>
	{
		public static readonly StructuralComparer Instance = new();

		public bool Equals( IrExpr a, IrExpr b )
		{
			if ( ReferenceEquals( a, b ) ) return true;
			if ( a is null || b is null ) return false;
			if ( a.Hash != 0 && b.Hash != 0 && a.Hash != b.Hash ) return false;

			return a.Equals( b );
		}

		public int GetHashCode( IrExpr expr ) =>
			expr is null ? 0 : expr.Hash != 0 ? expr.Hash : StructuralHash( expr );
	}
}