using Sandbox;
using Editor;
using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp;

/// <summary>
/// A Roslyn compilation of the code around a file, for go to definition and autocomplete.
/// s&amp;box doesn't hand out its compilation, so we build our own the same way it does: same
/// sources, parse options and generated global usings, referencing every assembly the editor
/// has loaded plus the compiled output of the project's other compilers.
///
/// Compilations are cached per source folder and only re-parse files that changed, so
/// asking again after a keystroke is cheap.
///
/// Usage: <see cref="TakeSnapshot"/> on the main thread (it reads s&amp;box state), then
/// <see cref="GetCompilation"/> on any thread.
/// </summary>
public static class RoslynProject
{
	/// <summary>
	/// Everything needed from s&amp;box, captured on the main thread.
	/// </summary>
	public class Snapshot
	{
		public string FilePath;
		public string SourceRoot;
		public string AssemblyName;
		public CSharpParseOptions ParseOptions;
		public string GeneratedCode;
		public List<(string Name, byte[] Image, string XmlDocs)> CompiledReferences = new();
		public List<string> IgnoreFolders = new();
		public Dictionary<string, string> OpenBuffers = new( StringComparer.OrdinalIgnoreCase );
	}

	/// <summary>
	/// Global usings s&amp;box adds, for when the file isn't part of a loaded compiler.
	/// </summary>
	const string FallbackGlobalUsings =
		"global using System;\n" +
		"global using System.Collections.Generic;\n" +
		"global using System.Linq;\n" +
		"global using System.Threading.Tasks;\n" +
		"global using Sandbox;\n" +
		"global using static Sandbox.Internal.GlobalGameNamespace;\n";

	const string FallbackEditorUsings =
		"global using Editor;\n" +
		"global using static Sandbox.Internal.GlobalToolsNamespace;\n";

	const string GeneratedPath = "<generated>";

	// Loading metadata is the slow part, so file references are kept between lookups
	private static readonly ConcurrentDictionary<string, (DateTime Stamp, MetadataReference Reference)> FileReferences = new( StringComparer.OrdinalIgnoreCase );

	public static Snapshot TakeSnapshot( string filePath, string projectRoot, IEnumerable<(string Path, string Text)> openBuffers )
	{
		var snapshot = new Snapshot { FilePath = Path.GetFullPath( filePath ) };

		foreach ( var (path, text) in openBuffers )
			snapshot.OpenBuffers[Path.GetFullPath( path )] = text;

		// Which s&box compiler builds this file? The one whose source folder holds it most specifically.
		var compilers = AllCompilers().ToList();
		Compiler owner = null;
		string ownerRoot = null;

		foreach ( var compiler in compilers )
		{
			var root = SourceRootOf( compiler );
			if ( root is null || !ProjectFiles.IsUnder( snapshot.FilePath, root ) ) continue;
			if ( ownerRoot is null || root.Length > ownerRoot.Length )
			{
				owner = compiler;
				ownerRoot = root;
			}
		}

		if ( owner is not null )
		{
			var config = owner.GetConfiguration();
			snapshot.SourceRoot = ownerRoot;
			snapshot.AssemblyName = owner.AssemblyName;
			snapshot.ParseOptions = config.GetParseOptions().WithDocumentationMode( DocumentationMode.Parse );
			snapshot.GeneratedCode = owner.GeneratedCode?.ToString();
			snapshot.IgnoreFolders = config.IgnoreFolders?.ToList() ?? new();
		}
		else
		{
			// Not compiled by s&box right now (e.g. a project that isn't loaded): guess the
			// source folder from the project layout and use s&box's usual global usings
			var relative = Path.GetRelativePath( projectRoot, snapshot.FilePath );
			var top = relative.Split( Path.DirectorySeparatorChar )[0];
			snapshot.SourceRoot = top == relative ? projectRoot : Path.Combine( projectRoot, top );
			snapshot.ParseOptions = new CSharpParseOptions( LanguageVersion.Preview, DocumentationMode.Parse );

			var isEditor = top.Equals( "Editor", StringComparison.OrdinalIgnoreCase );
			snapshot.GeneratedCode = FallbackGlobalUsings + (isEditor ? FallbackEditorUsings : "");
		}

		// Other compilers' output stands in for their source (e.g. the game's code, seen from the editor project)
		foreach ( var compiler in compilers )
		{
			if ( compiler == owner ) continue;
			var image = compiler.Output?.AssemblyData;
			if ( image is { Length: > 0 } ) snapshot.CompiledReferences.Add( (compiler.AssemblyName, image, compiler.Output.XmlDocumentation) );
		}

		return snapshot;
	}

