using System;
using System.Collections.Generic;
using System.Globalization;
using System.IO;
using System.Linq;
using System.Text;
using System.Text.RegularExpressions;

namespace ImportUnityPackage;

/// <summary>A reference to another Unity object: {fileID: N, guid: G, type: T}.</summary>
public readonly record struct UnityReference( long FileId, string Guid )
{
	// Unity writes all-zero GUIDs for built-in or missing references.
	public bool HasGuid => !string.IsNullOrEmpty( Guid ) && Guid.Any( c => c != '0' );
}

/// <summary>A node of Unity's YAML subset: a scalar, an ordered mapping (duplicate keys kept) or a sequence.</summary>
public sealed class YamlNode
{
	static readonly IReadOnlyList<KeyValuePair<string, YamlNode>> NoPairs = Array.Empty<KeyValuePair<string, YamlNode>>();
	static readonly IReadOnlyList<YamlNode> NoItems = Array.Empty<YamlNode>();
	YamlNode() { }
	internal static YamlNode Scalar( string value ) => new() { Value = value };
	internal static YamlNode Map( List<KeyValuePair<string, YamlNode>> pairs ) => new() { Pairs = pairs };
	internal static YamlNode List( List<YamlNode> items ) => new() { Items = items };

	/// <summary>Scalar text; null for mappings and sequences.</summary>
	public string Value { get; private init; }
	public IReadOnlyList<KeyValuePair<string, YamlNode>> Pairs { get; private init; } = NoPairs;
	public IReadOnlyList<YamlNode> Items { get; private init; } = NoItems;
	public bool IsList => Items != NoItems;

	/// <summary>The first value with this key, or null.</summary>
	public YamlNode this[string key]
	{
		get
		{
			foreach ( var pair in Pairs ) if ( pair.Key == key ) return pair.Value;
			return null;
		}
	}
	/// <summary>Every value with this key, in order. Unity 4 repeats keys such as "data".</summary>
	public IEnumerable<YamlNode> All( string key ) => Pairs.Where( p => p.Key == key ).Select( p => p.Value );
	/// <summary>The first value with this key anywhere below this node, in document order.</summary>
	public YamlNode Find( string key )
	{
		foreach ( var pair in Pairs )
		{
			if ( pair.Key == key ) return pair.Value;
			if ( pair.Value.Find( key ) is YamlNode nested ) return nested;
		}
		foreach ( var item in Items ) if ( item.Find( key ) is YamlNode nested ) return nested;
		return null;
	}
	public double? Number => double.TryParse( Value, NumberStyles.Float, CultureInfo.InvariantCulture, out var number ) && double.IsFinite( number ) ? number : null;
	public long? Integer => long.TryParse( Value, NumberStyles.Integer, CultureInfo.InvariantCulture, out var number ) ? number : null;
	public UnityReference Reference => new( this["fileID"]?.Integer ?? 0, this["guid"]?.Value ?? "" );
}

/// <summary>One object document: "--- !u!ClassId &amp;FileId [stripped]" followed by "Type: body".</summary>
public sealed record UnityYamlObject( int ClassId, long FileId, bool Stripped, string Type, YamlNode Body );

/// <summary>Reads Unity's text serialization (assets, prefabs, scenes and .meta files) into nodes.</summary>
public static class UnityYaml
{
	static readonly Regex Header = new( @"^---\s+!u!(?<class>-?[0-9]+)\s+&(?<id>-?[0-9]+)(?<stripped>\s+stripped)?" );
	readonly record struct Line( int Indent, string Content );

