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

namespace Editor.Prism.Compiler.Backends;

/// <summary>
/// One writable field of the engine's <c>Material</c> struct, as a surface graph sees it.
/// </summary>
/// <param name="Name">The canonical Prism output name, e.g. <c>Albedo</c>.</param>
/// <param name="Field">The HLSL field on <c>Material</c>.</param>
/// <param name="Type">The type the field expects.</param>
/// <param name="Saturate">True when the lit epilogue clamps the field to 0..1 before shading.</param>
/// <param name="Description">What the field means, for tooltips and the node library.</param>
public readonly record struct MaterialField(
	string Name, string Field, ShaderType Type, bool Saturate, string Description )
{
	/// <summary>The HLSL expression that writes this field of the generated material local.</summary>
	public string Reference => $"{SboxShaderTemplates.MaterialLocal}.{Field}";

	/// <inheritdoc/>
	public override string ToString() => $"{Type.Hlsl} {Reference}";
}

/// <summary>
/// Everything about binding a Prism graph to the s&amp;box material model: which surface outputs map
/// onto which <c>Material</c> fields, how the lit and unlit tails are written, and how a
/// <see cref="GlobalDecl"/> becomes a declaration the material editor understands.
/// <para>
/// The metadata grammar this emits is <c>system.fxc</c>'s, and it is unforgiving: <c>UiType( X )</c>
/// is unquoted while <c>UiGroup( "X" )</c> is quoted, <c>DefaultN</c> and <c>RangeN</c> suffixes must
/// match the channel count or the material editor silently ignores them, and <c>RangeN</c> packs all
/// minima before all maxima rather than min/max pairs.
/// </para>
/// </summary>
public static class SboxMaterialBinding
{
	// ---- the material model ----------------------------------------------

	/// <summary>
	/// The fields a surface graph writes, in the order the generated shader assigns them.
	/// Names match the output ports of the surface output node.
	/// </summary>
	public static readonly IReadOnlyList<MaterialField> Fields =
	[
		new( "Albedo", "Albedo", ShaderType.Float3, false, "Base colour, linear." ),
		new( "Opacity", "Opacity", ShaderType.Float, true, "Coverage. Drives alpha test and translucency." ),
		new( "Normal", "Normal", ShaderType.Float3, false, "Tangent-space normal; converted to world space before shading." ),
		new( "Roughness", "Roughness", ShaderType.Float, true, "Microfacet roughness." ),
		new( "Metalness", "Metalness", ShaderType.Float, true, "Dielectric to conductor blend." ),
		new( "AmbientOcclusion", "AmbientOcclusion", ShaderType.Float, true, "Baked occlusion multiplier." ),
		new( "Emission", "Emission", ShaderType.Float3, false, "Additive emissive colour." ),
		new( "Transmission", "Transmission", ShaderType.Float3, false, "Light transmitted through the surface." ),
		new( "TintMask", "TintMask", ShaderType.Float, true, "Where the per-instance tint applies." )
	];

	/// <summary>
	/// Geometry the material carries for tools visualisation and lighting. Written by the generated
	/// epilogue rather than by the graph.
	/// </summary>
	public static readonly IReadOnlyList<string> GeometryFields =
	[
		"WorldPosition", "WorldPositionWithOffset", "ScreenPosition",
		"WorldTangentU", "WorldTangentV", "LightmapUV", "TextureCoords"
	];

	/// <summary>The vertex-stage location a position offset is added to, before <c>FinalizeVertex</c>.</summary>
	public const string VertexPositionTarget = "i.vPositionWs.xyz";

	/// <summary>Look up a writable material field by its canonical Prism name.</summary>
	public static bool TryGetField( string name, out MaterialField field )
	{
		foreach ( var candidate in Fields )
		{
			if ( !string.Equals( candidate.Name, name, StringComparison.OrdinalIgnoreCase ) ) continue;

			field = candidate;
			return true;
		}

		field = default;
		return false;
	}

