using System;
using System.Buffers.Binary;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Text;

namespace ImportUnityPackage;

/// <summary>
/// Reads Unity's binary serialization (a SerializedFile), as the editor saves assets in binary mode, such as terrain data:
/// the header, the type trees that describe each class's fields, the object table and the external references. Objects
/// are read through their type tree into plain values: dictionaries for structures, lists for arrays of structures, typed
/// arrays (byte[], short[], float[] and so on) for arrays of numbers, and numbers, booleans and strings. Format versions
/// 12 to 22 (Unity 5 to 6) are supported when the file includes type trees, as editor files do.
/// </summary>
internal sealed class UnitySerializedFile
{
	public sealed class TypeNode
	{
		public string Type { get; init; }
		public string Name { get; init; }
		public bool IsArray { get; init; }
		/// <summary>Unity's "align after reading" meta flag.</summary>
		public bool Align { get; init; }
		public List<TypeNode> Children { get; } = new();
	}
	public sealed record ObjectInfo( long PathId, int ClassId, long Offset, long Size, TypeNode Tree );
	/// <summary>A file referenced by PPtr file index (1-based): the asset GUID as Unity writes it in text files.</summary>
	public sealed record External( string Guid, int Type, string Path );

	public int Version { get; private set; }
	public IReadOnlyList<ObjectInfo> Objects { get; private set; }
	public IReadOnlyList<External> Externals { get; private set; }
	readonly string path;
	bool bigEndian;
	UnitySerializedFile( string path ) { this.path = path; }

	const int MinimumVersion = 12, MaximumVersion = 22;

	/// <summary>Whether the file starts with a SerializedFile header (rather than Unity's text serialization).</summary>
	public static bool Is( string path )
	{
		try
		{
			using var stream = File.OpenRead( path );
			var header = new byte[48];
			var read = stream.Read( header, 0, header.Length );
			return read >= 20 && Header( header, stream.Length ) != null;
		}
		catch ( IOException ) { return false; }
	}

	static (int Version, long MetadataSize, long DataOffset, int Start, bool BigEndian)? Header( byte[] header, long length )
	{
		long metadataSize = BinaryPrimitives.ReadUInt32BigEndian( header ), fileSize = BinaryPrimitives.ReadUInt32BigEndian( header.AsSpan( 4 ) );
		var version = (int)BinaryPrimitives.ReadUInt32BigEndian( header.AsSpan( 8 ) );
		long dataOffset = BinaryPrimitives.ReadUInt32BigEndian( header.AsSpan( 12 ) );
		if ( version < 9 || version > 64 ) return null;
		var bigEndian = header[16] != 0;
		var start = 20;
		if ( version >= 22 )
		{
			metadataSize = BinaryPrimitives.ReadUInt32BigEndian( header.AsSpan( 20 ) );
			fileSize = BinaryPrimitives.ReadInt64BigEndian( header.AsSpan( 24 ) );
			dataOffset = BinaryPrimitives.ReadInt64BigEndian( header.AsSpan( 32 ) );
			start = 48;
		}
		if ( fileSize != length || dataOffset > fileSize || metadataSize <= 0 || start + metadataSize > fileSize ) return null;
		return (version, metadataSize, dataOffset, start, bigEndian);
	}