	/// <summary>Object documents in file order. A file without document headers (such as .meta) yields one object with no type.</summary>
	public static IReadOnlyList<UnityYamlObject> Parse( string text )
	{
		var objects = new List<UnityYamlObject>();
		var lines = new List<Line>();
		(int Class, long Id, bool Stripped)? header = null;
		void Flush()
		{
			if ( header == null && lines.Count == 0 ) return;
			var root = new Reader( lines ).Document();
			// Object documents hold a single "TypeName:" key whose value is the object body.
			if ( header is { } h ) objects.Add( new( h.Class, h.Id, h.Stripped, root.Pairs.FirstOrDefault().Key, root.Pairs.FirstOrDefault().Value ?? root ) );
			else objects.Add( new( 0, 0, false, null, root ) );
			lines.Clear();
		}
		foreach ( var raw in text.Split( '\n' ) )
		{
			var line = raw.TrimEnd( '\r' );
			if ( line.StartsWith( "---", StringComparison.Ordinal ) )
			{
				Flush();
				var match = Header.Match( line );
				header = match.Success && int.TryParse( match.Groups["class"].Value, NumberStyles.Integer, CultureInfo.InvariantCulture, out var classId ) &&
					long.TryParse( match.Groups["id"].Value, NumberStyles.Integer, CultureInfo.InvariantCulture, out var fileId )
					? (classId, fileId, match.Groups["stripped"].Success) : (0, 0, false);
				continue;
			}
			if ( line.StartsWith( '%' ) && header == null && lines.Count == 0 ) continue; // %YAML / %TAG directives
			var indent = 0;
			while ( indent < line.Length && line[indent] == ' ' ) indent++;
			if ( indent == line.Length ) continue;
			lines.Add( new( indent, line[indent..].TrimEnd() ) );
		}
		Flush();
		return objects;
	}

	/// <summary>Bodies of objects with this type name, including a headerless root that contains it as a key.</summary>
	public static IEnumerable<YamlNode> Objects( IEnumerable<UnityYamlObject> objects, string type ) =>
		objects.Select( o => o.Type == type ? o.Body : o.Type == null ? o.Body[type] : null ).Where( b => b != null );

	/// <summary>The root mapping of a single-document file such as a .meta.</summary>
	public static YamlNode Document( string text ) => Parse( text ).FirstOrDefault()?.Body;

	sealed class Reader( List<Line> source )
	{
		readonly Line[] lines = source.ToArray();
		int i;

		public YamlNode Document() => i < lines.Length ? Block( lines[i].Indent ) : YamlNode.Map( new() );

		static bool IsItem( Line line ) => line.Content == "-" || line.Content.StartsWith( "- ", StringComparison.Ordinal );
		// Unity's documents nest a few dozen levels; a deeper (malformed or crafted) file stops with an error instead of
		// overflowing the stack.
		int depth;
		YamlNode Block( int indent )
		{
			if ( ++depth > MaxDepth ) throw new InvalidDataException( "The YAML document is nested too deeply." );
			try { return IsItem( lines[i] ) ? Sequence( indent ) : Mapping( indent ); }
			finally { depth--; }
		}

		YamlNode Mapping( int indent )
		{
			var pairs = new List<KeyValuePair<string, YamlNode>>();
			while ( i < lines.Length && lines[i].Indent >= indent )
			{
				var line = lines[i];
				// Tolerate malformed lines rather than failing the whole file.
				if ( line.Indent > indent || IsItem( line ) || SplitKey( line.Content ) is not { } split ) { i++; continue; }
				pairs.Add( new( split.Key, AfterKey( split.Value, indent ) ) );
			}
			return YamlNode.Map( pairs );
		}

		YamlNode Sequence( int indent )
		{
			var items = new List<YamlNode>();
			while ( i < lines.Length && lines[i].Indent == indent && IsItem( lines[i] ) )
			{
				var rest = lines[i].Content[1..];
				var content = rest.TrimStart();
				if ( content.Length == 0 )
				{
					i++;
					items.Add( i < lines.Length && lines[i].Indent > indent ? Block( lines[i].Indent ) : YamlNode.Scalar( "" ) );
					continue;
				}
				// "- key: value" starts a mapping (or "- - x" a sequence) at the column after the dash.
				var column = indent + 1 + rest.Length - content.Length;
				var line = new Line( column, content );
				if ( IsItem( line ) || SplitKey( content ) != null ) { lines[i] = line; items.Add( Block( column ) ); }
				else { i++; items.Add( Inline( content, indent ) ); }
			}
			return YamlNode.List( items );
		}

