using System;
using System.Collections.Generic;
using System.Globalization;
using System.IO;
using System.IO.Compression;
using System.Linq;
using System.Security.Cryptography;
using System.Text;
using System.Text.Json;
using System.Text.Json.Nodes;

namespace ImportUnityPackage;

/// <summary>
/// A Unity TerrainData asset, read from Unity's binary serialization: the heightmap, holes, splat (paint) maps and their
/// terrain layers, tree prototypes and tree instances. Detail (grass) layers are counted but not read.
/// </summary>
internal sealed class UnityTerrainData
{
	/// <summary>Height samples per side; samples span the whole terrain, edge to edge.</summary>
	public int Resolution { get; private set; }
	/// <summary>Resolution² heights from 0 to <see cref="MaxHeight"/>, row by row along Unity Z, each row along X.</summary>
	public short[] Heights { get; private set; }
	/// <summary>Terrain size in Unity meters: X width, Y height range, Z length.</summary>
	public double[] Size { get; private set; }
	/// <summary>(Resolution - 1)² cells, row by row along Z; 0 is a hole. Null when the terrain has none.</summary>
	public byte[] Holes { get; private set; }
	public int AlphamapResolution { get; private set; }
	/// <summary>RGBA weights for layers 4n to 4n + 3, rows along Z from Unity's bottom row.</summary>
	public List<byte[]> Alphamaps { get; } = new();
	public List<UnityReference> Layers { get; } = new();
	public List<UnityReference> TreePrototypes { get; } = new();
	public List<Tree> Trees { get; } = new();
	public int DetailPrototypes { get; private set; }
	public List<string> Notes { get; } = new();

	/// <summary>A painted tree: position normalized to the terrain size, scales, and rotation about Y in radians.</summary>
	public readonly record struct Tree( double X, double Y, double Z, double Width, double Height, double Rotation, int Prototype );

	public const int MaxHeight = 32766;

	/// <summary>Drops the height, hole and splat maps once converted; placements only need sizes and trees.</summary>
	public void ReleaseMaps()
	{
		Heights = null; Holes = null;
		Alphamaps.Clear();
	}
	public const int ClassId = 156;

	/// <summary>Whether a Unity asset file holds terrain data: a binary file with a TerrainData object, or a text one.</summary>
	public static bool Is( string source )
	{
		try
		{
			if ( UnitySerializedFile.Is( source ) ) return UnitySerializedFile.Read( source ).Objects.Any( o => o.ClassId == ClassId );
			using var reader = new StreamReader( source );
			var buffer = new char[4096];
			var text = new string( buffer, 0, reader.Read( buffer, 0, buffer.Length ) );
			return text.StartsWith( "%YAML", StringComparison.Ordinal ) && text.Contains( "--- !u!156 ", StringComparison.Ordinal );
		}
		catch ( Exception ex ) when ( ex is IOException or InvalidDataException ) { return false; }
	}

