Editor/Core/ImportMerge.cs
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Security.Cryptography;
using System.Text;
using System.Text.Json;
using System.Threading;

namespace ImportUnityPackage;

public sealed class ImportFileChoice
{
	public string Path { get; init; }
	public string Status { get; init; }
	public string Detail { get; init; }
	public bool Conflict { get; init; }
	public bool Replace { get; set; }
	internal string ExistingHash { get; init; }
	internal string IncomingHash { get; init; }
	public string Action => ExistingHash == null ? "Add" : !Conflict ? "Reuse" : Replace ? "Replace" : "Keep existing";
}

public sealed class ImportedFileRecord
{
	public string Hash { get; set; }
	public string Guid { get; set; }
	public string Metadata { get; set; }
	public string SourcePath { get; set; }
	public string Conversion { get; set; }
	public string[] Packages { get; set; } = Array.Empty<string>();
}

/// <summary>A reviewed merge into Assets/Imported. Existing destinations are checked again before writing.</summary>
public sealed class ImportMergePlan : IDisposable
{
	const string Version = "original-folders-v1";
	readonly ImportPlan original;
	readonly string assetsRoot, stateRoot, indexPath, indexHash, scratch;
	readonly Dictionary<string, ImportedFileRecord> index;
	readonly Dictionary<string, UnityAsset> effective = new( StringComparer.OrdinalIgnoreCase );
	readonly Dictionary<string, string> snapshots = new( StringComparer.OrdinalIgnoreCase );
	ImportResult prepared;
	bool preserveRecovery;
	public string Destination { get; }
	public string Package { get; }
	public List<ImportFileChoice> Sources { get; } = new();
	public List<ImportFileChoice> Outputs { get; } = new();
	public string ReportPath { get; private set; }

	ImportMergePlan( ImportPlan plan, string assets )
	{
		original = plan;
		assetsRoot = System.IO.Path.GetFullPath( assets );
		Destination = System.IO.Path.Combine( assetsRoot, "Imported" );
		stateRoot = System.IO.Path.Combine( System.IO.Path.GetDirectoryName( assetsRoot ), ".importunitypackage" );
		indexPath = System.IO.Path.Combine( stateRoot, "index.json" );
		UnityImport.CheckDirectory( Destination );
		UnityImport.CheckDirectory( stateRoot );
		indexHash = HashFile( indexPath );
		index = File.Exists( indexPath )
			? new( JsonSerializer.Deserialize<Dictionary<string, ImportedFileRecord>>( File.ReadAllText( indexPath ) ) ?? throw new InvalidDataException( "Empty import index" ), StringComparer.OrdinalIgnoreCase )
			: new( StringComparer.OrdinalIgnoreCase );
		Package = System.IO.Path.GetFileName( plan.Archive.FileName );
		scratch = System.IO.Path.Combine( System.IO.Path.GetDirectoryName( assetsRoot ), ".unity-merge-" + Guid.NewGuid().ToString( "N" ) );
	}

	public static ImportMergePlan Create( ImportPlan plan, string assets, CancellationToken cancel = default )
	{
		var merge = new ImportMergePlan( plan, assets );
		try
		{
			foreach ( var item in plan.Assets )
			{
				cancel.ThrowIfCancellationRequested();
				var a = item.Asset;
				if ( a.Path.Split( '/' )[0].Equals( "_unity_generated", StringComparison.OrdinalIgnoreCase ) || a.Path.Equals( "unity-import-report.json", StringComparison.OrdinalIgnoreCase ) )
					throw new InvalidDataException( $"Reserved importer output path: {a.Path}" );
				var incoming = HashFile( a.Source, cancel );
				var existing = HashFile( merge.Target( a.Path ), cancel );
				merge.snapshots[a.Path] = existing;
				merge.index.TryGetValue( a.Path, out var record );
				var metadata = a.Metadata == null ? null : UnityImport.ReadText( a.Metadata );
				var sameSettings = record != null && record.Guid == a.Guid && record.Metadata == metadata;
				merge.Sources.Add( merge.Choice( a.Path, incoming, existing, existing != null && (existing != incoming || !sameSettings),
					record != null && record.Guid != a.Guid ? "Different Unity GUID at the same destination." : !sameSettings && existing != null ? "Unity metadata differs or this file has no import record." : "" ) );
				foreach ( var ext in new[] { ".vmat", ".tmat", ".vmdl" } )
				{
					var path = System.IO.Path.ChangeExtension( a.Path, ext );
					merge.snapshots.TryAdd( path, HashFile( merge.Target( path ), cancel ) );
				}
			}
			return merge;
		}
		catch { merge.Dispose(); throw; }
	}