	/// <summary>Reads the header, type trees, object table and external references; object data is read on demand.</summary>
	public static UnitySerializedFile Read( string path )
	{
		using var stream = File.OpenRead( path );
		var head = new byte[48];
		if ( stream.Read( head, 0, head.Length ) < 20 || Header( head, stream.Length ) is not { } header )
			throw new InvalidDataException( "Not a Unity binary serialized file." );
		if ( header.Version < MinimumVersion || header.Version > MaximumVersion )
			throw new InvalidDataException( $"Unity serialized file format {header.Version} is not supported (supported: {MinimumVersion}-{MaximumVersion})." );
		if ( header.MetadataSize > 64 * 1024 * 1024 ) throw new InvalidDataException( "The serialized file's metadata is too large." );
		var metadata = new byte[header.Start + header.MetadataSize];
		stream.Position = 0;
		stream.ReadExactly( metadata );
		var file = new UnitySerializedFile( path ) { Version = header.Version, bigEndian = header.BigEndian };
		var reader = new Reader( metadata, header.BigEndian ) { Position = header.Start };
		var v = header.Version;
		reader.CString(); // Unity version
		reader.Int32(); // target platform
		var typeTrees = reader.Byte() != 0;
		if ( !typeTrees ) throw new InvalidDataException( "The serialized file has no type trees (a player build file), which are needed to read it." );
		var types = new List<(int ClassId, TypeNode Tree)>();
		var typeCount = reader.Count();
		for ( var i = 0; i < typeCount; i++ )
		{
			var classId = reader.Int32();
			if ( v >= 16 ) reader.Byte(); // stripped
			if ( v >= 17 ) reader.Int16(); // script type index
			if ( v >= 16 ? classId == 114 : classId < 0 ) reader.Skip( 16 ); // script ID
			reader.Skip( 16 ); // type hash
			var tree = ReadTypeTree( reader, v );
			if ( v >= 21 ) reader.Skip( 4 * reader.Count() ); // type dependencies
			types.Add( (classId, tree) );
		}
		if ( v < 14 ) reader.Int32(); // 64-bit path IDs (format 14 always uses them)
		var objectCount = reader.Count();
		var objects = new List<ObjectInfo>( objectCount );
		for ( var i = 0; i < objectCount; i++ )
		{
			if ( v >= 14 ) reader.Align();
			var pathId = v >= 14 ? reader.Int64() : reader.Int32();
			var offset = (v >= 22 ? reader.Int64() : reader.UInt32()) + header.DataOffset;
			long size = reader.UInt32();
			var typeId = reader.Int32();
			(int ClassId, TypeNode Tree) type;
			if ( v < 16 )
			{
				var classId = reader.UInt16();
				type = types.FirstOrDefault( t => t.ClassId == typeId );
				if ( type.Tree == null ) type = (classId, null);
				if ( v >= 11 ) reader.Int16(); // script type index
				if ( v == 15 ) reader.Byte(); // stripped
			}
			else
			{
				if ( typeId < 0 || typeId >= types.Count ) throw new InvalidDataException( "A serialized object names an unknown type." );
				type = types[typeId];
				if ( v == 16 ) reader.Byte(); // stripped
			}
			if ( offset < 0 || size < 0 || offset + size > stream.Length ) throw new InvalidDataException( "A serialized object lies outside the file." );
			objects.Add( new( pathId, type.ClassId, offset, size, type.Tree ) );
		}
		var scriptCount = reader.Count();
		for ( var i = 0; i < scriptCount; i++ )
		{
			reader.Int32();
			if ( v >= 14 ) { reader.Align(); reader.Int64(); }
			else reader.Int32();
		}
		var externals = new List<External>();
		var externalCount = reader.Count();
		for ( var i = 0; i < externalCount; i++ )
		{
			reader.CString(); // unused
			var guid = reader.Bytes( 16 );
			var type = reader.Int32();
			externals.Add( new( Guid( guid ), type, reader.CString() ) );
		}
		file.Objects = objects;
		file.Externals = externals;
		return file;
	}

	/// <summary>Unity writes each GUID byte low nibble first in text files and .meta files.</summary>
	static string Guid( byte[] bytes ) => string.Concat( bytes.Select( b => "0123456789abcdef"[b & 15].ToString() + "0123456789abcdef"[b >> 4] ) );

