Editor/Core/UnityImport.cs
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Text.Json;
using System.Text.RegularExpressions;
using System.Threading;

namespace ImportUnityPackage;

public record ImportOptions( bool Materials, bool Models, bool Vmat, bool Tmat, bool Vmdl )
{
	public bool Automatic { get; init; }
	public static ImportOptions Auto( bool additionalTerrainMaterials = false ) => new( true, true, true, additionalTerrainMaterials, true ) { Automatic = true };
}
public record ImportIssue( string Asset, string Code, string Message )
{
	// Compatibility fallbacks are warnings; failing to produce a selected resource is an error.
	public string Severity => Code == "conversion-skipped" ? "Error" : "Warning";
}
public record ImportResult( string Directory, int AssetCount, int ConvertedCount, string[] Files, string[] Warnings )
{
	public ImportIssue[] Unresolved { get; init; } = Array.Empty<ImportIssue>();
	public string ReportPath { get; init; }
	public string[] ChangedFiles { get; init; } = Array.Empty<string>();
}
public delegate byte[] ExtractTextureChannel( string source, int channel, double scale, bool invert, CancellationToken cancel );

public static class UnityImport
{
	public static bool NeedsTextureResource( string path ) => Path.GetExtension( path ).ToLowerInvariant() is ".exr" or ".psd" or ".tif" or ".tiff";
	public static bool IsEnabled( UnityAsset asset, ImportOptions options ) => asset.Kind switch
	{
		UnityAssetKind.Material or UnityAssetKind.Texture => options.Materials,
		UnityAssetKind.Model or UnityAssetKind.ModelSupport => options.Models,
		_ => false
	};

	public static HashSet<UnityAsset> Selection( UnityArchive archive, ImportOptions options ) =>
		ImportCatalog.Read( archive ).CreatePlan( options ).Assets.Select( a => a.Asset ).ToHashSet();

	internal static string ReadText( string path )
	{
		if ( new FileInfo( path ).Length > 16 * 1024 * 1024 ) throw new InvalidDataException( "Material or metadata exceeds the 16 MiB text limit." );
		return File.ReadAllText( path );
	}

	public static ImportResult Run( UnityArchive archive, string assetsDirectory, ImportOptions options,
		IProgress<ImportProgress> progress, CancellationToken cancel, ExtractTextureChannel extractChannel = null )
		=> Run( ImportCatalog.Read( archive, cancel ).CreatePlan( options ), assetsDirectory, progress, cancel, extractChannel );

	public static ImportResult Run( ImportPlan plan, string assetsDirectory,
		IProgress<ImportProgress> progress, CancellationToken cancel, ExtractTextureChannel extractChannel = null )
	{
		using var merge = ImportMergePlan.Create( plan, assetsDirectory, cancel );
		merge.Prepare( progress, cancel, extractChannel );
		return merge.Commit( progress, cancel );
	}