	ImportFileChoice Choice( string path, string incoming, string existing, bool conflict, string detail = "" )
	{
		index.TryGetValue( path, out var record );
		var modified = existing != null && record != null && record.Hash != existing;
		var users = record?.Packages ?? Array.Empty<string>();
		return new()
		{
			Path = path, IncomingHash = incoming, ExistingHash = existing, Conflict = conflict,
			Status = existing == null ? "New" : modified ? "Locally modified" : conflict ? "Conflict" : "Identical",
			Detail = detail + (users.Length > 0 ? " Used by imports: " + string.Join( ", ", users ) + ". Replacing shared files affects their existing users." : "")
		};
	}

	public void Prepare( IProgress<ImportProgress> progress, CancellationToken cancel, ExtractTextureChannel extract = null )
	{
		if ( prepared != null ) throw new InvalidOperationException( "This import is already prepared." );
		CheckUnchanged( cancel );
		Directory.CreateDirectory( scratch );
		var chosen = Sources.ToDictionary( c => c.Path, StringComparer.OrdinalIgnoreCase );
		var copies = new List<UnityAsset>();
		foreach ( var asset in original.Archive.Assets )
		{
			cancel.ThrowIfCancellationRequested();
			if ( !chosen.TryGetValue( asset.Path, out var choice ) || choice.ExistingHash == null || choice.Replace )
			{
				copies.Add( new UnityAsset { Path = asset.Path, Guid = asset.Guid, Source = asset.Source, Metadata = asset.Metadata, Kind = asset.Kind, Selected = original.Find( asset )?.Explicit == true } );
				continue;
			}
			// Snapshot retained source data. Conversion must never use the rejected replacement's pixels or metadata.
			var source = System.IO.Path.Combine( scratch, asset.Guid + ".source" );
			File.Copy( Target( asset.Path ), source );
			if ( HashFile( source, cancel ) != choice.ExistingHash ) throw new IOException( $"Destination changed during review: {asset.Path}. Review the import again." );
			index.TryGetValue( asset.Path, out var record );
			string meta = null;
			if ( record?.Metadata != null ) { meta = source + ".meta"; File.WriteAllText( meta, record.Metadata ); }
			copies.Add( new UnityAsset { Path = asset.Path, Guid = asset.Guid, Source = source, Metadata = meta, Kind = asset.Kind, Selected = original.Find( asset )?.Explicit == true } );
		}
		// Include recorded dependencies needed by retained materials/models, even when the new package omits them.
		var guids = copies.Select( a => a.Guid ).ToHashSet( StringComparer.OrdinalIgnoreCase );
		var paths = copies.Select( a => a.Path ).ToHashSet( StringComparer.OrdinalIgnoreCase );
		var needed = new HashSet<string>( StringComparer.OrdinalIgnoreCase );
		void References( UnityAsset a )
		{
			if ( a.Kind == UnityAssetKind.Material ) needed.UnionWith( UnityMaterial.References( UnityImport.ReadText( a.Source ) ) );
			if ( a.Metadata != null ) needed.UnionWith( UnityMaterial.References( UnityImport.ReadText( a.Metadata ) ) );
		}
		foreach ( var asset in copies.Where( a => chosen.ContainsKey( a.Path ) ) ) References( asset );
		var pending = index.Where( x => x.Value.Guid != null && x.Value.SourcePath == x.Key ).ToList();
		while ( pending.Any( x => needed.Contains( x.Value.Guid ) && !guids.Contains( x.Value.Guid ) && !paths.Contains( x.Key ) ) )
		{
			var entry = pending.First( x => needed.Contains( x.Value.Guid ) && !guids.Contains( x.Value.Guid ) && !paths.Contains( x.Key ) );
			pending.Remove( entry );
			var (path, record) = entry;
			if ( !File.Exists( Target( path ) ) ) continue;
			var source = System.IO.Path.Combine( scratch, record.Guid + ".source" );
			File.Copy( Target( path ), source );
			var hash = HashFile( source, cancel );
			snapshots[path] = hash;
			string meta = null;
			if ( record.Metadata != null ) { meta = source + ".meta"; File.WriteAllText( meta, record.Metadata ); }
			var dependency = new UnityAsset { Path = path, Guid = record.Guid, Source = source, Metadata = meta, Kind = UnityArchive.Classify( path ) };
			copies.Add( dependency );
			guids.Add( record.Guid ); paths.Add( path );
			References( dependency );
		}
		using var view = UnityArchive.View( original.Archive.FileName, copies );
		foreach ( var choice in Sources.Where( c => c.ExistingHash != null && !c.Replace ) )
			if ( index.TryGetValue( choice.Path, out var prior ) && prior.Guid != null )
				{
					if ( copies.Any( a => a.Guid == prior.Guid && !a.Path.Equals( choice.Path, StringComparison.OrdinalIgnoreCase ) ) )
						throw new InvalidDataException( $"Retained source {choice.Path} has a GUID also supplied at another path. Resolve that source conflict before importing." );
					view.GuidAliases[prior.Guid] = copies.Single( a => a.Path == choice.Path ).Guid;
				}
		var catalog = ImportCatalog.Read( view, cancel );
		var effectivePlan = catalog.CreatePlan( original.Options );
		foreach ( var item in effectivePlan.Assets ) effective[item.Asset.Path] = item.Asset;
		prepared = UnityImport.PrepareFiles( effectivePlan, assetsRoot, progress, cancel, extract );
		foreach ( var file in prepared.Files )
		{
			var path = System.IO.Path.GetRelativePath( prepared.Directory, file ).Replace( '\\', '/' );
			if ( path == "unity-import-report.json" ) continue;
			var existing = HashFile( Target( path ), cancel );
			if ( snapshots.TryGetValue( path, out var previous ) && previous != existing ) throw new IOException( $"Destination changed during preparation: {path}. Review the import again." );
			snapshots[path] = existing;
			var incoming = HashFile( file, cancel );
			var owner = Owner( path );
			index.TryGetValue( path, out var previousRecord );
			var differentOwner = owner != null && previousRecord?.SourcePath != null &&
				(!previousRecord.SourcePath.Equals( owner.Path, StringComparison.OrdinalIgnoreCase ) || previousRecord.Guid != owner.Guid);
			var choice = Choice( path, incoming, existing, existing != null && (incoming != existing || differentOwner),
				differentOwner ? $"This output was previously generated from {previousRecord.SourcePath} (GUID {previousRecord.Guid})." : "" );
			if ( chosen.TryGetValue( path, out var sourceChoice ) ) choice.Replace = sourceChoice.Replace;
			else if ( effective.ContainsKey( path ) ) choice.Replace = false;
			Outputs.Add( choice );
		}
	}