	static TypeNode ReadTypeTree( Reader reader, int version )
	{
		var count = reader.Count();
		var stringSize = reader.Count();
		var nodeSize = version >= 19 ? 32 : 24;
		var nodes = reader.Bytes( count * nodeSize );
		var strings = reader.Bytes( stringSize );
		string Text( uint offset )
		{
			if ( (offset & 0x80000000) != 0 ) return CommonStrings.GetValueOrDefault( (int)(offset & 0x7FFFFFFF) ) ?? $"<common {offset & 0x7FFFFFFF}>";
			if ( offset >= strings.Length ) throw new InvalidDataException( "A type tree string lies outside its buffer." );
			var end = System.Array.IndexOf( strings, (byte)0, (int)offset );
			return Encoding.UTF8.GetString( strings, (int)offset, (end < 0 ? strings.Length : end) - (int)offset );
		}
		var node = new Reader( nodes, reader.BigEndian );
		var stack = new List<TypeNode>();
		TypeNode root = null;
		for ( var i = 0; i < count; i++ )
		{
			node.Position = i * nodeSize;
			node.UInt16(); // version
			int level = node.Byte();
			var flags = node.Byte();
			var type = Text( node.UInt32() );
			var name = Text( node.UInt32() );
			node.Int32(); // byte size
			node.Int32(); // index
			var meta = node.Int32();
			var current = new TypeNode { Type = type, Name = name, IsArray = (flags & 1) != 0, Align = (meta & 0x4000) != 0 };
			if ( level == 0 ) { if ( root != null ) throw new InvalidDataException( "A type tree has several roots." ); root = current; }
			else if ( level > stack.Count ) throw new InvalidDataException( "A type tree skips a level." );
			else stack[level - 1].Children.Add( current );
			if ( stack.Count > level ) stack.RemoveRange( level, stack.Count - level );
			stack.Add( current );
		}
		return root ?? throw new InvalidDataException( "A type tree is empty." );
	}

	/// <summary>Reads one object through its type tree.</summary>
	public Dictionary<string, object> ReadObject( ObjectInfo info )
	{
		if ( info.Tree == null ) throw new InvalidDataException( $"Object {info.PathId} has no type tree." );
		if ( info.Size > int.MaxValue ) throw new InvalidDataException( $"Object {info.PathId} is too large." );
		var data = new byte[info.Size];
		using ( var stream = File.OpenRead( path ) )
		{
			stream.Position = info.Offset;
			stream.ReadExactly( data );
		}
		var reader = new Reader( data, bigEndian );
		return Value( info.Tree, reader ) as Dictionary<string, object> ?? throw new InvalidDataException( $"Object {info.PathId} is not a structure." );
	}

	public ObjectInfo Find( long pathId ) => Objects.FirstOrDefault( o => o.PathId == pathId );

	static object Value( TypeNode node, Reader reader, int depth = 0 )
	{
		// Unity's type trees nest a few dozen levels; a deeper one is corrupt or crafted and stops with an error.
		if ( depth > 256 ) throw new InvalidDataException( "The serialized type tree is nested too deeply." );
		var align = node.Align;
		object value;
		switch ( node.Type )
		{
			case "SInt8": value = (long)(sbyte)reader.Byte(); break;
			case "UInt8" or "char": value = (long)reader.Byte(); break;
			case "bool": value = reader.Byte() != 0; break;
			case "SInt16" or "short": value = (long)reader.Int16(); break;
			case "UInt16" or "unsigned short": value = (long)reader.UInt16(); break;
			case "SInt32" or "int" or "Type*": value = (long)reader.Int32(); break;
			case "UInt32" or "unsigned int": value = (long)reader.UInt32(); break;
			case "SInt64" or "long long" or "FileSize": value = reader.Int64(); break;
			case "UInt64" or "unsigned long long": value = (long)reader.UInt64(); break;
			case "float": value = (double)reader.Single(); break;
			case "double": value = reader.Double(); break;
			case "string":
				value = Encoding.UTF8.GetString( reader.Bytes( reader.Count() ) );
				align |= node.Children.Count > 0 && node.Children[0].Align;
				break;
			case "TypelessData":
				value = reader.Bytes( reader.Count() );
				break;
			default:
				if ( node.Children.Count > 0 && node.Children[0].IsArray )
				{
					// vector, map, set and similar containers: one Array child holding size and data.
					var array = node.Children[0];
					align |= array.Align;
					value = ReadArray( array, reader, depth );
				}
				else if ( node.IsArray ) value = ReadArray( node, reader, depth );
				else
				{
					var fields = new Dictionary<string, object>( node.Children.Count, StringComparer.Ordinal );
					foreach ( var child in node.Children ) fields[child.Name] = Value( child, reader, depth + 1 );
					value = fields;
				}
				break;
		}
		if ( align ) reader.Align();
		return value;
	}

