using Editor.Prism.Compiler.Backends;
using Editor.Prism.Core;
using System.Text;
using System.Text.RegularExpressions;

namespace Editor.Prism.Toolchain;

/// <summary>Knobs for wrapping a raw buffer in a synthetic shader.</summary>
public sealed class ShaderProbeOptions
{
	/// <summary>File stem the probe is written under. Only ever seen in diagnostics.</summary>
	public string Name { get; set; } = "prism_probe";

	/// <summary>Which stage the buffer's code belongs to.</summary>
	public ShaderStage Stage { get; set; } = ShaderStage.Pixel;

	/// <summary>
	/// Pull in the standard lighting headers. Off by default: a fragment of maths does not need them and
	/// they cost compile time.
	/// </summary>
	public bool Lit { get; set; }

	/// <summary>Extra includes emitted into the program block ahead of the buffer.</summary>
	public IReadOnlyList<string> Includes { get; set; } = Array.Empty<string>();

	/// <summary>Extra <c>#define</c>s emitted into the COMMON block.</summary>
	public IReadOnlyList<string> Defines { get; set; } = Array.Empty<string>();

	/// <summary>
	/// Append a trivial entry point after the buffer. Turned off automatically when the buffer already
	/// declares one, because two <c>MainPs</c> is a worse error than none.
	/// </summary>
	public bool DeclareEntryPoint { get; set; } = true;
}

/// <summary>
/// A synthetic shader wrapping a raw buffer, plus everything needed to map a diagnostic in the wrapper
/// back to the buffer the user is actually editing.
/// </summary>
/// <param name="Text">The complete generated file.</param>
/// <param name="FileName">Suggested file name, extension included.</param>
/// <param name="BodyLine">1-based line of <paramref name="Text"/> where the buffer's first line lands.</param>
/// <param name="Stage">The stage the buffer was wrapped for.</param>
public sealed record ShaderProbe( string Text, string FileName, int BodyLine, ShaderStage Stage )
{
	/// <summary>Characters prefixed to every buffer line, which columns must be corrected by.</summary>
	public int BodyColumn { get; init; }

	/// <summary>How many lines the buffer had.</summary>
	public int BodyLineCount { get; init; }

	/// <summary>True when the buffer was already a complete shader and nothing was wrapped around it.</summary>
	public bool IsPassthrough { get; init; }

	/// <summary>True when a probe line falls inside the user's buffer rather than the scaffolding.</summary>
	public bool Contains( int probeLine ) =>
		probeLine >= BodyLine && ( BodyLineCount <= 0 || probeLine < BodyLine + BodyLineCount );

	/// <summary>Translate a span in the probe back to the buffer it came from.</summary>
	public SourceSpan MapBack( SourceSpan span, string bufferFile = null )
	{
		var file = bufferFile ?? span.File;

		if ( IsPassthrough ) return new SourceSpan( file, span.Line, span.Column, span.EndLine, span.EndColumn );

		var line = span.Line - BodyLine + 1;
		var endLine = span.EndLine - BodyLine + 1;
		var column = Math.Max( 1, span.Column - BodyColumn );
		var endColumn = Math.Max( column, span.EndColumn - BodyColumn );

		if ( line < 1 ) line = 1;
		if ( endLine < line ) endLine = line;

		return new SourceSpan( file, line, column, endLine, endColumn );
	}

	/// <summary>
	/// Translate a diagnostic back to the buffer. A diagnostic that landed in the scaffolding keeps its
	/// message but loses its bogus location and says where it really came from, because pointing a user
	/// at line 40 of a file they never wrote is worse than pointing at nothing.
	/// </summary>
	public Diagnostic MapBack( Diagnostic diagnostic, string bufferFile = null )
	{
		if ( diagnostic?.Span is not { } span ) return diagnostic;
		if ( IsPassthrough ) return diagnostic;

		if ( Contains( span.Line ) ) return diagnostic.WithSpan( MapBack( span, bufferFile ) );

		var note = "Reported in the scaffolding Prism wrapped around this file, not in the file itself.";

		return diagnostic with
		{
			Span = null,
			Detail = string.IsNullOrWhiteSpace( diagnostic.Detail ) ? note : diagnostic.Detail + "\n" + note
		};
	}

	/// <summary>Translate a batch of diagnostics back to the buffer.</summary>
	public IReadOnlyList<Diagnostic> MapBack( IEnumerable<Diagnostic> diagnostics, string bufferFile = null )
	{
		var results = new List<Diagnostic>();

		if ( diagnostics is null ) return results;

		foreach ( var diagnostic in diagnostics )
		{
			if ( diagnostic is null ) continue;

			results.Add( MapBack( diagnostic, bufferFile ) );
		}

		return results;
	}