	public ImportResult Commit( IProgress<ImportProgress> progress, CancellationToken cancel )
	{
		if ( prepared == null ) throw new InvalidOperationException( "Prepare this import before committing." );
		UnityImport.CheckDirectory( stateRoot );
		UnityImport.CheckDirectory( indexPath );
		Directory.CreateDirectory( stateRoot );
		using var mergeLock = new FileStream( System.IO.Path.Combine( stateRoot, "merge.lock" ), FileMode.OpenOrCreate, FileAccess.ReadWrite, FileShare.None );
		CheckUnchanged( cancel );
		var writes = Outputs.Where( c => c.ExistingHash == null || c.Conflict && c.Replace ).ToArray();
		var changes = new List<(string Target, string Backup)>();
		var newIndex = new Dictionary<string, ImportedFileRecord>( index, StringComparer.OrdinalIgnoreCase );
		foreach ( var output in Outputs )
		{
			index.TryGetValue( output.Path, out var old );
			var written = output.ExistingHash == null || output.Conflict && output.Replace;
			var hash = written ? output.IncomingHash : output.ExistingHash;
			effective.TryGetValue( output.Path, out var asset );
			var owner = Owner( output.Path );
			var sourceChoice = Sources.FirstOrDefault( c => c.Path.Equals( output.Path, StringComparison.OrdinalIgnoreCase ) );
			var acceptedSource = asset != null && (sourceChoice?.ExistingHash == null || sourceChoice.Replace || !sourceChoice.Conflict);
			// Do not adopt local edits or rejected versions as the importer's new baseline.
			newIndex[output.Path] = new ImportedFileRecord
			{
				Hash = written || old == null || sourceChoice?.Replace == true ? hash : old.Hash,
				Guid = acceptedSource ? asset.Guid : asset == null && written ? owner?.Guid : old?.Guid,
				Metadata = acceptedSource ? (asset.Metadata == null ? null : UnityImport.ReadText( asset.Metadata )) : old?.Metadata,
				SourcePath = asset != null ? asset.Path : written ? owner?.Path : old?.SourcePath,
				Conversion = written ? Version + ":" + original.Options : old?.Conversion,
				Packages = (old?.Packages ?? Array.Empty<string>()).Append( Package ).Distinct().ToArray()
			};
		}
		ReportPath = System.IO.Path.Combine( stateRoot, "reports", Guid.NewGuid().ToString( "N" ) + ".json" );
		var reportSource = System.IO.Path.Combine( scratch, "report.json" );
		var report = System.Text.Json.Nodes.JsonNode.Parse( File.ReadAllText( System.IO.Path.Combine( prepared.Directory, "unity-import-report.json" ) ) );
		report["Destination"] = Destination;
		report["Files"] = JsonSerializer.SerializeToNode( Outputs.Select( c => c.Path ).ToArray() );
		report["Merge"] = JsonSerializer.SerializeToNode( new
		{
			Sources = Sources.Select( c => new { c.Path, c.Status, c.Action, c.Detail } ),
			Files = Outputs.Select( c => new { c.Path, c.Status, c.Action, c.Detail } )
		} );
		File.WriteAllText( reportSource, report.ToJsonString( new JsonSerializerOptions { WriteIndented = true } ) );
		var indexSource = System.IO.Path.Combine( scratch, "index.json" );
		File.WriteAllText( indexSource, JsonSerializer.Serialize( newIndex, new JsonSerializerOptions { WriteIndented = true } ) );
		try
		{
			for ( var i = 0; i < writes.Length; i++ )
			{
				cancel.ThrowIfCancellationRequested();
				var choice = writes[i];
				progress?.Report( new( 0.96 + 0.03 * i / Math.Max( 1, writes.Length ), $"Writing {choice.Path}" ) );
				Write( Target( choice.Path ), System.IO.Path.Combine( prepared.Directory, choice.Path ), choice.ExistingHash );
			}
			cancel.ThrowIfCancellationRequested();
			Write( ReportPath, reportSource, null );
			Write( indexPath, indexSource, indexHash );
		}
		catch ( Exception failure )
		{
			var recovery = new List<string>();
			foreach ( var change in changes.AsEnumerable().Reverse() )
			{
				try
				{
					ImportStorage.Retry( () =>
					{
						if ( change.Backup == null ) File.Delete( change.Target );
						else File.Copy( change.Backup, change.Target, true );
					}, CancellationToken.None );
				}
				catch ( Exception ex ) { recovery.Add( $"{change.Target}: {ex.Message}" ); }
			}
			if ( recovery.Count > 0 )
			{
				preserveRecovery = true;
				throw new IOException( $"Merge failed and some files could not be restored. Backups and recovery.json are in {scratch}. " + string.Join( "; ", recovery ), failure );
			}
			throw;
		}
		progress?.Report( new( 1, "Files imported" ) );
		return new ImportResult( Destination, prepared.AssetCount, prepared.ConvertedCount,
			Outputs.Select( c => Target( c.Path ) ).ToArray(), prepared.Warnings )
		{
			Unresolved = prepared.Unresolved, ReportPath = ReportPath, ChangedFiles = writes.Select( c => Target( c.Path ) ).ToArray()
		};

		void Write( string target, string source, string expected )
		{
			UnityImport.CheckDirectory( target );
			if ( HashFile( target, cancel ) != expected ) throw new IOException( $"Destination changed since review: {target}. Review the import again." );
			Directory.CreateDirectory( System.IO.Path.GetDirectoryName( target ) );
			string backup = null;
			if ( File.Exists( target ) )
			{
				backup = System.IO.Path.Combine( scratch, "backup-" + changes.Count );
				File.Copy( target, backup );
			}
			File.WriteAllText( System.IO.Path.Combine( scratch, "recovery.json" ), JsonSerializer.Serialize( changes.Append( (Target: target, Backup: backup) ).Select( c => new { c.Target, c.Backup } ) ) );
			// Replace from a sibling temp file; readers never see a partially copied source.
			var temp = target + ".unity-write-" + Guid.NewGuid().ToString( "N" );
			try
			{
				File.Copy( source, temp );
				ImportStorage.Retry( () =>
				{
					if ( HashFile( target, cancel ) != expected ) throw new IOException( $"Destination changed while writing: {target}" );
					File.Move( temp, target, expected != null );
				}, cancel );
				changes.Add( (target, backup) );
			}
			finally { if ( File.Exists( temp ) ) File.Delete( temp ); }
		}
	}