	static object ReadArray( TypeNode array, Reader reader, int depth )
	{
		var count = reader.Count();
		if ( array.Children.Count < 2 ) throw new InvalidDataException( "A type tree array has no element type." );
		var element = array.Children[1];
		// Arrays of numbers are read in bulk into typed arrays.
		if ( element.Children.Count == 0 && !element.Align )
		{
			switch ( element.Type )
			{
				case "UInt8" or "char" or "bool" or "SInt8": return reader.Bytes( count );
				case "SInt16" or "short": return reader.Array<short>( count, 2 );
				case "UInt16" or "unsigned short": return reader.Array<ushort>( count, 2 );
				case "SInt32" or "int": return reader.Array<int>( count, 4 );
				case "UInt32" or "unsigned int": return reader.Array<uint>( count, 4 );
				case "float": return reader.Array<float>( count, 4 );
				case "SInt64" or "long long": return reader.Array<long>( count, 8 );
				case "UInt64" or "unsigned long long": return reader.Array<ulong>( count, 8 );
				case "double": return reader.Array<double>( count, 8 );
			}
		}
		var items = new List<object>( Math.Min( count, 1 << 20 ) );
		for ( var i = 0; i < count; i++ ) items.Add( Value( element, reader, depth + 1 ) );
		return items;
	}

	sealed class Reader
	{
		readonly byte[] data;
		public readonly bool BigEndian;
		public int Position;
		public Reader( byte[] data, bool bigEndian ) { this.data = data; BigEndian = bigEndian; }
		ReadOnlySpan<byte> Take( int count )
		{
			if ( count < 0 || Position + count > data.Length ) throw new InvalidDataException( "Unexpected end of serialized data." );
			var span = data.AsSpan( Position, count );
			Position += count;
			return span;
		}
		public byte Byte() => Take( 1 )[0];
		public short Int16() => BigEndian ? BinaryPrimitives.ReadInt16BigEndian( Take( 2 ) ) : BinaryPrimitives.ReadInt16LittleEndian( Take( 2 ) );
		public ushort UInt16() => BigEndian ? BinaryPrimitives.ReadUInt16BigEndian( Take( 2 ) ) : BinaryPrimitives.ReadUInt16LittleEndian( Take( 2 ) );
		public int Int32() => BigEndian ? BinaryPrimitives.ReadInt32BigEndian( Take( 4 ) ) : BinaryPrimitives.ReadInt32LittleEndian( Take( 4 ) );
		public uint UInt32() => BigEndian ? BinaryPrimitives.ReadUInt32BigEndian( Take( 4 ) ) : BinaryPrimitives.ReadUInt32LittleEndian( Take( 4 ) );
		public long Int64() => BigEndian ? BinaryPrimitives.ReadInt64BigEndian( Take( 8 ) ) : BinaryPrimitives.ReadInt64LittleEndian( Take( 8 ) );
		public ulong UInt64() => BigEndian ? BinaryPrimitives.ReadUInt64BigEndian( Take( 8 ) ) : BinaryPrimitives.ReadUInt64LittleEndian( Take( 8 ) );
		public float Single() => BigEndian ? BinaryPrimitives.ReadSingleBigEndian( Take( 4 ) ) : BinaryPrimitives.ReadSingleLittleEndian( Take( 4 ) );
		public double Double() => BigEndian ? BinaryPrimitives.ReadDoubleBigEndian( Take( 8 ) ) : BinaryPrimitives.ReadDoubleLittleEndian( Take( 8 ) );
		/// <summary>A non-negative 32-bit count or length.</summary>
		public int Count()
		{
			var count = Int32();
			if ( count < 0 || count > data.Length - Position ) throw new InvalidDataException( "A serialized count is out of range." );
			return count;
		}
		public byte[] Bytes( int count ) => Take( count ).ToArray();
		public void Skip( int count ) => Take( count );
		public void Align() => Position = Math.Min( data.Length, (Position + 3) & ~3 );
		public string CString()
		{
			var end = System.Array.IndexOf( data, (byte)0, Position );
			if ( end < 0 ) throw new InvalidDataException( "Unterminated string in serialized data." );
			var text = Encoding.UTF8.GetString( data, Position, end - Position );
			Position = end + 1;
			return text;
		}
		public T[] Array<T>( int count, int size ) where T : struct
		{
			var bytes = Take( checked(count * size) );
			var result = new T[count];
			var target = System.Runtime.InteropServices.MemoryMarshal.AsBytes( result.AsSpan() );
			bytes.CopyTo( target );
			if ( BigEndian && size > 1 )
				for ( var i = 0; i < target.Length; i += size ) target.Slice( i, size ).Reverse();
			return result;
		}
	}

