// Copyright Epic Games, Inc. All Rights Reserved.

using EpicGames.Core;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace UnrealBuildBase
{
	public class Rules
	{
		/// <summary>
		/// Enum for types of rules files. Should match extensions in RulesFileExtensions.
		/// </summary>
		public enum RulesFileType
		{
			/// <summary>
			/// .build.cs files
			/// </summary>
			Module,

			/// <summary>
			/// .target.cs files
			/// </summary>
			Target,

			/// <summary>
			/// .automation.csproj files
			/// </summary>
			AutomationModule,

			/// <summary>
			/// .ubtplugin.csproj files
			/// </summary>
			UbtPlugin,
		}

		/// <summary>
		/// Cached list of rules files in each directory of each type
		/// </summary>
		class RulesFileCache
		{
			public List<FileReference> ModuleRules = new List<FileReference>();
			public List<FileReference> TargetRules = new List<FileReference>();
			public List<FileReference> AutomationModules = new List<FileReference>();
			public List<FileReference> UbtPlugins = new List<FileReference>();
		}

		/// Map of root folders to a cached list of all UBT-related source files in that folder or any of its sub-folders.
		/// We cache these file names so we can avoid searching for them later on.
		static Dictionary<DirectoryReference, RulesFileCache> RootFolderToRulesFileCache = new Dictionary<DirectoryReference, RulesFileCache>();

		/// <summary>
		/// 
		/// </summary>
		/// <param name="RulesFileType"></param>
		/// <param name="GameFolders"></param>
		/// <param name="ForeignPlugins"></param>
		/// <param name="AdditionalSearchPaths"></param>
		/// <param name="bIncludeEngine"></param>
		/// <param name="bIncludeTempTargets">Whether to include targets generated by UAT to accomodate content-only projects that need to be compiled to include plugins</param>
		/// <returns></returns>
		public static List<FileReference> FindAllRulesSourceFiles(RulesFileType RulesFileType, List<DirectoryReference>? GameFolders, List<FileReference>? ForeignPlugins, List<DirectoryReference>? AdditionalSearchPaths, bool bIncludeEngine = true, bool bIncludeTempTargets = true)
		{
			List<DirectoryReference> Folders = new List<DirectoryReference>();

			// Add all engine source (including third party source)
			if (bIncludeEngine)
			{
				Folders.AddRange(Unreal.GetExtensionDirs(Unreal.EngineDirectory, "Source"));
			}

			// @todo plugin: Disallow modules from including plugin modules as dependency modules? (except when the module is part of that plugin)

			// Get all the root folders for plugins
			List<DirectoryReference> RootFolders = new List<DirectoryReference>();
			if (bIncludeEngine)
			{
				RootFolders.AddRange(Unreal.GetExtensionDirs(Unreal.EngineDirectory));
			}
			if (GameFolders != null)
			{
				RootFolders.AddRange(GameFolders.SelectMany(x => Unreal.GetExtensionDirs(x)));
			}

			// Find all the plugin source directories
			foreach (DirectoryReference RootFolder in RootFolders)
			{
				DirectoryReference PluginsFolder = DirectoryReference.Combine(RootFolder, "Plugins");
				foreach (FileReference PluginFile in PluginsBase.EnumeratePlugins(PluginsFolder))
				{
					Folders.Add(DirectoryReference.Combine(PluginFile.Directory, "Source"));
				}
			}

			// Add all the extra plugin folders
			if (ForeignPlugins != null)
			{
				foreach (FileReference ForeignPlugin in ForeignPlugins)
				{
					Folders.Add(DirectoryReference.Combine(ForeignPlugin.Directory, "Source"));
				}
			}

			// Add in the game folders to search
			if (GameFolders != null)
			{
				foreach (DirectoryReference GameFolder in GameFolders)
				{
					Folders.AddRange(Unreal.GetExtensionDirs(GameFolder, "Source"));

					if (bIncludeTempTargets)
					{
						DirectoryReference GameIntermediateSourceFolder = DirectoryReference.Combine(GameFolder, "Intermediate", "Source");
						Folders.Add(GameIntermediateSourceFolder);
					}
				}
			}

			// Process the additional search path, if sent in
			if (AdditionalSearchPaths != null)
			{
				foreach (DirectoryReference AdditionalSearchPath in AdditionalSearchPaths)
				{
					if (AdditionalSearchPath != null)
					{
						if (DirectoryReference.Exists(AdditionalSearchPath))
						{
							Folders.Add(AdditionalSearchPath);
						}
						else
						{
							throw new Exception($"Couldn't find AdditionalSearchPath for rules source files '{AdditionalSearchPath}'");
						}
					}
				}
			}

			return FindAllRulesFiles(Folders, RulesFileType);
		}

		/// <summary>
		/// Invalidate the cache for the givcen directory
		/// </summary>
		/// <param name="DirectoryPath">Directory to invalidate</param>
        public static void InvalidateRulesFileCache(string DirectoryPath)
        {
            DirectoryReference Directory = new DirectoryReference(DirectoryPath);
            RootFolderToRulesFileCache.Remove(Directory);
            DirectoryLookupCache.InvalidateCachedDirectory(Directory);
        }

		/// <summary>
		/// Prefetch multiple directories in parallel
		/// </summary>
		/// <param name="Directories">The directories to cache</param>
		public static void PrefetchRulesFiles(IEnumerable<DirectoryReference> Directories)
		{
			ThreadPoolWorkQueue? Queue = null;
			try
			{
				foreach(DirectoryReference Directory in Directories)
				{
					if(!RootFolderToRulesFileCache.ContainsKey(Directory))
					{
						RulesFileCache Cache = new RulesFileCache();
						RootFolderToRulesFileCache[Directory] = Cache;

						if(Queue == null)
						{
							Queue = new ThreadPoolWorkQueue();
						}

						DirectoryItem DirectoryItem = DirectoryItem.GetItemByDirectoryReference(Directory);
						Queue.Enqueue(() => FindAllRulesFilesRecursively(DirectoryItem, Cache, Queue));
					}
				}
			}
			finally
			{
				if(Queue != null)
				{
					Queue.Dispose();
					Queue = null;
				}
			}
		}

		/// <summary>
		/// Finds all the rules of the given type under a given directory
		/// </summary>
		/// <param name="Directory">Directory to search</param>
		/// <param name="Type">Type of rules to return</param>
		/// <returns>List of rules files of the given type</returns>
		public static IReadOnlyList<FileReference> FindAllRulesFiles(DirectoryReference Directory, RulesFileType Type)
		{
			return FindAllRulesFiles(new List<DirectoryReference> { Directory }, Type);
		}

		/// <summary>
		/// Finds all the rules of the given type under a given directories
		/// </summary>
		/// <param name="Directories">Directories to search</param>
		/// <param name="Type">Type of rules to return</param>
		/// <returns>List of rules files of the given type</returns>
		public static List<FileReference> FindAllRulesFiles(IEnumerable<DirectoryReference> Directories, RulesFileType Type)
		{
			List<(DirectoryReference, RulesFileCache)> Caches = new List<(DirectoryReference, RulesFileCache)>(Directories.Count());
			using (ThreadPoolWorkQueue Queue = new ThreadPoolWorkQueue())
			{
				foreach (DirectoryReference Directory in Directories)
				{
					RulesFileCache? Cache;
					if (!RootFolderToRulesFileCache.TryGetValue(Directory, out Cache))
					{
						Cache = new RulesFileCache();
						Queue.Enqueue(() => FindAllRulesFilesRecursively(DirectoryItem.GetItemByDirectoryReference(Directory), Cache, Queue));
					}
					Caches.Add((Directory, Cache));
				}
			}

			List<FileReference> Files = new List<FileReference>();
			
			foreach ((DirectoryReference Directory, RulesFileCache Cache) in Caches)
			{
				if (!RootFolderToRulesFileCache.ContainsKey(Directory))
				{
					Cache.ModuleRules.Sort((A, B) => A.FullName.CompareTo(B.FullName));
					Cache.TargetRules.Sort((A, B) => A.FullName.CompareTo(B.FullName));
					Cache.AutomationModules.Sort((A, B) => A.FullName.CompareTo(B.FullName));
					Cache.UbtPlugins.Sort((A, B) => A.FullName.CompareTo(B.FullName));
					RootFolderToRulesFileCache[Directory] = Cache;
				}

				// Get the list of files of the type we're looking for
				if (Type == RulesFileType.Module)
				{
					Files.AddRange(Cache.ModuleRules);
				}
				else if (Type == RulesFileType.Target)
				{
					Files.AddRange(Cache.TargetRules);
				}
				else if (Type == RulesFileType.AutomationModule)
				{
					Files.AddRange(Cache.AutomationModules);
				}
				else if (Type == RulesFileType.UbtPlugin)
				{
					Files.AddRange(Cache.UbtPlugins);
				}
				else
				{
					throw new Exception($"Unhandled rules type: {Type}");
				}
			}

			return Files;
		}

		/// <summary>
		/// Search through a directory tree for any rules files
		/// </summary>
		/// <param name="Directory">The root directory to search from</param>
		/// <param name="Cache">Receives all the discovered rules files</param>
		/// <param name="Queue">Queue for adding additional tasks to</param>
		private static void FindAllRulesFilesRecursively(DirectoryItem Directory, RulesFileCache Cache, ThreadPoolWorkQueue Queue)
		{
			// Scan all the files in this directory
			bool bSearchSubFolders = true;
			foreach (FileItem File in Directory.EnumerateFiles())
			{
				if (File.HasExtension(".build.cs"))
				{
					lock(Cache.ModuleRules)
					{
						Cache.ModuleRules.Add(File.Location);
					}
					bSearchSubFolders = false;
				}
				else if (File.HasExtension(".target.cs"))
				{
					lock(Cache.TargetRules)
					{
						Cache.TargetRules.Add(File.Location);
					}
				}
				else if (File.HasExtension(".automation.csproj"))
				{
					lock(Cache.AutomationModules)
					{
						Cache.AutomationModules.Add(File.Location);
					}
					bSearchSubFolders = false;
				}
				else if (File.HasExtension(".ubtplugin.csproj"))
				{
					lock(Cache.UbtPlugins)
					{
						Cache.UbtPlugins.Add(File.Location);
					}
					bSearchSubFolders = false;
				}
			}

			// If we didn't find anything to stop the search, search all the subdirectories too
			if (bSearchSubFolders)
			{
				foreach (DirectoryItem SubDirectory in Directory.EnumerateDirectories())
				{
					Queue.Enqueue(() => FindAllRulesFilesRecursively(SubDirectory, Cache, Queue));
				}
			}
		}


	}
}