	UnityAsset Owner( string path )
	{
		if ( effective.TryGetValue( path, out var source ) ) return source;
		return effective.Values.FirstOrDefault( a =>
			(a.Kind == UnityAssetKind.Material && (System.IO.Path.ChangeExtension( a.Path, ".vmat" ).Equals( path, StringComparison.OrdinalIgnoreCase ) || System.IO.Path.ChangeExtension( a.Path, ".tmat" ).Equals( path, StringComparison.OrdinalIgnoreCase ))) ||
			(a.Kind == UnityAssetKind.Model && System.IO.Path.ChangeExtension( a.Path, ".vmdl" ).Equals( path, StringComparison.OrdinalIgnoreCase )) ||
			(a.Kind == UnityAssetKind.Texture && (a.Path + ".vtex").Equals( path, StringComparison.OrdinalIgnoreCase )) );
	}

	void CheckUnchanged( CancellationToken cancel )
	{
		if ( HashFile( indexPath, cancel ) != indexHash ) throw new IOException( "Another import changed the import records. Review this import again." );
		foreach ( var (path, hash) in snapshots )
			if ( HashFile( Target( path ), cancel ) != hash ) throw new IOException( $"Destination changed since review: {path}. Review the import again." );
	}

	string Target( string path )
	{
		var safe = UnityArchive.SafePath( "Assets/" + path );
		var target = System.IO.Path.GetFullPath( System.IO.Path.Combine( Destination, safe ) );
		UnityImport.CheckDirectory( target );
		if ( Directory.Exists( target ) ) throw new IOException( $"A directory occupies the file destination: {path}" );
		return target;
	}

