using Editor.Prism.Compiler.Backends;
using Editor.Prism.Core;
using Editor.Prism.Preview;

namespace Editor.Prism.Compiler;

/// <summary>
/// The leaf expressions a backend can offer the debug-channel tail.
/// <para>
/// Every field is an expression in the backend's own language, or <c>null</c> when the generated
/// shader cannot answer that channel — a post-process graph has no vertex colour, a custom shading
/// model has no material struct, a Slang module only carries the interpolants the graph asked for.
/// A null field emits no case at all, so the channel falls through to the shaded result rather than
/// returning garbage. That fall-through is the contract
/// <see cref="PreviewChannelInfo.NeedsShaderSupport"/> promises the viewport.
/// </para>
/// </summary>
public sealed class PreviewChannelEnvironment
{
	/// <summary>Linear base colour, a 3-component expression.</summary>
	public string Albedo { get; set; }

	/// <summary>Coverage, a scalar expression.</summary>
	public string Opacity { get; set; }

	/// <summary>The tangent-space normal the graph produced, a 3-component expression.</summary>
	public string NormalTangent { get; set; }

	/// <summary>The shading normal in world space, a 3-component expression.</summary>
	public string NormalWorld { get; set; }

	/// <summary>Microfacet roughness, a scalar expression.</summary>
	public string Roughness { get; set; }

	/// <summary>Dielectric-to-conductor blend, a scalar expression.</summary>
	public string Metalness { get; set; }

	/// <summary>Baked occlusion, a scalar expression.</summary>
	public string AmbientOcclusion { get; set; }

	/// <summary>Emissive colour, a 3-component expression.</summary>
	public string Emission { get; set; }

	/// <summary>Transmitted light, a 3-component expression.</summary>
	public string Transmission { get; set; }

	/// <summary>Per-instance tint mask, a scalar expression.</summary>
	public string TintMask { get; set; }

	/// <summary>The first texture coordinate set, a 2-component expression.</summary>
	public string Uv0 { get; set; }

	/// <summary>The second texture coordinate set, a 2-component expression.</summary>
	public string Uv1 { get; set; }

	/// <summary>Interpolated vertex colour, a 3- or 4-component expression.</summary>
	public string VertexColor { get; set; }

	/// <summary>World position, a 3-component expression.</summary>
	public string WorldPosition { get; set; }

	/// <summary>
	/// The 2-component expression whose screen-space derivatives drive the derivative and mip-level
	/// channels. Normally the same expression as <see cref="Uv0"/>.
	/// </summary>
	public string DerivativeSource { get; set; }

	/// <summary>The uniform the channel index is read from.</summary>
	public string ChannelUniform { get; set; } = PrismConstants.ChannelAttribute;

	/// <summary>
	/// World-space distance the world-position channel wraps over. One inch-scale unit cube per
	/// 128 units reads well on a preview mesh without going to banding.
	/// </summary>
	public float WorldPositionPeriod { get; set; } = 128f;

	/// <summary>
	/// The texture size the mip-level channel assumes. Nothing in the shader knows which texture the
	/// user cares about, so a fixed reference makes the readout comparable between graphs.
	/// </summary>
	public float MipReferenceSize { get; set; } = 1024f;

	/// <summary>Mip level the channel's ramp saturates at.</summary>
	public float MipRange { get; set; } = 10f;
}

/// <summary>
/// The preview-only tails a generated shader grows in <see cref="CompileMode.Preview"/> and
/// <see cref="CompileMode.Thumbnail"/>: the per-node stage switch that makes node thumbnails cost one
/// compile instead of N, and the debug-channel switch behind the viewport's channel strip.
/// <para>
/// <b>The stage numbering is a contract, not an implementation detail.</b>
/// <see cref="NodeThumbnailService.BuildStageMap"/> recovers it from the emitted artifact's
/// <see cref="SourceMap"/> by numbering nodes in order of their first generated line, starting at
/// one. <see cref="BuildStageMap(SourceMap)"/> below is that same rule, and the backends run it over
/// the source map of a throw-away first pass so the numbers they bake into the switch are the ones
/// the preview will later compute. Zero always means "show the shader's own output".
/// </para>
/// <para>
/// Instrumentation is only ever added to the pixel entry point, and only for values that live there.
/// A node whose work the stage planner placed in the vertex program has a stage id but no case in the
/// switch: its temp simply does not exist in the pixel stage. Selecting it shows the shaded result
/// rather than a wrong colour, which is the same fall-through rule the channels use.
/// </para>
/// </summary>
public static class PreviewInstrumentation
{
	/// <summary>The uniform the preview writes a node's stage id into.</summary>
	public const string StageIdUniform = PrismConstants.StageIdAttribute;