	internal static ImportResult PrepareFiles( ImportPlan plan, string assetsDirectory,
		IProgress<ImportProgress> progress, CancellationToken cancel, ExtractTextureChannel extractChannel )
	{
		cancel.ThrowIfCancellationRequested();
		var archive = plan.Archive;
		var options = plan.Options;
		var selected = plan.Assets.Select( a => a.Asset ).ToArray();
		if ( selected.Length == 0 ) throw new InvalidOperationException( "Select at least one supported asset." );
		var destination = Path.Combine( Path.GetFullPath( assetsDirectory ), "Imported" );
		CheckDirectory( destination );
		var stage = Path.Combine( Path.GetDirectoryName( Path.GetFullPath( assetsDirectory ) ), ".unity-import-" + Guid.NewGuid().ToString( "N" ) );
		const string prefix = "Imported/";
		var warnings = new List<string>( archive.Warnings );
		var unresolved = new List<ImportIssue>( plan.Issues );
		warnings.AddRange( plan.Issues.Select( i => $"{i.Asset}: {i.Message}" ) );
		void Issue( string asset, string code, string message )
		{
			unresolved.Add( new( asset, code, message ) );
			warnings.Add( $"{asset}: {message}" );
		}
		var files = new List<string>();
		var outputs = new HashSet<string>( StringComparer.OrdinalIgnoreCase );
		var textures = selected.Where( a => a.Kind == UnityAssetKind.Texture ).ToDictionary( a => a.Guid, a => prefix + a.Path, StringComparer.OrdinalIgnoreCase );
		var textureSources = selected.Where( a => a.Kind == UnityAssetKind.Texture ).ToDictionary( a => a.Guid, StringComparer.OrdinalIgnoreCase );
		foreach ( var (alias, guid) in archive.GuidAliases )
		{
			if ( textures.TryGetValue( guid, out var texture ) ) textures[alias] = texture;
			if ( textureSources.TryGetValue( guid, out var source ) ) textureSources[alias] = source;
		}
		var sourceHashes = new Dictionary<string, string>();
		string SourceHash( UnityAsset a )
		{
			if ( !sourceHashes.TryGetValue( a.Source, out var hash ) ) sourceHashes[a.Source] = hash = ImportMergePlan.HashFile( a.Source, cancel );
			return hash;
		}
		var generatedChannels = new Dictionary<string, string>();
		var shaders = archive.Assets.Where( a => a.Path.EndsWith( ".shader", StringComparison.OrdinalIgnoreCase ) ).ToDictionary( a => a.Guid, StringComparer.OrdinalIgnoreCase );
		var shaderText = new Dictionary<string, string>( StringComparer.OrdinalIgnoreCase );
		var convertedMaterials = new Dictionary<string, UnityMaterial>( StringComparer.OrdinalIgnoreCase );
		var conversions = 0;
		var operation = "creating the staging folder";
		string currentAsset = null;
		Exception failure = null;
		var preserveStage = false;
		try
		{
			Directory.CreateDirectory( stage );
			string Output( string relative )
			{
				if ( !outputs.Add( relative ) ) throw new InvalidDataException( $"Two assets generate the same output: {relative}" );
				var full = Path.GetFullPath( Path.Combine( stage, relative ) );
				if ( !full.StartsWith( stage + Path.DirectorySeparatorChar, StringComparison.OrdinalIgnoreCase ) ) throw new InvalidDataException( "Invalid output path." );
				Directory.CreateDirectory( Path.GetDirectoryName( full ) );
				files.Add( relative );
				return full;
			}
			for ( var i = 0; i < selected.Length; i++ )
			{
				cancel.ThrowIfCancellationRequested();
				var asset = selected[i];
				operation = "writing imported files";
				currentAsset = asset.Path;
				var planned = plan.Find( asset );
				progress?.Report( new( 0.95 * i / selected.Length, $"Importing {asset.Path}" ) );
				using ( var source = File.OpenRead( asset.Source ) )
				using ( var target = new FileStream( Output( asset.Path ), FileMode.CreateNew ) )
						UnityArchive.Copy( source, target, cancel );
				if ( asset.Kind == UnityAssetKind.Texture && NeedsTextureResource( asset.Path ) )
				{
					// Keep source pixels/HDR data intact and give the editor a loadable texture resource.
					File.WriteAllText( Output( asset.Path + ".vtex" ), JsonSerializer.Serialize( new
					{
						Images = new[] { prefix + asset.Path }, InputColorSpace = "Linear", OutputColorSpace = "Linear",
						OutputFormat = asset.Path.EndsWith( ".exr", StringComparison.OrdinalIgnoreCase ) ? "RGBA16161616F" : "BC7",
						OutputMipAlgorithm = "Box", OutputTypeString = "2D"
					}, new JsonSerializerOptions { WriteIndented = true } ) );
					conversions++;
				}
				if ( planned.Vmat || planned.Tmat )
				{
					try
					{
						var material = UnityMaterial.Parse( ReadText( asset.Source ) );
						if ( material.ShaderGuid != null && shaders.TryGetValue( material.ShaderGuid, out var shader ) )
						{
							if ( !shaderText.TryGetValue( shader.Guid, out var source ) ) shaderText[shader.Guid] = source = ReadText( shader.Source );
							material.ConfigureShader( source );
						}
						if ( extractChannel != null )
						{
							material.PrepareChannels( (guid, channel, scale, invert) =>
							{
								if ( !textureSources.TryGetValue( guid, out var source ) )
								{
									Issue( asset.Path, "missing-texture", $"Missing or unsupported texture {guid}." );
									return null;
								}
								var key = $"{guid}_{channel}_{scale.ToString( System.Globalization.CultureInfo.InvariantCulture )}_{invert}";
								if ( !generatedChannels.TryGetValue( key, out var generated ) )
								{
									generated = $"_unity_generated/channels/{SourceHash( source )}_{key}.png";
									File.WriteAllBytes( Output( generated ), extractChannel( source.Source, channel, scale, invert, cancel ) );
									generatedChannels[key] = generated;
								}
								return prefix + generated;
							} );
						}
						string Resolve( string guid )
						{
							if ( textures.TryGetValue( guid, out var texture ) ) return texture;
							Issue( asset.Path, "missing-texture", $"Missing or unsupported texture {guid}." );
							return null;
						}
						if ( planned.Vmat ) { File.WriteAllText( Output( Path.ChangeExtension( asset.Path, ".vmat" ) ), material.ToVmat( Resolve ) ); conversions++; }
						convertedMaterials[asset.Guid] = material;
						if ( planned.Tmat ) { File.WriteAllText( Output( Path.ChangeExtension( asset.Path, ".tmat" ) ), material.ToTmat( Resolve ) ); conversions++; }
						warnings.AddRange( material.Warnings.Select( w => $"{asset.Path}: {w}" ) );
					}
					catch ( InvalidDataException ex ) { Issue( asset.Path, "conversion-skipped", $"Conversion skipped. {ex.Message}" ); }
				}
				if ( planned.Vmdl )
				{
					var meshPath = prefix + asset.Path;
					if ( asset.Path.EndsWith( ".fbx", StringComparison.OrdinalIgnoreCase ) && UnityFbxCompatibility.Normalize( asset.Source, cancel ) is byte[] normalized )
					{
						var repaired = $"_unity_generated/models/{SourceHash( asset )}_{asset.Guid}.fbx";
						File.WriteAllBytes( Output( repaired ), normalized );
						meshPath = prefix + repaired;
						warnings.Add( $"{asset.Path}: removed an exact duplicate vertex array after a closing brace in a derived FBX; original source preserved." );
					}
					var info = asset.ModelInfo;
					for ( var reference = 0; reference < info.PrefabMaterials.Count; reference++ )
						if ( archive.GuidAliases.TryGetValue( info.PrefabMaterials[reference], out var canonical ) ) info.PrefabMaterials[reference] = canonical;
					var scale = asset.Path.EndsWith( ".fbx", StringComparison.OrdinalIgnoreCase ) ? info.ImportScale( asset.Metadata == null ? null : ReadText( asset.Metadata ) ) : 1;
					var remaps = files.Where( f => f.EndsWith( ".vmat", StringComparison.OrdinalIgnoreCase ) )
						.GroupBy( Path.GetFileNameWithoutExtension, StringComparer.OrdinalIgnoreCase )
						.Where( g => g.Count() == 1 ).ToDictionary( g => g.Key.ToLowerInvariant() + ".vmat", g => prefix + g.Single(), StringComparer.OrdinalIgnoreCase );
					var meshName = Path.GetFileNameWithoutExtension( asset.Path ).ToLowerInvariant();
					if ( remaps.TryGetValue( meshName + ".vmat", out var matchingMaterial ) )
					{
						// Some exporters append an LOD suffix to slots that share the model's base material.
						for ( var lod = 0; lod <= 8; lod++ ) remaps.TryAdd( $"{meshName}_lod{lod}.vmat", matchingMaterial );
					}
					var prefabMaterials = selected.Where( a => info.PrefabMaterials.Contains( a.Guid, StringComparer.OrdinalIgnoreCase ) &&
						outputs.Contains( Path.ChangeExtension( a.Path, ".vmat" ) ) ).ToArray();
					foreach ( var slot in info.Materials )
					{
						var semantic = prefabMaterials.Where( a => UnityModel.Normalize( Path.GetFileNameWithoutExtension( a.Path ) ) == UnityModel.Normalize( slot ) ).ToArray();
						if ( semantic.Length == 1 ) remaps[slot.ToLowerInvariant() + ".vmat"] = prefix + Path.ChangeExtension( semantic[0].Path, ".vmat" );
					}
					foreach ( var (slot, albedoFile) in info.MaterialAlbedoFiles )
					{
						if ( remaps.ContainsKey( slot.ToLowerInvariant() + ".vmat" ) ) continue;
						var candidates = (prefabMaterials.Length > 0 ? prefabMaterials : selected.Where( a => a.Kind == UnityAssetKind.Material ).ToArray())
							.Where( a => convertedMaterials.TryGetValue( a.Guid, out var material ) && material.ColorTextureGuids.Any( guid =>
								textureSources.TryGetValue( guid, out var texture ) && Path.GetFileNameWithoutExtension( texture.Path ).Equals(
									Path.GetFileNameWithoutExtension( albedoFile ), StringComparison.OrdinalIgnoreCase ) ) ).ToArray();
						// Unity often keeps the old diffuse filename as the material name
						// after repacking its texture. Require a unique prefab-assigned material.
						if ( candidates.Length == 0 ) candidates = prefabMaterials.Where( a =>
							Path.GetFileNameWithoutExtension( a.Path ).Equals( Path.GetFileNameWithoutExtension( albedoFile ), StringComparison.OrdinalIgnoreCase ) ).ToArray();
						if ( candidates.Length == 1 && outputs.Contains( Path.ChangeExtension( candidates[0].Path, ".vmat" ) ) )
							remaps[slot.ToLowerInvariant() + ".vmat"] = prefix + Path.ChangeExtension( candidates[0].Path, ".vmat" );
					}
					var defaultMaterial = prefabMaterials.Length == 1 ? prefix + Path.ChangeExtension( prefabMaterials[0].Path, ".vmat" ) : null;
					if ( options.Materials && options.Vmat && defaultMaterial == null && info.Materials.Count == 0 && asset.Path.EndsWith( ".fbx", StringComparison.OrdinalIgnoreCase ) )
					{
						const string neutralPath = "_unity_generated/defaults/unassigned.vmat";
						if ( !outputs.Contains( neutralPath ) )
						{
							var neutral = new UnityMaterial();
							neutral.Colors["_Color"] = new[] { 0.5, 0.5, 0.5, 1.0 };
							File.WriteAllText( Output( neutralPath ), neutral.ToVmat( _ => null ) );
							conversions++;
						}
						defaultMaterial = prefix + neutralPath;
						Issue( asset.Path, "unassigned-material", "FBX contains no named material slots and no unambiguous prefab assignment was resolved; using a neutral material." );
					}
					if ( options.Materials && options.Vmat && defaultMaterial == null )
						foreach ( var slot in info.Materials.Where( s => !remaps.ContainsKey( s.ToLowerInvariant() + ".vmat" ) ) )
							Issue( asset.Path, "unassigned-material", $"No converted material was resolved for slot '{slot}'." );
					File.WriteAllText( Output( Path.ChangeExtension( asset.Path, ".vmdl" ) ), ModelDocument( meshPath, remaps, defaultMaterial, info.HighestDetailMeshes, scale ) );
					conversions++;
					warnings.Add( $"{asset.Path}: review scale, orientation, material assignments, collision and animations in ModelDoc." );
				}
			}
			var unsupported = archive.Assets.Count( a => a.Kind == UnityAssetKind.Unsupported );
			if ( unsupported > 0 ) warnings.Add( $"{unsupported} unsupported files (such as scripts, scenes and prefabs) were excluded." );
			var issues = unresolved.Distinct().ToArray();
			var report = new
			{
				Package = Path.GetFileName( archive.FileName ), Assets = selected.Length, Converted = conversions,
				Files = files.ToArray(), Unresolved = issues,
				Selection = plan.Assets.Select( a => new { Source = a.Asset.Path, a.Explicit, a.OutputLabel,
					RequiredBy = a.RequiredBy.Select( d => new { Source = d.Asset.Path, d.Reason } ).ToArray() } ).ToArray(),
				MaterialConversions = selected.Where( a => convertedMaterials.ContainsKey( a.Guid ) ).Select( a => new
				{
					Source = a.Path, convertedMaterials[a.Guid].ShaderName, convertedMaterials[a.Guid].ShaderGuid,
					Conversion = "Common material properties; shader programs are not translated",
					Vmat = plan.Find( a ).Vmat, Tmat = plan.Find( a ).Tmat
				} ).ToArray(),
				Warnings = warnings.Distinct().ToArray()
			};
			operation = "writing the import report";
			currentAsset = null;
			File.WriteAllText( Output( "unity-import-report.json" ), JsonSerializer.Serialize( report, new JsonSerializerOptions { WriteIndented = true } ) );
			cancel.ThrowIfCancellationRequested();
			var prepared = new ImportResult( stage, selected.Length, conversions, files.Select( f => Path.Combine( stage, f ) ).ToArray(), warnings.Distinct().ToArray() ) { Unresolved = issues };
			preserveStage = true;

			return prepared;
		}
		catch ( Exception ex )
		{
			failure = ex;
			if ( ex is IOException or UnauthorizedAccessException )
			{
				failure = new IOException( $"Import failed while {operation}" + (currentAsset == null ? "" : $" for '{currentAsset}'") +
					$". Staging folder: '{stage}'. Destination: '{destination}'. Windows error {ex.HResult & 0xffff}: {ex.Message}", ex );
				throw failure;
			}
			throw;
		}
		finally
		{
			if ( !preserveStage && Directory.Exists( stage ) )
			{
				try { ImportStorage.Retry( () => Directory.Delete( stage, true ), CancellationToken.None ); }
				catch ( Exception cleanup ) when ( cleanup is IOException or UnauthorizedAccessException )
				{
					// Preserve the original conversion/cancellation error if cleanup also fails.
					if ( failure == null ) throw new IOException( $"Could not remove staging folder '{stage}'. {cleanup.Message}", cleanup );
					failure.Data["StagingCleanupError"] = $"Could not remove '{stage}': {cleanup}";
				}
			}
		}
	}