	/// <summary>Reads terrain data; <paramref name="maps"/> also reads heights, holes and splat maps.</summary>
	public static UnityTerrainData Read( string source, bool maps = true )
	{
		if ( !UnitySerializedFile.Is( source ) )
			throw new InvalidDataException( "Terrain data saved in Unity's text serialization is not supported; re-save it in Unity with Mixed or binary asset serialization." );
		var file = UnitySerializedFile.Read( source );
		var info = file.Objects.FirstOrDefault( o => o.ClassId == ClassId ) ?? throw new InvalidDataException( "The asset holds no TerrainData." );
		var data = file.ReadObject( info );
		var result = new UnityTerrainData();
		UnityReference Reference( object pointer )
		{
			// Path IDs are 64-bit hashes, so they are read exactly rather than through a double.
			var fileIndex = Field( pointer, "m_FileID" ) is long f ? (int)f : 0;
			var pathId = Field( pointer, "m_PathID" ) is long p ? p : 0;
			if ( fileIndex == 0 ) return new( pathId, null );
			return fileIndex <= file.Externals.Count ? new( pathId, file.Externals[fileIndex - 1].Guid ) : default;
		}

		var heightmap = Field( data, "m_Heightmap" );
		var resolution = (int)(Field( heightmap, "m_Resolution" ) is { } r ? Number( r ) : Number( Field( heightmap, "m_Width" ) ));
		if ( Field( heightmap, "m_Height" ) is { } height && (int)Number( height ) != resolution ) throw new InvalidDataException( "Terrain heightmaps that are not square are not supported." );
		if ( resolution < 2 || resolution > 8193 ) throw new InvalidDataException( $"Terrain heightmap resolution {resolution} is out of range." );
		result.Resolution = resolution;
		var scale = Field( heightmap, "m_Scale" );
		result.Size = new[] { Number( Field( scale, "x" ) ) * (resolution - 1), Number( Field( scale, "y" ) ), Number( Field( scale, "z" ) ) * (resolution - 1) };
		if ( result.Size.Any( s => !(s > 0) || double.IsInfinity( s ) ) ) throw new InvalidDataException( "The terrain size is not positive." );

		var splat = Field( data, "m_SplatDatabase" );
		if ( Field( splat, "m_TerrainLayers" ) is List<object> layers ) result.Layers.AddRange( layers.Select( Reference ) );
		else if ( Field( splat, "m_Splats" ) is List<object> { Count: > 0 } ) result.Notes.Add( "Terrain textures from before Unity 2018.3 (splat prototypes rather than terrain layers) are not converted." );
		result.AlphamapResolution = (int)Number( Field( splat, "m_AlphamapResolution" ) ?? 0L );

		var details = Field( data, "m_DetailDatabase" );
		foreach ( var prototype in Field( details, "m_TreePrototypes" ) as List<object> ?? new() ) result.TreePrototypes.Add( Reference( Field( prototype, "prefab" ) ) );
		foreach ( var tree in Field( details, "m_TreeInstances" ) as List<object> ?? new() )
		{
			var position = Field( tree, "position" );
			result.Trees.Add( new( Number( Field( position, "x" ) ), Number( Field( position, "y" ) ), Number( Field( position, "z" ) ),
				Number( Field( tree, "widthScale" ) ?? 1.0 ), Number( Field( tree, "heightScale" ) ?? 1.0 ), Number( Field( tree, "rotation" ) ?? 0.0 ), (int)Number( Field( tree, "index" ) ?? 0L ) ) );
		}
		result.DetailPrototypes = (Field( details, "m_DetailPrototypes" ) as List<object>)?.Count ?? 0;
		if ( !maps ) return result;

		result.Heights = Field( heightmap, "m_Heights" ) as short[] ?? throw new InvalidDataException( "The terrain has no heights." );
		if ( result.Heights.Length != resolution * resolution ) throw new InvalidDataException( "The terrain heights do not match its resolution." );
		if ( Field( heightmap, "m_Holes" ) is byte[] holes && holes.Length == (resolution - 1) * (resolution - 1) && holes.Any( h => h == 0 ) ) result.Holes = holes;
		foreach ( var pointer in Field( splat, "m_AlphaTextures" ) as List<object> ?? new() )
		{
			var reference = Reference( pointer );
			if ( reference.Guid != null || file.Find( reference.FileId ) is not { } texture ) { result.Notes.Add( "A splat map is stored outside the terrain data; painted layers are incomplete." ); continue; }
			var pixels = Rgba( file.ReadObject( texture ), out var width, out var heightPixels, out var problem );
			if ( pixels == null || width != result.AlphamapResolution || heightPixels != result.AlphamapResolution )
			{
				result.Notes.Add( problem ?? "A splat map does not match the terrain's alphamap resolution; painted layers are incomplete." );
				result.Alphamaps.Add( null );
				continue;
			}
			result.Alphamaps.Add( pixels );
		}
		return result;
	}

	/// <summary>The top mip of an uncompressed Texture2D as RGBA bytes.</summary>
	static byte[] Rgba( Dictionary<string, object> texture, out int width, out int height, out string problem )
	{
		width = (int)Number( Field( texture, "m_Width" ) ); height = (int)Number( Field( texture, "m_Height" ) );
		problem = null;
		var format = (int)Number( Field( texture, "m_TextureFormat" ) );
		var image = Field( texture, "image data" ) as byte[] ?? Array.Empty<byte>();
		var pixels = width * height;
		if ( image.Length == 0 ) { problem = "A splat map's pixels are streamed from a separate file; painted layers are incomplete."; return null; }
		var bytesPerPixel = format switch { 1 => 1, 3 => 3, 4 or 5 or 14 => 4, _ => 0 };
		if ( bytesPerPixel == 0 || image.Length < pixels * bytesPerPixel ) { problem = $"A splat map uses texture format {format}, which is not read; painted layers are incomplete."; return null; }
		var rgba = new byte[pixels * 4];
		for ( var i = 0; i < pixels; i++ )
		{
			var s = i * bytesPerPixel; var d = i * 4;
			switch ( format )
			{
				case 1: rgba[d + 3] = image[s]; break; // Alpha8
				case 3: rgba[d] = image[s]; rgba[d + 1] = image[s + 1]; rgba[d + 2] = image[s + 2]; break; // RGB24
				case 4: Buffer.BlockCopy( image, s, rgba, d, 4 ); break; // RGBA32
				case 5: rgba[d] = image[s + 1]; rgba[d + 1] = image[s + 2]; rgba[d + 2] = image[s + 3]; rgba[d + 3] = image[s]; break; // ARGB32
				case 14: rgba[d] = image[s + 2]; rgba[d + 1] = image[s + 1]; rgba[d + 2] = image[s]; rgba[d + 3] = image[s + 3]; break; // BGRA32
			}
		}
		return rgba;
	}