	/// <summary>The uniform the preview writes the selected debug channel into.</summary>
	public const string ChannelUniform = PrismConstants.ChannelAttribute;

	/// <summary>The stage id that means "no override": show whatever the shader normally produces.</summary>
	public const int NoStage = 0;

	/// <summary>Prefix for the locals the debug-channel tail declares.</summary>
	public const string LocalPrefix = "vPrismDbg";

	/// <summary>True when a compile mode wants the preview instrumentation.</summary>
	public static bool IsEnabled( CompileMode mode ) => mode is CompileMode.Preview or CompileMode.Thumbnail;

	// ---- stage numbering ---------------------------------------------------

	/// <summary>
	/// Number every node that produced generated code, in order of its first generated line, starting
	/// at one.
	/// <para>
	/// This must stay identical to <see cref="NodeThumbnailService.BuildStageMap"/>: the backend bakes
	/// these numbers into the switch and the preview recomputes them from the shipped artifact. If the
	/// two ever disagree, every node thumbnail shows a different node's value — a bug that looks like
	/// a rendering problem and is not one.
	/// </para>
	/// </summary>
	public static IReadOnlyDictionary<NodeId, int> BuildStageMap( SourceMap map )
	{
		var stages = new Dictionary<NodeId, int>();

		if ( map is null ) return stages;

		var next = 1;

		foreach ( var entry in map.Entries )
		{
			if ( !entry.Node.IsValid ) continue;
			if ( stages.ContainsKey( entry.Node ) ) continue;

			stages[entry.Node] = next++;
		}

		return stages;
	}

	/// <summary>
	/// The node-to-stage-id map for a finished compile, recovered from the s&amp;box artifact's source
	/// map. This is what a panel calls to find the value to push into
	/// <see cref="StageIdUniform"/> for the selected node.
	/// </summary>
	public static IReadOnlyDictionary<NodeId, int> StageMap( CompileResult result ) =>
		BuildStageMap( result?.Artifact( PrismConstants.BackendHlsl )?.SourceMap );

	/// <summary>The stage id of one node, or <see cref="NoStage"/> when it produced no code.</summary>
	public static int StageIdOf( IReadOnlyDictionary<NodeId, int> stages, NodeId node ) =>
		stages is not null && node.IsValid && stages.TryGetValue( node, out var id ) ? id : NoStage;

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

	/// <summary>
	/// The two uniforms the instrumentation reads, in VFX annotation form.
	/// <para>
	/// Both are attribute-bound, so the preview and the thumbnail rig change what they display with a
	/// single <c>RenderAttributes.Set</c> and no recompile at all. They default to zero, which is the
	/// shaded result, so a shader compiled with instrumentation still renders normally when nothing
	/// pushes them.
	/// </para>
	/// </summary>
	public static IReadOnlyList<string> HlslDeclarations() =>
	[
		$"int {StageIdUniform} < Attribute( \"{StageIdUniform}\" ); Default( 0 ); >;",
		$"int {ChannelUniform} < Attribute( \"{ChannelUniform}\" ); Default( 0 ); >;"
	];

	/// <summary>The same two uniforms as plain Slang module-scope uniforms.</summary>
	public static IReadOnlyList<string> SlangDeclarations() =>
	[
		$"uniform int {StageIdUniform};",
		$"uniform int {ChannelUniform};"
	];

	// ---- the stage switch --------------------------------------------------