	/// <summary>True when this module drives the engine's standard shading model.</summary>
	public static bool UsesMaterial( IrModule module )
	{
		if ( module?.Meta is null ) return false;

		return module.Meta.Domain is ShaderDomain.Surface or ShaderDomain.PostProcess &&
			   module.Meta.ShadingModel != ShadingModel.Custom;
	}

	// ---- pixel prologue / epilogue ---------------------------------------

	/// <summary>
	/// Write the material declaration the graph's assignments target.
	/// <para>
	/// <c>Material::Init( i )</c> pulls world position, normal, tangents, lightmap UVs and texture
	/// coordinates straight off the pixel input, then the stock defaults are applied so a graph that
	/// wires nothing still renders the same white dielectric the built-in editor produces.
	/// </para>
	/// </summary>
	public static void WritePixelPrologue( HlslSourceBuilder builder, IrModule module )
	{
		if ( builder is null || !UsesMaterial( module ) ) return;

		builder.WriteBlock( SboxShaderTemplates.MaterialInit );
		builder.Blank();
	}

	/// <summary>
	/// Write the tail of the pixel entry point when the module's own body did not return.
	/// <para>
	/// The lit path finishes through <c>ShadingModelStandard::Shade</c>, which is what buys decals,
	/// alpha-to-coverage, clustered lighting, shadows, indirect lighting, SSAO, SSR, fog, wireframe,
	/// tools visualisation and the depth/normals G-buffer output for free. The unlit and post-process
	/// paths return a colour directly.
	/// </para>
	/// </summary>
	public static void WritePixelEpilogue( HlslSourceBuilder builder, IrModule module, HlslEmitter emitter )
	{
		if ( builder is null || module?.Meta is null ) return;

		var meta = module.Meta;

		if ( meta.ShadingModel == ShadingModel.Custom )
		{
			emitter?.Report( DiagnosticSeverity.Error, DiagnosticCode.NoOutput,
				"A custom shading model must return a colour from its pixel program.",
				"Prism emits no shading tail for ShadingModel.Custom, so the graph's output node has to end the function with a return." );
			builder.Write( "return float4( 0.0f, 0.0f, 0.0f, 1.0f );" );
			return;
		}

		if ( meta.Domain == ShaderDomain.Surface && meta.ShadingModel == ShadingModel.Lit )
		{
			builder.WriteBlock( SboxShaderTemplates.MaterialShade );
			return;
		}

		builder.WriteBlock( SboxShaderTemplates.MaterialUnlitReturn );
	}

	/// <summary>True when the last statement of a block already returns, so no tail should be added.</summary>
	public static bool EndsWithReturn( IrBlock block )
	{
		if ( block is null || block.Statements.Count == 0 ) return false;

		return block.Statements[^1] is IrReturn;
	}

	// ---- declarations -----------------------------------------------------

	/// <summary>
	/// Write every module-level declaration a stage reads, in declaration order.
	/// <para>
	/// Preview-only globals — the uniforms literals become so a slider drag needs no recompile — are
	/// stripped from a final build.
	/// </para>
	/// </summary>
	public static void WriteGlobals( HlslSourceBuilder builder, HlslEmitter emitter, ShaderStage stage )
	{
		if ( builder is null || emitter?.Module is null ) return;

		var module = emitter.Module;
		var wrote = false;
		var final = emitter.Options.Mode == CompileMode.Final;
		var placement = GlobalPlacement.For( module );

		foreach ( var global in module.Globals )
		{
			if ( global is null ) continue;
			if ( global.PreviewOnly && final ) continue;
			if ( !placement.ShouldDeclare( global, stage ) ) continue;

			foreach ( var line in Declare( global ) ) builder.Write( line );

			wrote = true;
		}

		if ( wrote ) builder.Blank();

		if ( WriteObjectTransformHelpers( builder, module, stage ) ) builder.Blank();
	}