	static object Field( object value, string name ) => value is Dictionary<string, object> fields ? fields.GetValueOrDefault( name ) : null;
	static double Number( object value ) => value switch { long l => l, double d => d, bool b => b ? 1 : 0, _ => throw new InvalidDataException( "A terrain field is missing or not a number." ) };

	/// <summary>The square s&amp;box terrain covering this terrain, in s&amp;box units.</summary>
	public sealed class Storage
	{
		public int Resolution { get; init; }
		public double Size { get; init; }
		public double Height { get; init; }
		public ushort[] Heights { get; init; }
		public uint[] Control { get; init; }
	}

	/// <summary>
	/// Resamples the terrain onto s&amp;box's square, power-of-two terrain. s&amp;box terrain-local X runs along Unity Z and
	/// local Y along Unity -X, from the Unity terrain's +X edge; each s&amp;box texel holds the surface at its centre, which
	/// the renderer interpolates between, so heights are sampled there from Unity's edge-to-edge samples. A terrain longer
	/// on one side than the other covers the square's remainder with holes. Painted weights keep each texel's two
	/// strongest layers with their blend, as s&amp;box's control map stores them; at most 32 layers can be addressed.
	/// </summary>
	public Storage ToStorage()
	{
		if ( Heights == null ) throw new InvalidOperationException( "Terrain maps were not read." );
		double width = Size[0], length = Size[2];
		var size = Math.Max( width, length );
		var spacing = Math.Min( width, length ) / (Resolution - 1);
		var resolution = 1;
		while ( resolution * 2 <= Math.Round( size / spacing ) && resolution < 4096 ) resolution *= 2;
		resolution = Math.Max( resolution, 32 );
		var cell = size / resolution;
		var heights = new ushort[resolution * resolution];
		var control = new uint[resolution * resolution];
		var layers = Math.Min( 32, Math.Max( Layers.Count, Alphamaps.Count * 4 ) );
		var weights = new int[Math.Max( 1, Alphamaps.Count * 4 )];
		var outside = 0;
		for ( var j = 0; j < resolution; j++ )
		{
			var x = width - (j + 0.5) * cell;
			for ( var i = 0; i < resolution; i++ )
			{
				var z = (i + 0.5) * cell;
				var index = i + j * resolution;
				var inside = x >= 0 && z <= length;
				// Unity sample coordinates of this point.
				var sx = Math.Clamp( x / width * (Resolution - 1), 0, Resolution - 1 );
				var sz = Math.Clamp( z / length * (Resolution - 1), 0, Resolution - 1 );
				heights[index] = (ushort)Math.Clamp( Math.Round( Sample( sx, sz ) / MaxHeight * ushort.MaxValue ), 0, ushort.MaxValue );
				var hole = !inside;
				if ( inside && Holes != null )
				{
					int cx = Math.Min( (int)sx, Resolution - 2 ), cz = Math.Min( (int)sz, Resolution - 2 );
					hole = Holes[cz * (Resolution - 1) + cx] == 0;
				}
				if ( !inside ) outside++;
				control[index] = Material( x / width, z / length, layers, weights ) | (hole ? 1u << 18 : 0);
			}
		}
		if ( outside > 0 ) Notes.Add( $"The terrain is {width:0.##} by {length:0.##} m; s&box terrains are square, so the rest of the {size:0.##} m square is a hole." );
		if ( Layers.Count > 32 ) Notes.Add( $"The terrain has {Layers.Count} layers; s&box terrain addresses 32, so the rest are not painted." );
		return new Storage { Resolution = resolution, Size = size * UnityTransform.InchesPerMeter, Height = Size[1] * UnityTransform.InchesPerMeter, Heights = heights, Control = control };
	}

	double Sample( double x, double z )
	{
		int x0 = (int)x, z0 = (int)z;
		int x1 = Math.Min( x0 + 1, Resolution - 1 ), z1 = Math.Min( z0 + 1, Resolution - 1 );
		double fx = x - x0, fz = z - z0;
		double H( int px, int pz ) => Math.Max( (short)0, Heights[pz * Resolution + px] );
		return (H( x0, z0 ) * (1 - fx) + H( x1, z0 ) * fx) * (1 - fz) + (H( x0, z1 ) * (1 - fx) + H( x1, z1 ) * fx) * fz;
	}