	/// <inheritdoc/>
	public override string ToString() => $"{FileName}: body at line {BodyLine} ({BodyLineCount} lines)";
}

/// <summary>
/// Wraps a raw <c>.hlsl</c> or <c>.slang</c> buffer in the smallest legal shader that will compile, so
/// the text editor can get real compiler diagnostics for a file the engine would otherwise refuse.
/// <para>
/// The engine only compiles VFX block files: a plain <c>.hlsl</c> is not a shader, it is an include.
/// Handing one to <c>EditorUtility.CompileShader</c> gets a block-header parse failure and no useful
/// message at all. So the buffer is spliced into a minimal <c>FEATURES/MODES/COMMON/VS/PS</c> skeleton
/// with a trivial entry point, compiled for real, and every diagnostic is mapped back by subtracting
/// the scaffolding's line and column offsets.
/// </para>
/// <para>
/// Two details are deliberate. The buffer is indented by one tab, because the native block masker is
/// whitespace sensitive and a stray <c>}</c> at column zero inside the buffer would otherwise tear the
/// file in half — the column offset is recorded and undone on the way back. And a buffer that already
/// declares its own entry point does not get a second one.
/// </para>
/// </summary>
public static class ShaderProbeBuilder
{
	const string NewLine = "\r\n";
	const string Indent = "\t";

	static readonly string[] s_blocks =
	{
		"HEADER", "MODES", "FEATURES", "COMMON", "VS", "PS", "GS", "CS", "PS_RENDER_STATE", "RTX"
	};

	static readonly Regex s_slangEntry = new(
		@"\[\s*shader\s*\(\s*""(?<stage>\w+)""\s*\)\s*\]",
		RegexOptions.Compiled | RegexOptions.CultureInvariant );

	static readonly Regex s_identifierCall = new(
		@"(?<name>[A-Za-z_]\w*)\s*\(", RegexOptions.Compiled | RegexOptions.CultureInvariant );

	/// <summary>
	/// Wrap a buffer, choosing the strategy from its extension or language id. Accepts
	/// <c>.hlsl</c>, <c>.slang</c>, <c>.shader</c>, or the bare words <c>hlsl</c>, <c>slang</c>,
	/// <c>vfx</c>.
	/// </summary>
	public static ShaderProbe For( string source, string extensionOrLanguage, ShaderProbeOptions options = null )
	{
		var kind = ( extensionOrLanguage ?? string.Empty ).Trim().TrimStart( '.' ).ToLowerInvariant();

		return kind switch
		{
			"slang" => ForSlang( source, options ),
			"shader" or "vfx" => ForShaderFile( source, options?.Name ),
			_ => ForHlsl( source, options )
		};
	}

	/// <summary>
	/// Wrap a raw HLSL buffer. A buffer that is already a complete block file is passed straight
	/// through, so the caller never has to decide.
	/// </summary>
	public static ShaderProbe ForHlsl( string source, ShaderProbeOptions options = null )
	{
		var text = source ?? string.Empty;

		if ( LooksLikeShaderFile( text ) ) return ForShaderFile( text, options?.Name );

		options ??= new ShaderProbeOptions();

		return PrismLog.Guard( "Building a shader probe", () => Build( text, options ),
			Passthrough( text, options.Name, options.Stage ) );
	}

	/// <summary>A buffer that is already a shader needs no wrapper; the identity probe keeps call sites simple.</summary>
	public static ShaderProbe ForShaderFile( string source, string name = null ) =>
		Passthrough( source ?? string.Empty, name, ShaderStage.Pixel );

	/// <summary>
	/// Prepare a Slang buffer for <c>slangc</c>. There is no scaffolding to add — slangc happily checks
	/// a module on its own — beyond pinning the language version when the buffer forgot to, which stops
	/// an upgraded compiler silently reinterpreting 2018 semantics as 2026 ones.
	/// </summary>
	public static ShaderProbe ForSlang( string source, ShaderProbeOptions options = null )
	{
		var text = source ?? string.Empty;
		var name = options?.Name ?? "prism_probe";
		var stage = options?.Stage ?? ShaderStage.Pixel;
		var fileName = $"{Sanitize( name )}.{PrismConstants.SlangExtension}";

		if ( text.TrimStart().StartsWith( "#language", StringComparison.Ordinal ) )
		{
			return new ShaderProbe( text, fileName, 1, stage )
			{
				IsPassthrough = true,
				BodyLineCount = CountLines( text )
			};
		}

		var header = $"#language slang {PrismConstants.SlangLanguageVersion}{NewLine}";

		return new ShaderProbe( header + text, fileName, 2, stage )
		{
			BodyColumn = 0,
			BodyLineCount = CountLines( text )
		};
	}