	/// <summary>
	/// Write the object-transform helpers when this stage's code asks for one of the object matrices.
	/// <para>
	/// They live here rather than in the block writer because every path that emits declarations for a
	/// stage goes through <see cref="WriteGlobals"/> — the <c>.shader</c> writer and the standalone
	/// emitter the text editor probes with — and a helper that only one of them emitted would compile
	/// in one place and not the other.
	/// </para>
	/// </summary>
	/// <returns>True when anything was written.</returns>
	public static bool WriteObjectTransformHelpers( HlslSourceBuilder builder, IrModule module, ShaderStage stage )
	{
		if ( builder is null || module?.Meta is null ) return false;

		var body = HlslIntrinsics.ObjectTransformHelpers( stage, module.Meta.Domain );

		if ( string.IsNullOrEmpty( body ) ) return false;
		if ( !NeedsObjectTransform( module, stage ) ) return false;

		builder.WriteBlock( body );

		return true;
	}

	/// <summary>True when a stage's code reads one of the object-space transform builtins.</summary>
	public static bool NeedsObjectTransform( IrModule module, ShaderStage stage )
	{
		foreach ( var expr in IrWalk.Expressions( module, stage ) )
		{
			foreach ( var node in IrExprUtil.Walk( expr ) )
			{
				if ( node is not IrBuiltinRef reference ) continue;

				if ( reference.Id is Builtin.ObjectToWorld or Builtin.WorldToObject or Builtin.ObjectScale )
					return true;
			}
		}

		return false;
	}

	/// <summary>
	/// Turn one declaration into the lines a <c>.shader</c> needs for it. Textures produce two: the
	/// artist-facing input slot and the GPU texture that packs it.
	/// </summary>
	public static IReadOnlyList<string> Declare( GlobalDecl decl )
	{
		if ( decl is null ) return Array.Empty<string>();

		// One rename, here, before any of the six declaration shapes below reads the name — and the same
		// rename the expression printer applies to every reference. The engine's block parser is ASCII
		// only and rejects the whole file, with no line number, over a single accented letter.
		var safe = SboxShaderTemplates.SafeIdentifier( decl.Name );

		if ( !string.Equals( safe, decl.Name, StringComparison.Ordinal ) ) decl = decl with { Name = safe };

		return decl.Kind switch
		{
			GlobalKind.Texture => DeclareTexture( decl ),
			GlobalKind.Sampler => [DeclareSampler( decl )],
			GlobalKind.Constant => [DeclareConstant( decl )],
			GlobalKind.Buffer or GlobalKind.RwBuffer or GlobalKind.RwTexture => [DeclareResource( decl )],
			_ => [DeclareUniform( decl )]
		};
	}

	static string DeclareUniform( GlobalDecl decl )
	{
		var annotations = new List<string>();
		var components = Math.Clamp( decl.Type.Components, 1, 4 );
		var ui = decl.Ui;

		if ( !string.IsNullOrEmpty( decl.AttributeName ) )
		{
			// Attribute-bound parameters are pushed from C# at runtime. The stock compiler drops UI
			// metadata entirely for these, because the material editor never shows them.
			annotations.Add( $"Attribute( \"{SboxShaderTemplates.QuoteSafe( decl.AttributeName )}\" )" );

			if ( decl.Default.HasValue ) annotations.Add( DefaultAnnotation( decl.Default.Value, components, decl.Type ) );
		}
		else
		{
			var control = UiTypeToken( decl, ui );

			if ( !string.IsNullOrEmpty( control ) ) annotations.Add( $"UiType( {control} )" );

			if ( decl.Default.HasValue ) annotations.Add( DefaultAnnotation( decl.Default.Value, components, decl.Type ) );

			if ( ui is not null && ui.Control == UiControl.Slider && decl.Type.IsNumeric && !decl.Type.IsBoolean )
			{
				annotations.Add( RangeAnnotation( ui.Min, ui.Max, components ) );

				if ( ui.Step > 0 ) annotations.Add( $"UiStep( {Value( ui.Step )} )" );
			}

			var group = FormatUiGroup( ui );

			if ( !string.IsNullOrEmpty( group ) ) annotations.Add( $"UiGroup( \"{group}\" )" );
		}

		return $"{HlslIntrinsics.DeclarationType( decl )} {decl.Name}{Annotations( annotations )};";
	}