	/// <summary>The control value at a normalized terrain position: the two strongest layers and their blend.</summary>
	uint Material( double u, double v, int layers, int[] weights )
	{
		if ( Alphamaps.Count == 0 || AlphamapResolution <= 0 ) return 0;
		var px = Math.Clamp( (int)Math.Floor( u * AlphamapResolution ), 0, AlphamapResolution - 1 );
		var py = Math.Clamp( (int)Math.Floor( v * AlphamapResolution ), 0, AlphamapResolution - 1 );
		var offset = (py * AlphamapResolution + px) * 4;
		for ( var m = 0; m < Alphamaps.Count; m++ )
			for ( var c = 0; c < 4; c++ ) weights[m * 4 + c] = Alphamaps[m]?[offset + c] ?? 0;
		int first = -1, second = -1;
		for ( var l = 0; l < Math.Min( layers, weights.Length ); l++ )
		{
			if ( first < 0 || weights[l] > weights[first] ) { second = first; first = l; }
			else if ( second < 0 || weights[l] > weights[second] ) second = l;
		}
		if ( first < 0 ) return 0;
		if ( second < 0 || weights[second] == 0 ) second = first;
		var total = weights[first] + (second == first ? 0 : weights[second]);
		var blend = total == 0 || second == first ? 0 : (uint)Math.Clamp( weights[second] * 255 / total, 0, 255 );
		return (uint)first & 31 | ((uint)second & 31) << 5 | blend << 10;
	}

	/// <summary>
	/// The s&amp;box terrain resource: its JSON and the binary sidecar (<c>.terrain_d</c>) holding the deflated height and
	/// control maps, in the engine's blob layout. <paramref name="materials"/> are the TMAT paths of the layers, in order.
	/// <paramref name="heightBlend"/> blends layers by their height maps, as Unity's terrain shaders do only when their
	/// material enables it (<c>_EnableHeightBlend</c>); otherwise layers blend by weight.
	/// </summary>
	public static (string Json, byte[] Blob) TerrainFile( Storage storage, IReadOnlyList<string> materials, bool heightBlend = false )
	{
		var payload = new MemoryStream();
		using ( var writer = new BinaryWriter( payload, Encoding.UTF8, true ) )
		{
			writer.Write( 1 ); // blob version
			writer.Write( 2 ); // maps
			void Map( string name, byte[] raw, byte elementSize )
			{
				writer.Write( name.Length ); writer.Write( Encoding.ASCII.GetBytes( name ) );
				writer.Write( elementSize );
				writer.Write( raw.Length / elementSize );
				var compressed = new MemoryStream();
				using ( var deflate = new DeflateStream( compressed, CompressionLevel.Optimal, true ) ) deflate.Write( raw );
				writer.Write( (int)compressed.Length );
				writer.Write( compressed.GetBuffer(), 0, (int)compressed.Length );
			}
			Map( "heightmap", System.Runtime.InteropServices.MemoryMarshal.AsBytes( storage.Heights.AsSpan() ).ToArray(), 2 );
			Map( "splatmap", System.Runtime.InteropServices.MemoryMarshal.AsBytes( storage.Control.AsSpan() ).ToArray(), 4 );
		}
		var bytes = payload.ToArray();
		var hash = SHA256.HashData( bytes );
		hash[7] = (byte)((hash[7] & 0x0F) | 0x80); hash[8] = (byte)((hash[8] & 0x3F) | 0x80);
		var id = new Guid( hash.AsSpan( 0, 16 ) );
		var blob = new MemoryStream();
		using ( var writer = new BinaryWriter( blob ) )
		{
			const int headerSize = 8 + 32;
			writer.Write( 1 ); // file version
			writer.Write( 1 ); // blobs
			writer.Write( id.ToByteArray() );
			writer.Write( 1 );
			writer.Write( (long)headerSize );
			writer.Write( bytes.Length );
			writer.Write( bytes );
		}
		var json = new JsonObject
		{
			["Maps"] = new JsonObject { ["$blob"] = id.ToString() },
			["Resolution"] = storage.Resolution,
			["TerrainSize"] = Math.Round( storage.Size, 3 ),
			["TerrainHeight"] = Math.Round( storage.Height, 3 ),
			["Materials"] = new JsonArray( materials.Select( m => (JsonNode)JsonValue.Create( m ) ).ToArray() ),
			["MaterialSettings"] = new JsonObject { ["HeightBlendEnabled"] = heightBlend, ["HeightBlendSharpness"] = 0.87 },
			["ResourceVersion"] = 3,
			["__references"] = new JsonArray(),
			["__version"] = 3
		};
		return (json.ToJsonString( new JsonSerializerOptions { WriteIndented = true } ), blob.ToArray());
	}
}