		YamlNode AfterKey( string rest, int indent )
		{
			i++;
			if ( rest.Length > 0 ) return Inline( rest, indent );
			if ( i < lines.Length && lines[i].Indent > indent ) return Block( lines[i].Indent );
			// Unity writes a key's sequence items at the key's own indentation.
			if ( i < lines.Length && lines[i].Indent == indent && IsItem( lines[i] ) ) return Sequence( indent );
			return YamlNode.Scalar( "" );
		}

		/// <summary>A value starting on an already consumed line; continuation lines are more indented than the owner.</summary>
		YamlNode Inline( string text, int indent )
		{
			string More() => i < lines.Length && lines[i].Indent > indent ? lines[i++].Content : null;
			switch ( text[0] )
			{
				case '{' or '[':
				{
					var flow = new StringBuilder( text );
					var scan = new FlowScan();
					scan.Feed( text );
					while ( !scan.Balanced && More() is string next ) { flow.Append( ' ' ).Append( next ); scan.Feed( " " + next ); }
					var position = 0;
					return Flow( flow.ToString(), ref position );
				}
				case '"' or '\'':
				{
					var quoted = new StringBuilder( text );
					var scan = new QuoteScan( text[0] );
					scan.Feed( text, 1 );
					while ( !scan.Closed && More() is string next ) { quoted.Append( '\n' ).Append( next ); scan.Feed( "\n" + next, 0 ); }
					var end = 0;
					return YamlNode.Scalar( Quoted( quoted.ToString(), ref end ) );
				}
				case '|' or '>':
					var block = new List<string>();
					while ( More() is string next ) block.Add( next );
					return YamlNode.Scalar( string.Join( text[0] == '|' ? "\n" : " ", block ) );
				default:
					var plain = new StringBuilder( text );
					while ( More() is string next ) plain.Append( ' ' ).Append( next );
					return YamlNode.Scalar( plain.ToString() );
			}
		}
	}

	/// <summary>"key: value" or "key:" outside quotes and flow collections; null for any other line.</summary>
	static (string Key, string Value)? SplitKey( string content )
	{
		if ( content[0] is '{' or '[' ) return null;
		if ( content[0] is '"' or '\'' )
		{
			var end = 0;
			var key = Quoted( content, ref end );
			if ( end < content.Length && content[end] == ':' && (end + 1 == content.Length || content[end + 1] == ' ') ) return (key, content[(end + 1)..].Trim());
			return null;
		}
		var colon = content.IndexOf( ": ", StringComparison.Ordinal );
		if ( colon > 0 ) return (content[..colon], content[(colon + 2)..].Trim());
		if ( content.EndsWith( ':' ) && content.Length > 1 ) return (content[..^1], "");
		return null;
	}

	const int MaxDepth = 512;

	/// <summary>Whether a flow collection's brackets are closed, fed one piece at a time.</summary>
	sealed class FlowScan
	{
		int depth;
		char quote;
		bool escaped;
		public bool Balanced => depth <= 0 && quote == '\0';
		public void Feed( string text )
		{
			foreach ( var c in text )
			{
				if ( escaped ) { escaped = false; continue; }
				if ( quote != '\0' ) { if ( c == '\\' && quote == '"' ) escaped = true; else if ( c == quote ) quote = '\0'; }
				else if ( c is '"' or '\'' ) quote = c;
				else if ( c is '{' or '[' ) depth++;
				else if ( c is '}' or ']' ) depth--;
			}
		}
	}

	/// <summary>Whether a quoted scalar is closed, fed one piece at a time.</summary>
	sealed class QuoteScan( char quote )
	{
		bool escaped, pendingQuote;
		public bool Closed { get; private set; }
		public void Feed( string text, int start )
		{
			for ( var p = start; p < text.Length && !Closed; p++ )
			{
				var c = text[p];
				if ( escaped ) { escaped = false; continue; }
				// A single quote doubled inside single quotes is a literal quote.
				if ( pendingQuote ) { pendingQuote = false; if ( c == '\'' ) continue; Closed = true; return; }
				if ( quote == '"' && c == '\\' ) escaped = true;
				else if ( c == quote ) { if ( quote == '\'' ) pendingQuote = true; else Closed = true; }
			}
			// Pieces are lines, so a doubled quote never spans two: a quote ending this piece closes the scalar.
			if ( pendingQuote ) Closed = true;
		}
	}