	static IReadOnlyList<string> DeclareTexture( GlobalDecl decl )
	{
		var lines = new List<string>( 2 );
		var input = InputTextureName( decl.Name );
		var colorSpace = decl.Srgb ? "Srgb" : "Linear";
		var group = FormatUiGroup( decl.Ui );
		var creator = decl.Type.Object switch
		{
			ObjectKind.Texture3D => "CreateInputTexture3D",
			ObjectKind.TextureCube or ObjectKind.TextureCubeArray => "CreateInputTextureCube",
			_ => "CreateInputTexture2D"
		};

		string fallback;

		if ( !string.IsNullOrEmpty( decl.DefaultAsset ) )
		{
			fallback = $"DefaultFile( \"{SboxShaderTemplates.QuoteSafe( decl.DefaultAsset )}\" )";
		}
		else
		{
			var channels = TextureChannels( decl );
			fallback = DefaultAnnotation( decl.Default ?? ConstValue.One, channels, ShaderType.Vec( ScalarKind.Float, channels ) );
		}

		if ( !string.IsNullOrEmpty( decl.AttributeName ) )
		{
			// A texture bound to a render attribute has no material-editor slot at all.
			lines.Add( $"{HlslIntrinsics.DeclarationType( decl )} {decl.Name} < Attribute( \"{SboxShaderTemplates.QuoteSafe( decl.AttributeName )}\" ); SrgbRead( {Boolean( decl.Srgb )} ); >;" );
			return lines;
		}

		lines.Add( $"{creator}( {input}, {colorSpace}, 8, \"\", \"\", \"{group}\", {fallback} );" );
		lines.Add( $"{HlslIntrinsics.DeclarationType( decl )} {decl.Name} < Channel( RGBA, Box( {input} ), {colorSpace} ); OutputFormat( BC7 ); SrgbRead( {Boolean( decl.Srgb )} ); >;" );

		return lines;
	}

	static string DeclareSampler( GlobalDecl decl )
	{
		// Sampler state is expressed as raw annotations because the VFX grammar has far more knobs
		// (Filter, AddressU/V/W, MipBias, MaxAniso, BorderColor, ComparisonFunc) than UiHints models.
		var annotations = decl.Ui?.Options is { Count: > 0 } options
			? options.Where( x => !string.IsNullOrWhiteSpace( x ) ).ToList()
			: new List<string>();

		if ( decl.Type.Object != ObjectKind.SamplerComparisonState )
		{
			return $"SamplerState {decl.Name}{Annotations( annotations )};";
		}

		// A comparison sampler is the hardware depth-compare unit, and it only works with a COMPARISON
		// filter and a comparison function. Both are mandatory rather than stylistic: a comparison
		// sampler declared with an ordinary filter is a driver-level error on some targets and silently
		// returns the raw depth on others. The defaults mirror the engine's own ShadowDepthPCFSampler.
		if ( !HasAnnotation( annotations, "Filter" ) ) annotations.Insert( 0, ComparisonFilter );
		if ( !HasAnnotation( annotations, "ComparisonFunc" ) ) annotations.Add( ComparisonFunction );

		return $"SamplerComparisonState {decl.Name}{Annotations( annotations )};";
	}

	/// <summary>The filter a comparison sampler defaults to: hardware PCF, bilinear between taps.</summary>
	public const string ComparisonFilter = "Filter( COMPARISON_MIN_MAG_MIP_LINEAR )";

	/// <summary>
	/// The comparison a comparison sampler defaults to. <c>LESS_EQUAL</c> is the ordinary
	/// "is the sample nearer than the stored depth" test; the engine's shadow sampler uses
	/// <c>GREATER_EQUAL</c> because its depth buffer is reversed.
	/// </summary>
	public const string ComparisonFunction = "ComparisonFunc( LESS_EQUAL )";

	/// <summary>True when an annotation list already sets a named piece of sampler state.</summary>
	static bool HasAnnotation( IReadOnlyList<string> annotations, string name )
	{
		foreach ( var annotation in annotations )
		{
			if ( annotation is null ) continue;
			if ( annotation.TrimStart().StartsWith( name, StringComparison.OrdinalIgnoreCase ) ) return true;
		}

		return false;
	}