	/// <summary>Unity's built-in type tree strings, addressed by offset (with the high bit set) from every type tree.</summary>
	static readonly Dictionary<int, string> CommonStrings = BuildCommonStrings(
		"AABB|AnimationClip|AnimationCurve|AnimationState|Array|Base|BitField|bitset|bool|char|ColorRGBA|Component|" +
		"data|deque|double|dynamic_array|FastPropertyName|first|float|Font|GameObject|Generic Mono|GradientNEW|GUID|" +
		"GUIStyle|int|list|long long|map|Matrix4x4f|MdFour|MonoBehaviour|MonoScript|m_ByteSize|m_Curve|" +
		"m_EditorClassIdentifier|m_EditorHideFlags|m_Enabled|m_ExtensionPtr|m_GameObject|m_Index|m_IsArray|m_IsStatic|" +
		"m_MetaFlag|m_Name|m_ObjectHideFlags|m_PrefabInternal|m_PrefabParentObject|m_Script|m_StaticEditorFlags|m_Type|" +
		"m_Version|Object|pair|PPtr<Component>|PPtr<GameObject>|PPtr<Material>|PPtr<MonoBehaviour>|PPtr<MonoScript>|" +
		"PPtr<Object>|PPtr<Prefab>|PPtr<Sprite>|PPtr<TextAsset>|PPtr<Texture>|PPtr<Texture2D>|PPtr<Transform>|Prefab|" +
		"Quaternionf|Rectf|RectInt|RectOffset|second|set|short|size|SInt16|SInt32|SInt64|SInt8|staticvector|string|" +
		"TextAsset|TextMesh|Texture|Texture2D|Transform|TypelessData|UInt16|UInt32|UInt64|UInt8|unsigned int|" +
		"unsigned long long|unsigned short|vector|Vector2f|Vector3f|Vector4f|m_ScriptingClassIdentifier|Gradient|Type*|" +
		"int2_storage|int3_storage|BoundsInt|m_CorrespondingSourceObject|m_PrefabInstance|m_PrefabAsset|FileSize|" +
		"Hash128|RenderingLayerMask|fixed_array|EntityId|LoadableObjectId|LoadableSceneId" );

	static Dictionary<int, string> BuildCommonStrings( string names )
	{
		var result = new Dictionary<int, string>();
		var offset = 0;
		foreach ( var name in names.Split( '|' ) ) { result[offset] = name; offset += Encoding.UTF8.GetByteCount( name ) + 1; }
		return result;
	}
}