using Editor.Prism.Core;

namespace Editor.Prism.Compiler.Ir;

/// <summary>
/// A literal value, stored as up to four components. Integers and booleans use the same storage —
/// the owning <see cref="ShaderType"/> says how to spell them. Nothing in a shader graph needs more
/// range than a double provides.
/// </summary>
public readonly record struct ConstValue( double X, double Y, double Z, double W )
{
	/// <summary>All components zero.</summary>
	public static readonly ConstValue Zero = default;

	/// <summary>All components one.</summary>
	public static readonly ConstValue One = new( 1, 1, 1, 1 );

	/// <summary>A scalar.</summary>
	public static ConstValue From( float value ) => new( value, value, value, value );

	/// <summary>An integer.</summary>
	public static ConstValue From( int value ) => new( value, value, value, value );

	/// <summary>A boolean, stored as 0 or 1.</summary>
	public static ConstValue From( bool value ) => value ? One : Zero;

	/// <summary>A two-component vector.</summary>
	public static ConstValue From( Vector2 value ) => new( value.x, value.y, 0, 0 );

	/// <summary>A three-component vector.</summary>
	public static ConstValue From( Vector3 value ) => new( value.x, value.y, value.z, 0 );

	/// <summary>A four-component vector.</summary>
	public static ConstValue From( Vector4 value ) => new( value.x, value.y, value.z, value.w );

	/// <summary>A colour.</summary>
	public static ConstValue From( Color value ) => new( value.r, value.g, value.b, value.a );

	/// <summary>Component by index. Out-of-range indices return the last component.</summary>
	public double this[int index] => index switch
	{
		0 => X,
		1 => Y,
		2 => Z,
		_ => W
	};

	/// <summary>True when the first <paramref name="components"/> components are all zero.</summary>
	public bool IsZero( int components ) => AllEqual( 0, components );

	/// <summary>True when the first <paramref name="components"/> components are all one.</summary>
	public bool IsOne( int components ) => AllEqual( 1, components );

	/// <summary>True when the first <paramref name="components"/> components all equal <paramref name="value"/>.</summary>
	public bool AllEqual( double value, int components )
	{
		for ( int i = 0; i < Math.Clamp( components, 1, 4 ); i++ )
		{
			if ( this[i] != value ) return false;
		}

		return true;
	}

	/// <summary>The value as a Vector4, for interop with engine types.</summary>
	public Vector4 ToVector4() => new( (float)X, (float)Y, (float)Z, (float)W );

	/// <inheritdoc/>
	public override string ToString() => $"({X}, {Y}, {Z}, {W})";
}

/// <summary>
/// A typed expression node.
/// <para>
/// The IR is an expression tree, not a full CFG — we do not need one, and a tree keeps hash-consed
/// CSE, constant folding and printing simple. <see cref="Hash"/> is filled by the builder and is
/// deliberately excluded from equality, as is <see cref="Pure"/>, which is a function of the node's
/// operation rather than independent state.
/// </para>
/// </summary>
public abstract record IrExpr( ShaderType Type )
{
	/// <summary>Structural hash, filled by the hash-consing builder. Not part of equality.</summary>
	public int Hash { get; init; }

	/// <summary>
	/// False for expressions with observable side effects or with results that depend on
	/// neighbouring lanes: sampling with implicit derivatives, atomics, wave intrinsics.
	/// Only pure expressions participate in CSE and hoisting.
	/// </summary>
	public bool Pure { get; init; } = true;

	/// <summary>Sub-expressions, in evaluation order. Empty for leaves.</summary>
	public virtual IReadOnlyList<IrExpr> Children => Array.Empty<IrExpr>();

	/// <summary>Compares the node kind and result type; derived records add their own operands.</summary>
	public virtual bool Equals( IrExpr other ) =>
		other is not null && EqualityContract == other.EqualityContract && Type == other.Type;

	/// <inheritdoc/>
	public override int GetHashCode() => HashCode.Combine( EqualityContract, Type );
}

/// <summary>A literal.</summary>
public sealed record IrConst( ShaderType Type, ConstValue Value ) : IrExpr( Type );

/// <summary>A reference to a named local, usually a temp the emitter bound earlier.</summary>
public sealed record IrVar( ShaderType Type, string Name ) : IrExpr( Type );

/// <summary>A reference to a module-level uniform, texture, sampler or buffer.</summary>
public sealed record IrGlobalRef( ShaderType Type, GlobalDecl Decl ) : IrExpr( Type );

/// <summary>
/// A reference to an environment-provided value. Lowered per stage by the backend, so nodes never
/// have to know that world position is a different expression in the vertex and pixel stages.
/// </summary>
public sealed record IrBuiltinRef( ShaderType Type, Builtin Id ) : IrExpr( Type );

/// <summary>A call to a canonical intrinsic, spelled per backend.</summary>
public sealed record IrCall( ShaderType Type, Intrinsic Id, IrExpr[] Args ) : IrExpr( Type )
{
	/// <inheritdoc/>
	public override IReadOnlyList<IrExpr> Children => Args ?? Array.Empty<IrExpr>();

	/// <summary>Structural comparison, including the argument list.</summary>
	public bool Equals( IrCall other ) =>
		base.Equals( other ) && Id == other.Id && IrExprUtil.SequenceEqual( Args, other.Args );

	/// <inheritdoc/>
	public override int GetHashCode() =>
		HashCode.Combine( base.GetHashCode(), (int)Id, IrExprUtil.SequenceHash( Args ) );
}

