using System;
using System.Collections.Generic;
using System.Linq;
using System.Text.Json;
using System.Text.RegularExpressions;

namespace ImportUnityPackage;

/// <summary>
/// Property declarations of a package's own shader: ShaderLab <c>Properties</c> blocks and Shader Graph files.
/// Explicit markers (<c>[MainTexture]</c>, <c>[MainColor]</c>, <c>[Normal]</c>, <c>"bump"</c> defaults, Shader Graph normal-map
/// textures) are evidence of a property's role; display and reference names are only inferences.
/// </summary>
internal sealed class UnityShaderDeclarations
{
	internal sealed record Property( string Name, string Display, bool Texture, bool Color, bool Normal, bool MainTexture, bool MainColor, bool Hidden = false );
	internal List<Property> Properties { get; } = new();
	/// <summary>Every declared property name, of any type.</summary>
	internal HashSet<string> Declared { get; } = new( StringComparer.Ordinal );
	internal string Name { get; private set; }
	internal bool AlphaClip { get; private set; }
	internal bool Transparent { get; private set; }
	internal bool TwoSided { get; private set; }
	internal bool Graph { get; private set; }
	/// <summary>
	/// Color properties whose alpha weights a blend in the shader code (<c>lerp( a, b, _MainColor.a * mask )</c>): the color
	/// is mixed in through a mask (a texture channel or vertex color), so it is neither a tint of the whole surface nor its
	/// opacity.
	/// </summary>
	internal HashSet<string> MaskedColors { get; } = new( StringComparer.Ordinal );

	// Many materials share one shader: the declarations are read once per distinct source (they are not changed after).
	static readonly System.Runtime.CompilerServices.ConditionalWeakTable<string, UnityShaderDeclarations> parsed = new();

	internal static UnityShaderDeclarations Parse( string source ) => parsed.GetValue( source, Read );

	static UnityShaderDeclarations Read( string source )
	{
		var result = new UnityShaderDeclarations();
		if ( source.TrimStart().StartsWith( '{' ) ) result.ReadGraph( source );
		else result.ReadShaderLab( source );
		return result;
	}

	void ReadShaderLab( string source )
	{
		Name = Regex.Match( source, "Shader\\s+\"([^\"]+)\"" ).Groups[1].Value;
		var timeout = TimeSpan.FromSeconds( 2 );
		foreach ( Match match in Regex.Matches( source, @"(?m)^\s*((?:\[[^\]\r\n]*\]\s*)*)([A-Za-z_][A-Za-z0-9_]*)\s*\(\s*""([^""]*)""\s*,\s*([A-Za-z0-9]+)[^)\r\n]*\)\s*(?:=\s*""?([A-Za-z]*))?", RegexOptions.None, timeout ) )
		{
			var attributes = match.Groups[1].Value;
			var type = match.Groups[4].Value;
			bool Has( string attribute ) => Regex.IsMatch( attributes, @"\[\s*" + attribute + @"\s*\]", RegexOptions.IgnoreCase );
			var texture = type.Equals( "2D", StringComparison.OrdinalIgnoreCase );
			Properties.Add( new( match.Groups[2].Value, match.Groups[3].Value, texture, type.Equals( "Color", StringComparison.OrdinalIgnoreCase ),
				texture && (Has( "Normal" ) || match.Groups[5].Value.Equals( "bump", StringComparison.OrdinalIgnoreCase )), Has( "MainTexture" ), Has( "MainColor" ), Has( "HideInInspector" ) ) );
		}
		// Any "_Name ( "Display", ..." line declares a property, whatever its type.
		Declared.UnionWith( Regex.Matches( source, @"(?m)^\s*(?:\[[^\]\r\n]*\]\s*)*([A-Za-z_][A-Za-z0-9_]*)\s*\(\s*""[^""]*""\s*,", RegexOptions.None, timeout ).Select( m => m.Groups[1].Value ) );
		AlphaClip = Regex.IsMatch( source, "\"RenderType\"\\s*=\\s*\"TransparentCutout\"" );
		Transparent = Regex.IsMatch( source, "\"RenderType\"\\s*=\\s*\"Transparent\"" );
		TwoSided = Regex.IsMatch( source, @"(?m)^\s*Cull\s+Off\s*$", RegexOptions.IgnoreCase );
		// Only the blend weight (a lerp's third argument) counts: "lerp( a, b, t ) * _Color.a" is an ordinary opacity use.
		foreach ( Match call in Regex.Matches( source, @"\blerp\s*\(", RegexOptions.None, timeout ) )
			if ( LerpWeight( source, call.Index + call.Length ) is { } weight )
				foreach ( Match use in Regex.Matches( weight, @"\*\s*\(?\s*([A-Za-z_][A-Za-z0-9_]*)\.[aw]\b|\b([A-Za-z_][A-Za-z0-9_]*)\.[aw]\s*\*", RegexOptions.None, timeout ) )
					MaskedColors.Add( use.Groups[1].Success ? use.Groups[1].Value : use.Groups[2].Value );
		MaskedColors.IntersectWith( Properties.Where( p => p.Color ).Select( p => p.Name ) );
	}