	/// <summary>
	/// The entry points a Slang buffer declares, read off its <c>[shader("...")]</c> attributes. Feed
	/// these to <see cref="SlangcValidator"/>, which needs an explicit entry point and stage per
	/// invocation.
	/// </summary>
	public static IReadOnlyList<SlangEntryPoint> DiscoverSlangEntryPoints( string source )
	{
		var results = new List<SlangEntryPoint>();

		if ( string.IsNullOrWhiteSpace( source ) ) return results;

		PrismLog.Guard( "Discovering Slang entry points", () =>
		{
			foreach ( Match match in s_slangEntry.Matches( source ) )
			{
				var stage = StageOf( match.Groups["stage"].Value );

				if ( stage == ShaderStage.None ) continue;

				var call = s_identifierCall.Match( source, match.Index + match.Length );

				if ( !call.Success ) continue;

				var name = call.Groups["name"].Value;

				if ( results.Any( x => string.Equals( x.Name, name, StringComparison.Ordinal ) ) ) continue;

				results.Add( new SlangEntryPoint( name, stage ) );
			}
		} );

		return results;
	}

	/// <summary>True when the text is already a VFX block file rather than a plain include.</summary>
	public static bool LooksLikeShaderFile( string text )
	{
		if ( string.IsNullOrWhiteSpace( text ) ) return false;

		foreach ( var line in text.Split( '\n' ) )
		{
			var trimmed = line.Trim();

			foreach ( var block in s_blocks )
			{
				if ( string.Equals( trimmed, block, StringComparison.Ordinal ) ) return true;
			}
		}

		return false;
	}

	/// <summary>True when the buffer already declares the entry point for a stage.</summary>
	public static bool DeclaresEntryPoint( string text, ShaderStage stage )
	{
		if ( string.IsNullOrWhiteSpace( text ) ) return false;

		var name = stage.EntryPoint();

		return !string.IsNullOrEmpty( name ) &&
			Regex.IsMatch( text, $@"\b{Regex.Escape( name )}\s*\(", RegexOptions.CultureInvariant );
	}

	// ---- construction ----------------------------------------------------