	private static IEnumerable<Compiler> AllCompilers()
	{
		var seen = new HashSet<Compiler>();

		if ( Project.CompileGroup is { } group )
			foreach ( var c in group.Compilers )
				if ( c is not null && seen.Add( c ) ) yield return c;

		foreach ( var project in EditorUtility.Projects.GetAll() )
		{
			var editorCompiler = project.EditorCompiler;
			if ( editorCompiler is not null && seen.Add( editorCompiler ) ) yield return editorCompiler;

			var compiler = project.Compiler;
			if ( compiler is not null && seen.Add( compiler ) ) yield return compiler;
		}
	}

	private static string SourceRootOf( Compiler compiler )
	{
		try
		{
			var root = compiler.FileSystem?.GetFullPath( "/" );
			return string.IsNullOrEmpty( root ) || !Directory.Exists( root ) ? null : Path.GetFullPath( root ).TrimEnd( '\\', '/' );
		}
		catch ( Exception )
		{
			return null;
		}
	}

	// --- CACHED COMPILATIONS ---

	private class CachedFile
	{
		public DateTime Stamp;
		public string Buffer;
		public SyntaxTree Tree;
	}

	private class Cache
	{
		public string AssemblyName;
		public CSharpParseOptions ParseOptions;
		public string GeneratedCode;
		public List<byte[]> CompiledImages;
		public CSharpCompilation Compilation;
		public readonly Dictionary<string, CachedFile> Files = new( StringComparer.OrdinalIgnoreCase );
	}

	private static readonly object CacheLock = new();
	private static readonly Dictionary<string, Cache> Caches = new( StringComparer.OrdinalIgnoreCase );

	/// <summary>
	/// The compilation for the snapshot's source folder, brought up to date, and the syntax tree of
	/// the snapshot's file (null if that file isn't part of it). Safe to call off the main thread.
	/// </summary>
	public static (CSharpCompilation Compilation, SyntaxTree Tree) GetCompilation( Snapshot snapshot )
	{
		lock ( CacheLock )
		{
			if ( !Caches.TryGetValue( snapshot.SourceRoot, out var cache ) || !SameSetup( cache, snapshot ) )
			{
				cache = Create( snapshot );
				Caches[snapshot.SourceRoot] = cache;
			}
			else if ( !cache.CompiledImages.SequenceEqual( snapshot.CompiledReferences.Select( x => x.Image ) ) )
			{
				// Another compiler rebuilt (hotload): swap in its new output
				cache.CompiledImages = snapshot.CompiledReferences.Select( x => x.Image ).ToList();
				cache.Compilation = cache.Compilation.WithReferences( References( snapshot ) );
			}

			Update( cache, snapshot );

			cache.Files.TryGetValue( snapshot.FilePath, out var target );
			return (cache.Compilation, target?.Tree);
		}
	}

	private static bool SameSetup( Cache cache, Snapshot snapshot ) =>
		cache.AssemblyName == snapshot.AssemblyName
		&& cache.GeneratedCode == snapshot.GeneratedCode
		&& Equals( cache.ParseOptions, snapshot.ParseOptions );

	private static Cache Create( Snapshot snapshot )
	{
		var cache = new Cache
		{
			AssemblyName = snapshot.AssemblyName,
			ParseOptions = snapshot.ParseOptions,
			GeneratedCode = snapshot.GeneratedCode,
			CompiledImages = snapshot.CompiledReferences.Select( x => x.Image ).ToList()
		};

		var trees = new List<SyntaxTree>();
		if ( !string.IsNullOrWhiteSpace( snapshot.GeneratedCode ) )
			trees.Add( CSharpSyntaxTree.ParseText( snapshot.GeneratedCode, snapshot.ParseOptions, GeneratedPath ) );

		cache.Compilation = CSharpCompilation.Create(
			snapshot.AssemblyName ?? "HotCode.Navigation",
			trees,
			References( snapshot ),
			new CSharpCompilationOptions( OutputKind.DynamicallyLinkedLibrary, allowUnsafe: true, nullableContextOptions: NullableContextOptions.Disable ) );

		return cache;
	}