	/// <summary>Reads a quoted scalar starting at position, folding line breaks and decoding escapes.</summary>
	static string Quoted( string text, ref int position )
	{
		var quote = text[position++];
		var result = new StringBuilder();
		while ( position < text.Length )
		{
			var c = text[position++];
			if ( c == quote )
			{
				if ( quote == '\'' && position < text.Length && text[position] == '\'' ) { result.Append( '\'' ); position++; continue; }
				break;
			}
			if ( c == '\n' )
			{
				// A folded line break becomes a space; continuation indentation was already removed.
				while ( result.Length > 0 && result[^1] == ' ' ) result.Length--;
				result.Append( ' ' );
				continue;
			}
			if ( c != '\\' || quote != '"' || position >= text.Length ) { result.Append( c ); continue; }
			var e = text[position++];
			switch ( e )
			{
				case 'n': result.Append( '\n' ); break;
				case 't': result.Append( '\t' ); break;
				case 'r': result.Append( '\r' ); break;
				case '0': result.Append( '\0' ); break;
				case '\n': break; // escaped line break: join without a space
				case 'x' or 'u' or 'U':
					var length = e == 'x' ? 2 : e == 'u' ? 4 : 8;
					if ( position + length <= text.Length && int.TryParse( text.AsSpan( position, length ), NumberStyles.HexNumber, CultureInfo.InvariantCulture, out var code ) )
					{
						if ( code is >= 0 and < 0xD800 or > 0xDFFF and <= 0x10FFFF ) result.Append( char.ConvertFromUtf32( code ) );
						else result.Append( '\uFFFD' );
						position += length;
					}
					break;
				default: result.Append( e ); break; // \\ \" \/ and others
			}
		}
		return result.ToString();
	}

	static YamlNode Flow( string text, ref int p, int depth = 0 )
	{
		if ( depth > MaxDepth ) throw new InvalidDataException( "The YAML document is nested too deeply." );
		void Skip( ref int at ) { while ( at < text.Length && text[at] == ' ' ) at++; }
		Skip( ref p );
		if ( p >= text.Length ) return YamlNode.Scalar( "" );
		if ( text[p] == '{' )
		{
			p++;
			var pairs = new List<KeyValuePair<string, YamlNode>>();
			while ( true )
			{
				Skip( ref p );
				if ( p >= text.Length ) break;
				if ( text[p] == '}' ) { p++; break; }
				if ( text[p] == ',' ) { p++; continue; }
				string key;
				if ( text[p] is '"' or '\'' ) key = Quoted( text, ref p );
				else
				{
					var start = p;
					while ( p < text.Length && !(text[p] == ':' && (p + 1 == text.Length || text[p + 1] is ' ' or ',' or '}')) && text[p] is not (',' or '}') ) p++;
					key = text[start..p].Trim();
				}
				Skip( ref p );
				var value = YamlNode.Scalar( "" );
				if ( p < text.Length && text[p] == ':' ) { p++; value = Flow( text, ref p, depth + 1 ); }
				pairs.Add( new( key, value ) );
			}
			return YamlNode.Map( pairs );
		}
		if ( text[p] == '[' )
		{
			p++;
			var items = new List<YamlNode>();
			while ( true )
			{
				Skip( ref p );
				if ( p >= text.Length ) break;
				// A stray '}' also ends the list (malformed text must not loop forever).
				if ( text[p] is ']' or '}' ) { p++; break; }
				if ( text[p] == ',' ) { p++; continue; }
				var before = p;
				items.Add( Flow( text, ref p, depth + 1 ) );
				if ( p == before ) p++;
			}
			return YamlNode.List( items );
		}
		if ( text[p] is '"' or '\'' ) return YamlNode.Scalar( Quoted( text, ref p ) );
		var begin = p;
		while ( p < text.Length && text[p] is not (',' or '}' or ']') ) p++;
		return YamlNode.Scalar( text[begin..p].Trim() );
	}
}