	static string DeclareConstant( GlobalDecl decl )
	{
		var value = decl.Default.HasValue
			? HlslIntrinsics.Literal( decl.Type, decl.Default.Value )
			: HlslIntrinsics.Fallback( decl.Type );

		return $"static const {decl.Type.Hlsl} {decl.Name} = {value};";
	}

	static string DeclareResource( GlobalDecl decl )
	{
		var annotations = new List<string>();

		if ( !string.IsNullOrEmpty( decl.AttributeName ) )
		{
			annotations.Add( $"Attribute( \"{SboxShaderTemplates.QuoteSafe( decl.AttributeName )}\" )" );
		}

		return $"{HlslIntrinsics.DeclarationType( decl )} {decl.Name}{Annotations( annotations )};";
	}

	// ---- annotation grammar ----------------------------------------------

	/// <summary>
	/// The <c>DefaultN</c> annotation for a value, with the suffix matching the channel count.
	/// <para>
	/// The built-in generator always writes <c>Default4</c> regardless of the type; the material editor
	/// then silently ignores the metadata. Prism does not reproduce that.
	/// </para>
	/// </summary>
	public static string DefaultAnnotation( ConstValue value, int components, ShaderType type )
	{
		components = Math.Clamp( components, 1, 4 );

		if ( type.IsBoolean ) return $"Default( {( value.X != 0 ? "1" : "0" )} )";

		var parts = new string[components];

		for ( int i = 0; i < components; i++ ) parts[i] = Value( value[i] );

		return components == 1
			? $"Default( {parts[0]} )"
			: $"Default{components}( {string.Join( ", ", parts )} )";
	}

	/// <summary>
	/// The <c>RangeN</c> annotation. <c>RangeN</c> packs <em>all minima and then all maxima</em>, not
	/// min/max pairs — getting that backwards produces a slider with nonsense bounds.
	/// </summary>
	public static string RangeAnnotation( float min, float max, int components )
	{
		components = Math.Clamp( components, 1, 4 );

		if ( components == 1 ) return $"Range( {Value( min )}, {Value( max )} )";

		var parts = new List<string>( components * 2 );

		for ( int i = 0; i < components; i++ ) parts.Add( Value( min ) );
		for ( int i = 0; i < components; i++ ) parts.Add( Value( max ) );

		return $"Range{components}( {string.Join( ", ", parts )} )";
	}

	/// <summary>
	/// The <c>UiGroup</c> payload: <c>"Primary,PrimaryPriority/ItemPriority"</c>. A group string that
	/// already contains a comma or a slash is passed through untouched, so an authored group survives.
	/// </summary>
	public static string FormatUiGroup( UiHints ui )
	{
		if ( ui is null || string.IsNullOrWhiteSpace( ui.Group ) ) return string.Empty;

		var group = SboxShaderTemplates.QuoteSafe( ui.Group );

		if ( group.Contains( ',' ) || group.Contains( '/' ) ) return group;

		return $"{group},10/{ui.Order}";
	}

	/// <summary>
	/// The <c>UiType</c> token, unquoted as the grammar requires, or empty when the type's default
	/// control is right. Only <c>Slider</c>, <c>Color</c> and <c>CheckBox</c> exist.
	/// </summary>
	public static string UiTypeToken( GlobalDecl decl, UiHints ui )
	{
		if ( decl is not null && decl.Type.IsBoolean ) return "CheckBox";

		return ui?.Control switch
		{
			UiControl.Slider => "Slider",
			UiControl.Color => "Color",
			UiControl.Toggle => "CheckBox",
			_ => string.Empty
		};
	}