	/// <summary>The third argument of a call whose arguments start at <paramref name="start"/>, or null.</summary>
	static string LerpWeight( string source, int start )
	{
		var depth = 0;
		var argument = 0;
		var weightStart = -1;
		for ( var i = start; i < source.Length && i < start + 4096; i++ )
		{
			var c = source[i];
			if ( c == '(' ) depth++;
			else if ( c == ')' )
			{
				if ( depth == 0 ) return weightStart < 0 ? null : source[weightStart..i];
				depth--;
			}
			else if ( c == ',' && depth == 0 && ++argument == 2 ) weightStart = i + 1;
			else if ( c is ';' or '{' or '}' ) return null;
		}
		return null;
	}

	void ReadGraph( string source )
	{
		Graph = true;
		// Shader Graph 10+ stores one JSON document per object; older versions nest properties as JSON strings.
		var documents = new List<JsonElement>();
		var reader = new Utf8JsonReader( System.Text.Encoding.UTF8.GetBytes( source ),
			new JsonReaderOptions { AllowTrailingCommas = true, CommentHandling = JsonCommentHandling.Skip, AllowMultipleValues = true } );
		try
		{
			while ( reader.Read() )
			{
				if ( reader.TokenType != JsonTokenType.StartObject ) continue;
				using var document = JsonDocument.ParseValue( ref reader );
				documents.Add( document.RootElement.Clone() );
			}
		}
		catch ( JsonException ) { }
		foreach ( var root in documents )
		{
			if ( root.TryGetProperty( "m_SerializedProperties", out var legacy ) && legacy.ValueKind == JsonValueKind.Array )
				foreach ( var item in legacy.EnumerateArray() )
					if ( item.TryGetProperty( "JSONnodeData", out var data ) && data.ValueKind == JsonValueKind.String )
					{
						try { using var nested = JsonDocument.Parse( data.GetString() ); Add( nested.RootElement, item.TryGetProperty( "typeInfo", out var info ) && info.TryGetProperty( "fullName", out var full ) ? full.GetString() : "" ); }
						catch ( JsonException ) { }
					}
			var type = root.TryGetProperty( "m_Type", out var typeName ) ? typeName.GetString() ?? "" : "";
			Add( root, type );
			if ( type.EndsWith( ".UniversalTarget", StringComparison.Ordinal ) || type.EndsWith( ".BuiltInTarget", StringComparison.Ordinal ) )
			{
				AlphaClip |= root.TryGetProperty( "m_AlphaClip", out var clip ) && clip.ValueKind == JsonValueKind.True;
				Transparent |= root.TryGetProperty( "m_SurfaceType", out var surface ) && surface.TryGetInt32( out var kind ) && kind == 1;
				TwoSided |= root.TryGetProperty( "m_RenderFace", out var face ) && face.TryGetInt32( out var faces ) && faces == 0;
			}
		}
	}

	void Add( JsonElement element, string type )
	{
		if ( !type.EndsWith( "ShaderProperty", StringComparison.Ordinal ) ) return;
		var texture = type.EndsWith( "Texture2DShaderProperty", StringComparison.Ordinal );
		var color = type.EndsWith( "ColorShaderProperty", StringComparison.Ordinal );
		string Text( string key ) => element.TryGetProperty( key, out var value ) && value.ValueKind == JsonValueKind.String ? value.GetString() : null;
		bool Flag( string key ) => element.TryGetProperty( key, out var value ) && value.ValueKind == JsonValueKind.True;
		var name = Text( "m_OverrideReferenceName" ) is { Length: > 0 } custom ? custom : Text( "m_DefaultReferenceName" ) ?? Text( "m_Name" );
		if ( string.IsNullOrEmpty( name ) ) return;
		Declared.Add( name );
		if ( !texture && !color ) return;
		// Texture2DShaderProperty.DefaultType 3 is the normal-map ("bump") default.
		var normal = texture && element.TryGetProperty( "m_DefaultType", out var defaultType ) && defaultType.TryGetInt32( out var value ) && value == 3;
		Properties.Add( new( name, Text( "m_Name" ) ?? name, texture, color, normal, Flag( "isMainTexture" ), Flag( "isMainColor" ), Flag( "m_Hidden" ) || Flag( "hidden" ) ) );
	}

}