	/// <summary>
	/// Re-parses files that changed on disk or in an open tab, adds new ones and drops deleted ones.
	/// </summary>
	private static void Update( Cache cache, Snapshot snapshot )
	{
		var current = new HashSet<string>( StringComparer.OrdinalIgnoreCase );
		var added = new List<SyntaxTree>();
		var removed = new List<SyntaxTree>();

		foreach ( var file in SourceFiles( snapshot ) )
		{
			current.Add( file );
			snapshot.OpenBuffers.TryGetValue( file, out var buffer );

			DateTime stamp;
			try { stamp = File.GetLastWriteTimeUtc( file ); }
			catch ( Exception e ) when ( e is IOException or UnauthorizedAccessException ) { continue; }

			cache.Files.TryGetValue( file, out var cached );

			var unchanged = cached is not null && (buffer is not null
				? cached.Buffer is not null && (ReferenceEquals( cached.Buffer, buffer ) || cached.Buffer == buffer)
				: cached.Buffer is null && cached.Stamp == stamp);

			if ( unchanged ) continue;

			string text = buffer;
			if ( text is null )
			{
				try { text = File.ReadAllText( file ); }
				catch ( Exception e ) when ( e is IOException or UnauthorizedAccessException ) { continue; }
			}

			var tree = CSharpSyntaxTree.ParseText( TextBuffer.Normalize( text ), snapshot.ParseOptions, file );
			if ( cached is not null ) removed.Add( cached.Tree );
			added.Add( tree );
			cache.Files[file] = new CachedFile { Stamp = stamp, Buffer = buffer, Tree = tree };
		}

		foreach ( var gone in cache.Files.Keys.Where( k => !current.Contains( k ) ).ToList() )
		{
			removed.Add( cache.Files[gone].Tree );
			cache.Files.Remove( gone );
		}

		if ( removed.Count > 0 ) cache.Compilation = cache.Compilation.RemoveSyntaxTrees( removed );
		if ( added.Count > 0 ) cache.Compilation = cache.Compilation.AddSyntaxTrees( added );
	}

	private static IEnumerable<string> SourceFiles( Snapshot snapshot )
	{
		List<string> files;
		try
		{
			files = Directory.EnumerateFiles( snapshot.SourceRoot, "*.cs", SearchOption.AllDirectories ).ToList();
		}
		catch ( Exception e ) when ( e is IOException or UnauthorizedAccessException )
		{
			yield break;
		}

		foreach ( var file in files )
		{
			var relative = Path.GetRelativePath( snapshot.SourceRoot, file );
			if ( CodeFileNode.IsInHiddenFolder( relative ) ) continue;
			if ( snapshot.IgnoreFolders.Any( f => relative.StartsWith( f.Trim( '/', '\\' ) + Path.DirectorySeparatorChar, StringComparison.OrdinalIgnoreCase ) ) ) continue;
			yield return Path.GetFullPath( file );
		}
	}

	private static List<MetadataReference> References( Snapshot snapshot )
	{
		var references = new List<MetadataReference>();
		var names = new HashSet<string>( StringComparer.OrdinalIgnoreCase );

		// Never reference the assembly we're compiling from source, or every type would be ambiguous
		if ( snapshot.AssemblyName is not null ) names.Add( snapshot.AssemblyName );

		foreach ( var (name, image, xmlDocs) in snapshot.CompiledReferences )
		{
			if ( name is null || !names.Add( name ) ) continue;
			var docs = string.IsNullOrEmpty( xmlDocs ) ? null : XmlDocs.FromText( xmlDocs );
			references.Add( MetadataReference.CreateFromImage( image, documentation: docs ) );
		}

		foreach ( var assembly in AppDomain.CurrentDomain.GetAssemblies() )
		{
			if ( assembly.IsDynamic ) continue;

			string location;
			try { location = assembly.Location; }
			catch ( NotSupportedException ) { continue; }

			if ( string.IsNullOrEmpty( location ) || !File.Exists( location ) ) continue;
			if ( !names.Add( assembly.GetName().Name ?? location ) ) continue;

			references.Add( FileReference( location ) );
		}

		return references;
	}

	private static MetadataReference FileReference( string path )
	{
		var stamp = File.GetLastWriteTimeUtc( path );
		if ( FileReferences.TryGetValue( path, out var cached ) && cached.Stamp == stamp )
			return cached.Reference;

		// s&box (and .NET) ship Foo.xml next to Foo.dll; with it, hints and hovers can show the docs
		var xml = Path.ChangeExtension( path, ".xml" );
		var docs = File.Exists( xml ) ? XmlDocs.FromFile( xml ) : null;
		var reference = MetadataReference.CreateFromFile( path, documentation: docs );
		FileReferences[path] = (stamp, reference);
		return reference;
	}
}