	/// <summary>
	/// How many channels a texture's default value describes.
	/// <para>
	/// <see cref="GlobalDecl"/> carries no explicit channel count for a texture, so this reads the
	/// signals that do exist: an explicit colour control or sRGB encoding means three channels, a
	/// value identical across all four components is a single-channel data map, and an unused fourth
	/// component means three. Everything else is RGBA.
	/// </para>
	/// </summary>
	public static int TextureChannels( GlobalDecl decl )
	{
		if ( decl is null ) return 4;
		if ( decl.Ui?.Control == UiControl.Color ) return 3;
		if ( decl.Srgb ) return 3;

		var value = decl.Default ?? ConstValue.One;

		if ( value.AllEqual( value.X, 4 ) ) return 1;
		if ( value.W is 0 or 1 ) return 3;

		return 4;
	}

	static string InputTextureName( string name )
	{
		if ( string.IsNullOrWhiteSpace( name ) ) return "TextureInput";

		var trimmed = name;

		if ( trimmed.StartsWith( "g_t", StringComparison.Ordinal ) ) trimmed = trimmed[3..];
		else if ( trimmed.StartsWith( "g_", StringComparison.Ordinal ) ) trimmed = trimmed[2..];

		if ( trimmed.Length == 0 ) return "TextureInput";
		if ( trimmed.StartsWith( "Texture", StringComparison.Ordinal ) ) return trimmed;

		return "Texture" + char.ToUpperInvariant( trimmed[0] ) + trimmed[1..];
	}

	static string Annotations( IReadOnlyList<string> annotations )
	{
		if ( annotations is null || annotations.Count == 0 ) return string.Empty;

		var builder = new StringBuilder( " < " );

		foreach ( var annotation in annotations )
		{
			builder.Append( annotation ).Append( "; " );
		}

		builder.Append( '>' );

		return builder.ToString();
	}

	/// <summary>
	/// Spell a metadata value. The annotation grammar takes plain decimals, and the value has to be
	/// rounded through <c>float</c> first — a literal authored as <c>0.62f</c> widened to a double
	/// prints as <c>0.6200000047683716</c>, which is technically correct and completely unreadable.
	/// </summary>
	static string Value( double value )
	{
		var text = HlslIntrinsics.Number( value, ScalarKind.Float );

		return text.EndsWith( "f", StringComparison.Ordinal ) ? text[..^1] : text;
	}

	static string Boolean( bool value ) => value ? "true" : "false";
}

/// <summary>
/// Traversal of an <see cref="IrModule"/> restricted to what one stage actually emits.
/// <para>
/// A stage's block contains its entry point plus every stage-agnostic function, so "does this stage
/// reference X" is a question about that exact set and not about the module as a whole. Getting the
/// set wrong is what makes a texture appear in both the VS and the PS block.
/// </para>
/// </summary>
public static class IrWalk
{
	/// <summary>Every function whose body is written into a stage's block.</summary>
	public static IEnumerable<IrFunction> Functions( IrModule module, ShaderStage stage )
	{
		if ( module is null ) yield break;

		foreach ( var function in module.Functions )
		{
			if ( function is null ) continue;

			if ( function.IsEntryPoint )
			{
				if ( function.Stage == stage ) yield return function;
				continue;
			}

			// A function with no stage of its own is emitted into every block, exactly as
			// HlslEmitter.WriteFunctions does, so whatever it reads is read by every stage.
			if ( function.Stage != ShaderStage.None && function.Stage != stage ) continue;

			yield return function;
		}
	}

	/// <summary>Every statement in a block, including those nested inside branches, loops and scopes.</summary>
	public static IEnumerable<IrStmt> Statements( IrBlock block )
	{
		if ( block is null ) yield break;

		foreach ( var statement in block.Statements )
		{
			yield return statement;

			foreach ( var child in Blocks( statement ) )
			{
				foreach ( var nested in Statements( child ) ) yield return nested;
			}
		}
	}

	/// <summary>The sub-blocks a statement owns.</summary>
	public static IEnumerable<IrBlock> Blocks( IrStmt statement )
	{
		switch ( statement )
		{
			case IrIf branch:
				if ( branch.Then is not null ) yield return branch.Then;
				if ( branch.Else is not null ) yield return branch.Else;
				break;

			case IrPreprocessorIf guard:
				if ( guard.Then is not null ) yield return guard.Then;
				if ( guard.Else is not null ) yield return guard.Else;
				break;

			case IrFor loop when loop.Body is not null:
				yield return loop.Body;
				break;

			case IrWhile loop when loop.Body is not null:
				yield return loop.Body;
				break;

			case IrScope scope when scope.Body is not null:
				yield return scope.Body;
				break;
		}
	}