	static ShaderProbe Build( string body, ShaderProbeOptions options )
	{
		var stage = options.Stage is ShaderStage.Vertex or ShaderStage.Pixel ? options.Stage : ShaderStage.Pixel;
		var declare = options.DeclareEntryPoint && !DeclaresEntryPoint( body, stage );

		var builder = new StringBuilder();
		var line = 0;

		void Write( string text = "" )
		{
			builder.Append( text ).Append( NewLine );
			line++;
		}

		Write( "// Generated by Prism so a raw HLSL buffer gets real compiler diagnostics. Do not edit." );
		Write();

		Write( SboxShaderTemplates.BlockFeatures );
		Write( "{" );
		Write( Indent + $"#include \"{SboxShaderTemplates.IncludeFeatures}\"" );
		Write( "}" );
		Write();

		Write( SboxShaderTemplates.BlockModes );
		Write( "{" );
		Write( Indent + SboxShaderTemplates.ModeStatement( SboxShaderTemplates.ModeForward ) );
		Write( "}" );
		Write();

		Write( SboxShaderTemplates.BlockCommon );
		Write( "{" );
		Write( Indent + $"#include \"{SboxShaderTemplates.IncludeShared}\"" );

		foreach ( var define in options.Defines ?? Array.Empty<string>() )
		{
			if ( string.IsNullOrWhiteSpace( define ) ) continue;

			Write( Indent + "#define " + define.Trim() );
		}

		Write( "}" );
		Write();

		Write( $"struct {SboxShaderTemplates.StructVertexInput}" );
		Write( "{" );
		Write( Indent + $"#include \"{SboxShaderTemplates.IncludeVertexInput}\"" );
		Write( "};" );
		Write();

		Write( $"struct {SboxShaderTemplates.StructPixelInput}" );
		Write( "{" );
		Write( Indent + $"#include \"{SboxShaderTemplates.IncludePixelInput}\"" );
		Write( "};" );
		Write();

		var bodyLine = 0;
		var bodyLines = 0;

		// The vertex block always exists: a shader with no VS does not link, even when the buffer under
		// test is pixel-stage code.
		Write( ShaderStage.Vertex.BlockName() );
		Write( "{" );
		Write( Indent + $"#include \"{SboxShaderTemplates.IncludeVertex}\"" );

		if ( stage == ShaderStage.Vertex )
		{
			Write();
			bodyLine = line + 1;
			bodyLines = Splice( builder, body, ref line );
			Write();
		}

		if ( stage != ShaderStage.Vertex || declare )
		{
			Write();
			Write( Indent + $"{SboxShaderTemplates.StructPixelInput} {PrismConstants.EntryPointVertex}" +
				$"( {SboxShaderTemplates.StructVertexInput} {SboxShaderTemplates.VertexInputLocal} )" );
			Write( Indent + "{" );
			Write( Indent + Indent + $"{SboxShaderTemplates.StructPixelInput} {SboxShaderTemplates.PixelInputLocal} " +
				$"= ProcessVertex( {SboxShaderTemplates.VertexInputLocal} );" );
			Write( Indent + Indent + SboxShaderTemplates.SurfaceVertexEpilogue );
			Write( Indent + "}" );
		}

		Write( "}" );
		Write();

		Write( ShaderStage.Pixel.BlockName() );
		Write( "{" );
		Write( Indent + $"#include \"{SboxShaderTemplates.IncludePixel}\"" );

		foreach ( var include in options.Includes ?? Array.Empty<string>() )
		{
			if ( string.IsNullOrWhiteSpace( include ) ) continue;

			Write( Indent + $"#include \"{include.Trim().Trim( '"' )}\"" );
		}

		if ( stage == ShaderStage.Pixel )
		{
			Write();
			bodyLine = line + 1;
			bodyLines = Splice( builder, body, ref line );
		}

		if ( stage != ShaderStage.Pixel || declare )
		{
			Write();
			Write( Indent + $"float4 {PrismConstants.EntryPointPixel}" +
				$"( {SboxShaderTemplates.StructPixelInput} {SboxShaderTemplates.PixelInputLocal} ) : SV_Target0" );
			Write( Indent + "{" );

			Write( options.Lit
				? Indent + Indent + $"return ShadingModelStandard::Shade( Material::Init( {SboxShaderTemplates.PixelInputLocal} ) );"
				: Indent + Indent + "return float4( 1.0, 1.0, 1.0, 1.0 );" );

			Write( Indent + "}" );
		}

		Write( "}" );

		return new ShaderProbe( builder.ToString(),
			$"{Sanitize( options.Name )}.{PrismConstants.ShaderExtension}",
			bodyLine <= 0 ? 1 : bodyLine, stage )
		{
			BodyColumn = Indent.Length,
			BodyLineCount = bodyLines
		};
	}

	/// <summary>
	/// Append the user's buffer, one tab in. The indent keeps a stray brace at column zero from tearing
	/// the block file in half; <see cref="ShaderProbe.BodyColumn"/> undoes it on the way back.
	/// </summary>
	static int Splice( StringBuilder builder, string body, ref int line )
	{
		var lines = ( body ?? string.Empty ).Replace( "\r\n", "\n" ).Split( '\n' );

		foreach ( var text in lines )
		{
			builder.Append( text.Length == 0 ? string.Empty : Indent ).Append( text ).Append( NewLine );
			line++;
		}

		return lines.Length;
	}

	static ShaderProbe Passthrough( string text, string name, ShaderStage stage ) =>
		new( text, $"{Sanitize( name ?? "prism_probe" )}.{PrismConstants.ShaderExtension}", 1, stage )
		{
			IsPassthrough = true,
			BodyLineCount = CountLines( text )
		};

	static int CountLines( string text ) =>
		string.IsNullOrEmpty( text ) ? 0 : text.Replace( "\r\n", "\n" ).Split( '\n' ).Length;

	static ShaderStage StageOf( string slang ) => slang switch
	{
		"vertex" => ShaderStage.Vertex,
		"fragment" or "pixel" => ShaderStage.Pixel,
		"geometry" => ShaderStage.Geometry,
		"compute" => ShaderStage.Compute,
		_ => ShaderStage.None
	};

	static string Sanitize( string name )
	{
		if ( string.IsNullOrWhiteSpace( name ) ) return "prism_probe";

		var builder = new StringBuilder( name.Length );

		foreach ( var c in name )
		{
			builder.Append( char.IsLetterOrDigit( c ) || c is '_' or '-' ? c : '_' );
		}

		var cleaned = builder.ToString().Trim( '_', '-' );

		return cleaned.Length == 0 ? "prism_probe" : cleaned;
	}
}