	internal static void CheckDirectory( string path )
	{
		for ( var current = new DirectoryInfo( Path.GetFullPath( path ) ); current != null; current = current.Parent )
			if ( (Directory.Exists( current.FullName ) || File.Exists( current.FullName )) && File.GetAttributes( current.FullName ).HasFlag( FileAttributes.ReparsePoint ) )
				throw new IOException( $"Import destination cannot pass through a symbolic link: {current.FullName}" );
	}

	static string ModelDocument( string mesh, Dictionary<string, string> remaps, string defaultMaterial, string[] highestDetailMeshes, double scale ) => $$"""
		<!-- kv3 encoding:text:version{e21c7f3c-8a33-41c5-9977-a76d3a32aa0d} format:modeldoc29:version{3cec427c-1b0e-4d48-a90a-0436f33a6041} -->
		{
			rootNode =
			{
				_class = "RootNode"
				children =
				[
					{
						_class = "MaterialGroupList"
						children =
						[
							{
								_class = "DefaultMaterialGroup"
								remaps = [ {{string.Join( ", ", remaps.Select( r => "{ from = " + UnityMaterial.Quote( r.Key ) + " to = " + UnityMaterial.Quote( r.Value ) + " }" ) )}} ]
								use_global_default = {{(defaultMaterial != null ? "true" : "false")}}
								global_default_material = {{UnityMaterial.Quote( defaultMaterial ?? "" )}}
							}
						]
					},
					{
						_class = "RenderMeshList"
						children =
						[
							{
								_class = "RenderMeshFile"
								filename = {{UnityMaterial.Quote( mesh )}}
								import_filter =
								{
									exclude_by_default = {{(highestDetailMeshes.Length > 0 ? "true" : "false")}}
									exception_list = [ {{string.Join( ", ", highestDetailMeshes.Select( UnityMaterial.Quote ) )}} ]
								}
								import_scale = {{scale.ToString( "0.#########", System.Globalization.CultureInfo.InvariantCulture )}}
								import_translation = [ 0.0, 0.0, 0.0 ]
								import_rotation = [ 0.0, 0.0, 0.0 ]
							}
						]
					}
				]
				model_archetype = ""
			}
		}
		""";
}