	/// <summary>Every root expression a statement evaluates.</summary>
	public static IEnumerable<IrExpr> Expressions( IrStmt statement )
	{
		switch ( statement )
		{
			case IrDecl decl when decl.Init is not null:
				yield return decl.Init;
				break;

			case IrAssign assign:
				if ( assign.Target is not null ) yield return assign.Target;
				if ( assign.Value is not null ) yield return assign.Value;
				break;

			case IrIf branch when branch.Cond is not null:
				yield return branch.Cond;
				break;

			case IrFor loop when loop.Count is not null:
				yield return loop.Count;
				break;

			case IrWhile loop when loop.Cond is not null:
				yield return loop.Cond;
				break;

			case IrReturn ret when ret.Value is not null:
				yield return ret.Value;
				break;

			case IrExprStmt expression when expression.Value is not null:
				yield return expression.Value;
				break;
		}
	}

	/// <summary>Every root expression the code of one stage evaluates.</summary>
	public static IEnumerable<IrExpr> Expressions( IrModule module, ShaderStage stage )
	{
		foreach ( var function in Functions( module, stage ) )
		{
			foreach ( var statement in Statements( function.Body ) )
			{
				foreach ( var expression in Expressions( statement ) ) yield return expression;
			}
		}
	}
}

/// <summary>
/// Which blocks of a <c>.shader</c> each module-level declaration belongs in.
/// <para>
/// <see cref="GlobalDecl.Stages"/> is what a declaration <em>claims</em>, and nothing upstream ever
/// narrows it from <see cref="StageMask.All"/> — so trusting it emits every texture, sampler and
/// uniform into the VS block <em>and</em> the PS block. That is not merely untidy: textures and
/// samplers are a scarce, hard-limited per-program resource, and a graph that samples sixteen
/// textures in the pixel program would burn sixteen vertex-stage slots it never touches.
/// </para>
/// <para>
/// So the referencing set is computed from the IR instead. A declaration nothing references is still
/// emitted once — an unwired blackboard parameter has to keep its slot in the material editor — but
/// into one block rather than all of them.
/// </para>
/// </summary>
public sealed class GlobalPlacement
{
	readonly Dictionary<string, StageMask> _referenced = new( StringComparer.Ordinal );
	readonly StageMask _emitted;
	readonly ShaderStage _primary;

	GlobalPlacement( IrModule module )
	{
		_emitted = module.Meta?.Stages ?? StageMask.VertexPixel;

		if ( _emitted == StageMask.None ) _emitted = StageMask.VertexPixel;

		// Where an unreferenced declaration goes. The pixel program is the one a material parameter is
		// meaningful for, and the only block a compute shader has is its own.
		_primary = _emitted.Contains( ShaderStage.Pixel ) ? ShaderStage.Pixel
			: _emitted.Contains( ShaderStage.Compute ) ? ShaderStage.Compute
			: _emitted.Contains( ShaderStage.Vertex ) ? ShaderStage.Vertex
			: ShaderStage.Geometry;

		var names = new HashSet<string>( StringComparer.Ordinal );

		foreach ( var global in module.Globals )
		{
			if ( global?.Name is not null ) names.Add( global.Name );
		}

		foreach ( var stage in ShaderStages.All )
		{
			if ( stage == ShaderStage.None || !_emitted.Contains( stage ) ) continue;

			Collect( module, stage );
		}

		CollectHelpers( module, names );
		PairSamplers( module );
	}

	/// <summary>Analyse a module. Cheap enough to run once per emitted block.</summary>
	public static GlobalPlacement For( IrModule module ) =>
		module is null ? null : new GlobalPlacement( module );

	/// <summary>The stages whose code mentions a declaration by name.</summary>
	public StageMask ReferencedBy( string name ) =>
		name is not null && _referenced.TryGetValue( name, out var mask ) ? mask : StageMask.None;