	/// <summary>
	/// True when a value can be shown as a colour. Matrices, textures, samplers, buffers and structs
	/// cannot, so they get no case in the switch and fall through to the shaded result.
	/// </summary>
	public static bool CanShow( ShaderType type ) => type.IsScalarOrVector && !type.IsVoid;

	/// <summary>
	/// Widen an expression of any showable type to an opaque <c>float4</c>. Booleans and integers
	/// convert rather than bitcast, so a boolean mask reads as black and white.
	/// <para>
	/// Alpha is always one, including for a four-component value whose own alpha is dropped. A preview
	/// mesh and a thumbnail are both drawn over a transparent background, so honouring a node's alpha
	/// would make half the thumbnails in a graph invisible; the alpha of a surface is what the
	/// <see cref="PreviewChannel.Opacity"/> channel is for.
	/// </para>
	/// </summary>
	public static string AsColor( string expression, ShaderType type )
	{
		if ( string.IsNullOrEmpty( expression ) || !CanShow( type ) ) return null;

		return Math.Clamp( type.Components, 1, 4 ) switch
		{
			1 => $"float4( ( {expression} ).xxx, 1.0f )",
			2 => $"float4( ( {expression} ), 0.0f, 1.0f )",
			3 => $"float4( ( {expression} ), 1.0f )",
			_ => $"float4( ( {expression} ).xyz, 1.0f )"
		};
	}

	/// <summary>
	/// The one-line switch case for a node's value:
	/// <c>if ( g_iPrismStageId == 7 ) return float4( t7, 1.0f );</c>
	/// Returns null when the value cannot be shown, in which case nothing is emitted.
	/// </summary>
	public static string StageCase( int stageId, string expression, ShaderType type )
	{
		if ( stageId <= NoStage ) return null;

		var colour = AsColor( expression, type );

		return colour is null ? null : $"if ( {StageIdUniform} == {stageId} ) return {colour};";
	}

	/// <summary>The comment that introduces the two preview uniforms in the generated pixel program.</summary>
	public static IReadOnlyList<string> Banner() =>
	[
		"// Prism preview instrumentation. Every node's value and every debug channel is reachable from",
		$"// this one shader: {StageIdUniform} picks a node, {ChannelUniform} picks a channel, and zero",
		"// means the shader's own output. Both are render attributes, so switching costs no recompile."
	];

	// ---- the debug channels ------------------------------------------------

