using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;
using UnityEditor;
using Object = UnityEngine.Object;

public static class DependencyTracker
{
    private static Dictionary<string, List<string>> dependantCache;

    static DependencyTracker()
    {
        EditorApplication.projectChanged += () => dependantCache = null;
    }

    [MenuItem("Assets/Select Dependants", priority = 31)]
    private static void SelectDependants()
    {
        Object[] selection = Selection.objects;
        if (selection.Length == 0)
            return;

        if (dependantCache != null)
            SelectDependants(selection);
        else
            StartCoroutine(GatherDependencies(), () => SelectDependants(selection));
    }

    private static void SelectDependants(Object[] selection)
    {
        var assetPaths = selection.Select(i => AssetDatabase.GetAssetPath(i));

        var dependantPaths = selection.Length == 1 && AssetDatabase.IsValidFolder(assetPaths.First())
            ? SelectFolderDependants(assetPaths.First())
            : EnumerateFileDependants(assetPaths);

        Selection.objects = dependantPaths
            .Select(i => AssetDatabase.LoadMainAssetAtPath(i))
            .Concat(selection)
            .ToArray();
    }

    private static IEnumerable<string> EnumerateFileDependants(IEnumerable<string> assetPaths)
    {
        foreach (var assetPath in assetPaths)
            if (dependantCache.TryGetValue(assetPath, out var dependantPaths))
                foreach (string dependantPath in dependantPaths)
                    yield return dependantPath;
    }

    // Find assets inside a folder that are referenced by assets outside it.
    private static IEnumerable<string> SelectFolderDependants(string folderPath)
    {
        foreach (var assetPath in dependantCache.Keys)
            if (assetPath.StartsWith(folderPath))
                foreach (var dependantPath in dependantCache[assetPath])
                    if (!dependantPath.StartsWith(folderPath))
                        yield return dependantPath;
    }

    private static IEnumerator GatherDependencies()
    {
        EditorUtility.DisplayProgressBar("Gathering Dependencies", "", 0);
        yield return null;

        var assets = AssetDatabase.FindAssets(null, new[] { "Assets" })
            .Distinct()
            .ToList();

        var lastYieldTime = DateTime.Now;

        var dependencies = new List<KeyValuePair<string, string>>();
        for (int i = 0; i < assets.Count; i++)
        {
            string assetPath = AssetDatabase.GUIDToAssetPath(assets[i]);

            var now = DateTime.Now;
            if (now - lastYieldTime > TimeSpan.FromMilliseconds(100))
            {
                lastYieldTime = now;
                EditorUtility.DisplayProgressBar("Gathering Dependencies", assetPath,
                    (float)i / assets.Count);
                yield return null;
            }

            foreach (var dependencyPath in AssetDatabase.GetDependencies(assetPath))
                if (dependencyPath != assetPath)
                    dependencies.Add(new KeyValuePair<string, string>(dependencyPath, assetPath));
        }

        EditorUtility.DisplayProgressBar("Gathering Dependencies", "Building dictionary", 1);
        yield return null;

        dependantCache = dependencies
            .GroupBy(i => i.Key)
            .ToDictionary(i => i.Key, i => i
                .Select(j => j.Value)
                .ToList());

        EditorUtility.ClearProgressBar();
    }

    private static void StartCoroutine(IEnumerator routine, Action onCompleted = null)
    {
        EditorApplication.CallbackFunction routineFunc = null;
        routineFunc = () =>
        {
            if (!routine.MoveNext())
            {
                EditorApplication.update -= routineFunc;
                if (onCompleted != null)
                    onCompleted();
            }
        };

        EditorApplication.update += routineFunc;
    }
}