/// <summary>A call to a shared helper function emitted once per module.</summary>
public sealed record IrHelperCall( ShaderType Type, HelperFunction Fn, IrExpr[] Args ) : IrExpr( Type )
{
	/// <inheritdoc/>
	public override IReadOnlyList<IrExpr> Children => Args ?? Array.Empty<IrExpr>();

	/// <summary>
	/// Structural comparison. Helpers are deduplicated by name and a same-name/different-body
	/// collision is a hard compile error, so the name is the identity — comparing by reference here
	/// would silently defeat CSE between two nodes that declare the same helper independently.
	/// </summary>
	public bool Equals( IrHelperCall other ) =>
		base.Equals( other ) && Fn?.Name == other.Fn?.Name && IrExprUtil.SequenceEqual( Args, other.Args );

	/// <inheritdoc/>
	public override int GetHashCode() =>
		HashCode.Combine( base.GetHashCode(), Fn?.Name, IrExprUtil.SequenceHash( Args ) );
}

/// <summary>A binary operation.</summary>
public sealed record IrBinary( ShaderType Type, BinaryOp Op, IrExpr L, IrExpr R ) : IrExpr( Type )
{
	/// <inheritdoc/>
	public override IReadOnlyList<IrExpr> Children => [L, R];
}

/// <summary>A unary operation.</summary>
public sealed record IrUnary( ShaderType Type, UnaryOp Op, IrExpr V ) : IrExpr( Type )
{
	/// <inheritdoc/>
	public override IReadOnlyList<IrExpr> Children => [V];
}

/// <summary>A swizzle, e.g. <c>.xy</c>, <c>.rgb</c> or <c>.xxxw</c>.</summary>
public sealed record IrSwizzle( ShaderType Type, IrExpr V, string Mask ) : IrExpr( Type )
{
	/// <inheritdoc/>
	public override IReadOnlyList<IrExpr> Children => [V];
}

/// <summary>Construction of a vector, matrix or struct from parts.</summary>
public sealed record IrConstruct( ShaderType Type, IrExpr[] Parts ) : IrExpr( Type )
{
	/// <inheritdoc/>
	public override IReadOnlyList<IrExpr> Children => Parts ?? Array.Empty<IrExpr>();

	/// <summary>Structural comparison, including the part list.</summary>
	public bool Equals( IrConstruct other ) =>
		base.Equals( other ) && IrExprUtil.SequenceEqual( Parts, other.Parts );

	/// <inheritdoc/>
	public override int GetHashCode() => HashCode.Combine( base.GetHashCode(), IrExprUtil.SequenceHash( Parts ) );
}

/// <summary>A conversion between types.</summary>
public sealed record IrCast( ShaderType Type, IrExpr V, CastKind Kind ) : IrExpr( Type )
{
	/// <summary>Value written into components invented by a <see cref="CastKind.Pad"/>.</summary>
	public float Fill { get; init; }

	/// <inheritdoc/>
	public override IReadOnlyList<IrExpr> Children => [V];
}

/// <summary>
/// A component-wise select. Always lowered to <c>select( c, a, b )</c>, never to <c>?:</c> —
/// the ternary does not work component-wise on vectors in HLSL 2021 or Slang.
/// </summary>
public sealed record IrSelect( ShaderType Type, IrExpr C, IrExpr A, IrExpr B ) : IrExpr( Type )
{
	/// <inheritdoc/>
	public override IReadOnlyList<IrExpr> Children => [C, A, B];
}

/// <summary>An indexing expression, e.g. a buffer load or a matrix row.</summary>
public sealed record IrIndex( ShaderType Type, IrExpr V, IrExpr I ) : IrExpr( Type )
{
	/// <inheritdoc/>
	public override IReadOnlyList<IrExpr> Children => [V, I];
}

/// <summary>A struct member access.</summary>
public sealed record IrMember( ShaderType Type, IrExpr V, string Field ) : IrExpr( Type )
{
	/// <inheritdoc/>
	public override IReadOnlyList<IrExpr> Children => [V];
}

/// <summary>Shared structural comparison helpers for IR nodes that carry arrays.</summary>
public static class IrExprUtil
{
	/// <summary>Element-wise structural comparison of two argument arrays.</summary>
	public static bool SequenceEqual( IrExpr[] a, IrExpr[] b )
	{
		if ( ReferenceEquals( a, b ) ) return true;
		if ( a is null || b is null ) return false;
		if ( a.Length != b.Length ) return false;

		for ( int i = 0; i < a.Length; i++ )
		{
			if ( !Equals( a[i], b[i] ) ) return false;
		}

		return true;
	}

	/// <summary>Order-sensitive structural hash of an argument array.</summary>
	public static int SequenceHash( IrExpr[] items )
	{
		if ( items is null ) return 0;

		var hash = items.Length;

		foreach ( var item in items )
		{
			hash = HashCode.Combine( hash, item?.GetHashCode() ?? 0 );
		}

		return hash;
	}

	/// <summary>Walk an expression tree depth-first, yielding every node including the root.</summary>
	public static IEnumerable<IrExpr> Walk( IrExpr root )
	{
		if ( root is null ) yield break;

		yield return root;

		foreach ( var child in root.Children )
		{
			foreach ( var descendant in Walk( child ) )
			{
				yield return descendant;
			}
		}
	}
}