	public static string HashFile( string path, CancellationToken cancel = default )
	{
		if ( !File.Exists( path ) ) return null;
		using var input = File.OpenRead( path );
		using var hash = IncrementalHash.CreateHash( HashAlgorithmName.SHA256 );
		var buffer = new byte[128 * 1024];
		int count;
		while ( (count = input.Read( buffer, 0, buffer.Length )) > 0 ) { cancel.ThrowIfCancellationRequested(); hash.AppendData( buffer, 0, count ); }
		return Convert.ToHexString( hash.GetHashAndReset() ).ToLowerInvariant();
	}

	public void Dispose()
	{
		if ( preserveRecovery ) return;
		foreach ( var directory in new[] { prepared?.Directory, scratch }.Where( p => p != null && Directory.Exists( p ) ) )
		{
			// Both are private, generated staging directories, never the shared destination.
			var parent = System.IO.Path.GetDirectoryName( assetsRoot );
			if ( System.IO.Path.GetDirectoryName( directory ) != parent || !System.IO.Path.GetFileName( directory ).StartsWith( ".unity-", StringComparison.Ordinal ) )
				throw new IOException( "Refusing to clean a directory outside importer staging." );
			try { ImportStorage.Retry( () => Directory.Delete( directory, true ), CancellationToken.None ); }
			catch ( IOException ) { /* Completed imports must not become failures because temporary cleanup was locked. */ }
			catch ( UnauthorizedAccessException ) { }
		}
	}
}