	/// <summary>
	/// The debug-channel tail: one <c>if</c> per channel the shader can answer, in the order
	/// <see cref="PreviewChannels"/> declares them, plus the locals the derivative channels need.
	/// <para>
	/// Everything is flat — one statement per line, no nesting — so a caller only has to know how to
	/// write a line, and the same list serves the block-file writer and the Slang emitter. An empty
	/// result means this shader can answer no channel at all and nothing should be written.
	/// </para>
	/// <para>
	/// Channels the engine renders on its own — overdraw, shading complexity, wireframe — are absent
	/// by design: <see cref="PreviewChannels.Apply(PreviewChannel, CameraComponent)"/> serves those
	/// through <c>SceneCameraDebugMode</c> and needs nothing from the generated shader.
	/// </para>
	/// </summary>
	public static IReadOnlyList<string> ChannelLines( PreviewChannelEnvironment env )
	{
		if ( env is null ) return Array.Empty<string>();

		var uniform = string.IsNullOrEmpty( env.ChannelUniform ) ? ChannelUniform : env.ChannelUniform;
		var derivatives = !string.IsNullOrEmpty( env.DerivativeSource );
		var cases = new List<string>();

		Case( PreviewChannel.Albedo, Rgb( env.Albedo ) );
		Case( PreviewChannel.Opacity, Grey( env.Opacity ) );
		Case( PreviewChannel.NormalWorld, Signed( env.NormalWorld ) );
		Case( PreviewChannel.NormalTangent, Signed( env.NormalTangent ) );
		Case( PreviewChannel.Roughness, Grey( env.Roughness ) );
		Case( PreviewChannel.Metalness, Grey( env.Metalness ) );
		Case( PreviewChannel.AmbientOcclusion, Grey( env.AmbientOcclusion ) );
		Case( PreviewChannel.Emission, Rgb( env.Emission ) );
		Case( PreviewChannel.Transmission, Rgb( env.Transmission ) );
		Case( PreviewChannel.TintMask, Grey( env.TintMask ) );
		Case( PreviewChannel.Uv0, Uv( env.Uv0 ) );
		Case( PreviewChannel.Uv1, Uv( env.Uv1 ) );
		Case( PreviewChannel.VertexColor, Rgb( env.VertexColor, ".xyz" ) );
		Case( PreviewChannel.WorldPosition, Wrapped( env.WorldPosition, env.WorldPositionPeriod ) );

		if ( derivatives )
		{
			Case( PreviewChannel.Derivatives,
				$"float4( saturate( abs( {LocalPrefix}Ddx ) * 128.0f ), saturate( length( {LocalPrefix}Ddy ) * 128.0f ), 1.0f )" );

			Case( PreviewChannel.MipLevel,
				$"float4( saturate( {LocalPrefix}Mip / {Number( env.MipRange )} ).xxx, 1.0f )" );
		}

		if ( cases.Count == 0 ) return Array.Empty<string>();

		var lines = new List<string>( cases.Count + 8 )
		{
			"// Prism preview: the debug channels behind the viewport's channel strip.",
			$"// {uniform} is 0 for the shaded result. A channel this shader cannot answer has no case",
			"// here at all, so selecting it keeps showing the shaded result rather than a wrong colour."
		};

		if ( derivatives )
		{
			var reference = Number( env.MipReferenceSize );

			// Screen-space derivatives have to be taken in uniform control flow, so they are computed
			// before the switch rather than inside the two cases that read them.
			lines.Add( $"float2 {LocalPrefix}Ddx = ddx( {env.DerivativeSource} );" );
			lines.Add( $"float2 {LocalPrefix}Ddy = ddy( {env.DerivativeSource} );" );
			lines.Add( $"float2 {LocalPrefix}MipDx = {LocalPrefix}Ddx * {reference};" );
			lines.Add( $"float2 {LocalPrefix}MipDy = {LocalPrefix}Ddy * {reference};" );
			lines.Add( $"float {LocalPrefix}Mip = 0.5f * log2( max( max( dot( {LocalPrefix}MipDx, {LocalPrefix}MipDx ), " +
				$"dot( {LocalPrefix}MipDy, {LocalPrefix}MipDy ) ), 1.0e-8f ) );" );
		}

		lines.AddRange( cases );

		return lines;

		void Case( PreviewChannel channel, string colour )
		{
			if ( string.IsNullOrEmpty( colour ) ) return;

			cases.Add( $"if ( {uniform} == {(int)channel} ) return {colour};" );
		}
	}

	static string Rgb( string expression, string swizzle = null ) =>
		string.IsNullOrEmpty( expression )
			? null
			: $"float4( ( {expression} ){swizzle}, 1.0f )";

	static string Grey( string expression ) =>
		string.IsNullOrEmpty( expression )
			? null
			: $"float4( saturate( {expression} ).xxx, 1.0f )";

	static string Signed( string expression ) =>
		string.IsNullOrEmpty( expression )
			? null
			: $"float4( ( {expression} ) * 0.5f + 0.5f, 1.0f )";

	static string Uv( string expression ) =>
		string.IsNullOrEmpty( expression )
			? null
			: $"float4( frac( {expression} ), 0.0f, 1.0f )";

	static string Wrapped( string expression, float period ) =>
		string.IsNullOrEmpty( expression )
			? null
			: $"float4( frac( ( {expression} ) / {Number( period )} ), 1.0f )";

	/// <summary>
	/// Spell a constant the way both backends spell one, so the generated text stays byte-identical
	/// between regenerations of an unchanged graph.
	/// </summary>
	static string Number( float value ) => HlslIntrinsics.Number( value, ScalarKind.Float );
}