	/// <summary>The blocks a declaration should actually be written into.</summary>
	public StageMask StagesFor( GlobalDecl decl )
	{
		if ( decl is null ) return StageMask.None;

		var declared = decl.Stages == StageMask.None ? StageMask.All : decl.Stages;
		var referenced = ReferencedBy( decl.Name );

		if ( referenced != StageMask.None )
		{
			var narrowed = declared & referenced & _emitted;

			// A declaration that claims one stage and is read in another is a bug upstream, not a reason
			// to drop it: emit it where it is read so the generated code still compiles.
			return narrowed != StageMask.None ? narrowed : referenced & _emitted;
		}

		var fallback = declared & _primary.ToMask();

		return fallback != StageMask.None ? fallback : declared & _emitted;
	}

	/// <summary>True when a declaration belongs in a given stage's block.</summary>
	public bool ShouldDeclare( GlobalDecl decl, ShaderStage stage ) => StagesFor( decl ).Contains( stage );

	void Collect( IrModule module, ShaderStage stage )
	{
		foreach ( var expression in IrWalk.Expressions( module, stage ) )
		{
			foreach ( var node in IrExprUtil.Walk( expression ) )
			{
				if ( node is IrGlobalRef reference && reference.Decl is not null ) Mark( reference.Decl.Name, stage );
			}
		}
	}

	/// <summary>
	/// Mark the declarations named inside helper bodies.
	/// <para>
	/// A helper's body is author-supplied text rather than IR, so the only way to know whether it
	/// touches a declaration is to read it. Missing one would emit code naming something the block
	/// never declared, which is a compile error rather than a cosmetic problem, so this deliberately
	/// errs towards declaring too much — and it marks every stage the helper is emitted into, because
	/// <see cref="HlslEmitter.WriteHelpers"/> writes a helper into a block whether or not that block's
	/// code calls it.
	/// </para>
	/// <para>
	/// One pass over each body, matching whole identifiers against a set. Scanning per declaration
	/// instead would be a body-length times declaration-count product on the preview compile path.
	/// </para>
	/// </summary>
	void CollectHelpers( IrModule module, HashSet<string> names )
	{
		if ( names.Count == 0 ) return;

		foreach ( var helper in module.Helpers )
		{
			if ( helper is null ) continue;

			var stages = helper.Stages & _emitted;

			if ( stages == StageMask.None ) continue;

			var body = helper.BodyFor( PrismConstants.BackendHlsl );

			if ( string.IsNullOrEmpty( body ) ) continue;

			foreach ( var identifier in Identifiers( body ) )
			{
				if ( !names.Contains( identifier ) ) continue;

				foreach ( var stage in stages.Stages() ) Mark( identifier, stage );
			}
		}
	}

	/// <summary>Every maximal identifier in a body, in order and with duplicates.</summary>
	static IEnumerable<string> Identifiers( string body )
	{
		var start = -1;

		for ( int i = 0; i <= body.Length; i++ )
		{
			var word = i < body.Length && IsWord( body[i] );

			if ( word && start < 0 )
			{
				start = i;
				continue;
			}

			if ( word || start < 0 ) continue;

			yield return body[start..i];

			start = -1;
		}
	}

	/// <summary>
	/// A texture's sampler has to live wherever the texture does: the sample call names both, and the
	/// pairing is recorded on the declaration rather than showing up as a reference.
	/// </summary>
	void PairSamplers( IrModule module )
	{
		foreach ( var global in module.Globals )
		{
			if ( global?.SamplerName is null ) continue;

			var stages = ReferencedBy( global.Name );

			if ( stages == StageMask.None ) continue;

			foreach ( var stage in stages.Stages() ) Mark( global.SamplerName, stage );
		}
	}

	void Mark( string name, ShaderStage stage )
	{
		if ( string.IsNullOrEmpty( name ) ) return;

		_referenced.TryGetValue( name, out var mask );
		_referenced[name] = mask | stage.ToMask();
	}

	static bool IsWord( char c ) => char.IsLetterOrDigit( c ) || c == '_';
}