using System.Collections.Generic;
using System.Linq;
using UnityEditor;
using UnityEditor.IMGUI.Controls;
using UnityEngine;

public class AssetBundleAssetDuplicateTool : EditorWindow
{
	private static readonly Dictionary<string, string> s_AssetABNameCache = new Dictionary<string, string>(1000);

	private HashSet<string> m_ABNames = new HashSet<string>();
	private Dictionary<string, List<string>> m_AssetMap = new Dictionary<string, List<string>>(100);
	private AssetListTree m_Tree;
	private TreeViewState m_TreeState;
	private HashSet<string> m_NewABNames = new HashSet<string>();
	private string m_NewABName = null;


	[MenuItem("Giant Tools/AssetBundles/处理重复依赖")]
	public static void OpenWindow()
	{
		var tWnd = EditorWindow.GetWindow<AssetBundleAssetDuplicateTool>();
		tWnd.Show();
	}



	private void OnEnable()
	{
		RefreshABNames();
		if (null == m_TreeState)
			m_TreeState = new TreeViewState();
		if (null == m_Tree)
			m_Tree = new AssetListTree(m_TreeState, m_AssetMap, m_ABNames, MoveAllToAssetBundle);
	}



	private void OnGUI()
	{
		if(GUILayout.Button("扫描",GUILayout.Height(40)))
		{
			Scan();
			EditorGUIUtility.ExitGUI();
		}
		GUILayout.Space(20);
		GUILayout.BeginHorizontal();
		{
			if (EditorGUILayout.DropdownButton(new GUIContent("AssetBundles"), FocusType.Passive, GUILayout.Width(100)))
			{
				GenericMenu tMenu = new GenericMenu();
				foreach (var tABName in m_NewABNames)
					tMenu.AddItem(new GUIContent(tABName), false, OnMenuClicked, tABName);
				tMenu.AddSeparator("");
				var tABNames = AssetDatabase.GetAllAssetBundleNames();
				for (int tIdx = 0, tCount = tABNames.Length; tIdx < tCount; tIdx++)
					tMenu.AddItem(new GUIContent(tABNames[tIdx]), false, OnMenuClicked, tABNames[tIdx]);
				tMenu.DropDown(GUILayoutUtility.GetRect(position.width, 0));
			}
			m_NewABName = GUILayout.TextField(m_NewABName, GUILayout.MaxWidth(400));
			if(GUILayout.Button("新增", GUILayout.Width(60)) && !string.IsNullOrEmpty(m_NewABName) && m_NewABNames.Add(m_NewABName))
			{
				m_NewABName = null;
				RefreshABNames();
			}
			if (GUILayout.Button("一键移动", GUILayout.Width(60)) && !string.IsNullOrEmpty(m_NewABName))
			{
				MoveAutoIncludesToBundle(m_NewABName);
				m_NewABName = null;
				RefreshABNames();
				GUIUtility.ExitGUI();
			}
		}
		GUILayout.EndHorizontal();
		GUILayout.Space(20);
		if (m_AssetMap.Count > 0)
			m_Tree.OnGUI(GUILayoutUtility.GetRect(position.width, position.height - 50));
	}



	private void OnMenuClicked(object userData)
	{
		m_NewABName = userData as string;
	}



	private void Scan()
	{
		m_AssetMap.Clear();
		s_AssetABNameCache.Clear();
		var tAllMap = new Dictionary<string, HashSet<string>>(1000);
		var tABNames = AssetDatabase.GetAllAssetBundleNames();
		for (int tIdx = 0, tCount = tABNames.Length; tIdx < tCount; tIdx++)
		{
			var tABName = tABNames[tIdx];
			m_ABNames.Add(tABName);
			var tAssets = AssetDatabase.GetAssetPathsFromAssetBundle(tABName);
			ScanAssetBundle(tABName, tAssets, tAllMap);
			ScanAssetBundleDependencies(tABName, tAssets, tAllMap);
		}
		float tProgress = 1;
		foreach (var tPair in tAllMap)
		{
			if (!IsAssetInPackage(tPair.Key) && tPair.Value.Count > 1)
			{
				m_AssetMap.Add(tPair.Key, tPair.Value.ToList());
				EditorUtility.DisplayProgressBar($"发现{m_AssetMap.Count}个被重复包含的资源", tPair.Key, tProgress++ / tAllMap.Count);
			}
		}
		EditorUtility.ClearProgressBar();
		Debug.Log($"扫描结束，共找到{m_AssetMap.Count}个被重复包含的资源");
		m_Tree.SetSelection(new List<int>());
		m_Tree.Reload();
	}



	/// <summary>
	/// 将自动包含的资源移动到合适的AB包。
	/// 对于只被场景依赖的资源，因为Unity不允许资源和场景在同一个包里，所以让Unity自己去处理
	/// </summary>
	/// <param name="bundleName">AB包名称</param>
	public static void MoveAutoIncludesToBundle(string bundleName)
	{
		s_AssetABNameCache.Clear();
		var tAsset2ABMap = new Dictionary<string, HashSet<string>>(1000);
		var tABNames = AssetDatabase.GetAllAssetBundleNames();
		string tABName;
		for (int tIdx = 0, tCount = tABNames.Length; tIdx < tCount; tIdx++)
		{
			tABName = tABNames[tIdx];
			var tAssets = AssetDatabase.GetAssetPathsFromAssetBundle(tABName);
			ScanAssetBundle(tABName, tAssets, tAsset2ABMap);
			ScanAssetBundleDependencies(tABName, tAssets, tAsset2ABMap);
		}
		EditorUtility.ClearProgressBar();
		var tSceneABSet = new HashSet<string>();
		var tModifiedAssets = new Dictionary<string, AssetImporter>(1000);
		//第一遍扫描场景文件
		foreach (var tPair in tAsset2ABMap)
		{
			var tAssetName = tPair.Key;
			if(tAssetName.EndsWith(".unity", System.StringComparison.OrdinalIgnoreCase))
			{
				var tImporter = AssetImporter.GetAtPath(tAssetName);
				var tSceneABName = tImporter.assetBundleName;
				if (tPair.Value.Count != 1 /*多处依赖，基本不会出现在场景资源上*/ || tSceneABName != tPair.Value.ToArray()[0] /*没有显式指定AssetBundleName，应该也很少出现*/)
				{
					tSceneABName = System.IO.Path.GetFileNameWithoutExtension(tAssetName);
					tImporter.assetBundleName = tSceneABName;
					tModifiedAssets.Add(tAssetName, tImporter);
				}
				tSceneABSet.Add(tSceneABName);
			}
		}
		//第二遍扫描其他资源，确保场景和资源不在同一个AB包中，他们的组织形式不一样
		foreach (var tPair in tAsset2ABMap)
		{
			var tAssetName = tPair.Key;
			//跳过场景资源和Package资源
			if(IsAssetInPackage(tAssetName))
			{
				var tImporter = AssetImporter.GetAtPath(tAssetName);
				if (null != tImporter && !string.IsNullOrEmpty(tImporter.assetBundleName))
				{
					tImporter.assetBundleName = string.Empty;
					tModifiedAssets.Add(tAssetName, tImporter);
				}
			}
			else if (!tAssetName.EndsWith(".unity", System.StringComparison.OrdinalIgnoreCase))
			{
				var tABSet = tPair.Value;
				var tImporter = AssetImporter.GetAtPath(tAssetName);
				if (tABSet.Count != 1)
				{
					//被多处依赖的资源，直接分配到指定AB包
					tImporter.assetBundleName = bundleName;
					tModifiedAssets.Add(tAssetName, tImporter);
				}
				else
				{
					//因为当前资源只被一个场景AB包引用，所以只需要判断是不是被显式指定到某个场景AB包中
					var tAssetABName = tImporter.assetBundleName;
					if(tSceneABSet.Contains(tAssetABName))
					{
						tImporter.assetBundleName = null;
						tModifiedAssets.Add(tAssetName, tImporter);
					}
				}
			}
		}
		//应用修改
		foreach (var tPair in tModifiedAssets)
			EditorUtility.SetDirty(tPair.Value);
		AssetDatabase.SaveAssets();
		AssetDatabase.Refresh();
	}



	private static void ScanAssetBundle(string abName, string[] assets, Dictionary<string, HashSet<string>> allMap)
	{
		string tAssetABName;
		for (int tAssetIdx = 0, tAssetCount = assets.Length; tAssetIdx < tAssetCount; tAssetIdx++)
		{
			var tAssetName = assets[tAssetIdx];
			if (IgnoreAsset(tAssetName)) continue;
			EditorUtility.DisplayProgressBar($"解析{abName}包的资源", tAssetName, (tAssetIdx + 1.0f) / tAssetCount);
			HashSet<string> tABSet = null;
			if (!allMap.TryGetValue(tAssetName, out tABSet))
			{
				tABSet = new HashSet<string>();
				allMap[tAssetName] = tABSet;
			}
			if (!s_AssetABNameCache.TryGetValue(tAssetName, out tAssetABName))
			{
				tAssetABName = AssetImporter.GetAtPath(tAssetName).assetBundleName;
				s_AssetABNameCache[tAssetName] = tAssetABName;
			}
			if (string.IsNullOrEmpty(tAssetABName))
				tABSet.Add(abName);
			else
				tABSet.Add(tAssetABName);
		}
	}



	private static void ScanAssetBundleDependencies(string abName, string[] assets, Dictionary<string, HashSet<string>> allMap)
	{
		string tAssetABName;
		var tDepends = AssetDatabase.GetDependencies(assets, true);
		for (int tDepIdx = 0, tDepCount = tDepends.Length; tDepIdx < tDepCount; tDepIdx++)
		{
			var tDepAssetName = tDepends[tDepIdx];
			if (IgnoreAsset(tDepAssetName)) continue;
			EditorUtility.DisplayProgressBar($"解析{abName}包资源的依赖", tDepAssetName, (tDepIdx + 1.0f) / tDepCount);
			HashSet<string> tABSet = null;
			if (!allMap.TryGetValue(tDepAssetName, out tABSet))
			{
				tABSet = new HashSet<string>();
				allMap[tDepAssetName] = tABSet;
			}
			if (!s_AssetABNameCache.TryGetValue(tDepAssetName, out tAssetABName))
			{
				tAssetABName = AssetImporter.GetAtPath(tDepAssetName).assetBundleName;
				s_AssetABNameCache[tDepAssetName] = tAssetABName;
			}
			if (string.IsNullOrEmpty(tAssetABName))
				tABSet.Add(abName);
			else
				tABSet.Add(tAssetABName);
		}
	}



	private void RefreshABNames()
	{
		m_ABNames.Clear();
		foreach (var tABName in m_NewABNames)
			m_ABNames.Add(tABName);
		var tABNames = AssetDatabase.GetAllAssetBundleNames();
		for (int tIdx = 0, tCount = tABNames.Length; tIdx < tCount; tIdx++)
		{
			var tABName = tABNames[tIdx];
			m_ABNames.Add(tABName);
		}
	}



	private static bool IgnoreAsset(string assetPath)
	{
		return string.IsNullOrEmpty(assetPath) || assetPath.EndsWith(".cs");
	}



	private static bool IsAssetInPackage(string assetPath)
	{
		return (assetPath.StartsWith("Packages/", System.StringComparison.Ordinal));
	}



	private void MoveAllToAssetBundle(string bundleName, List<string> assets)
	{
		if(!string.IsNullOrEmpty(bundleName) && null != assets && assets.Count > 0)
		{
			for (int tIdx = 0, tCount = assets.Count; tIdx < tCount; tIdx++)
			{
				var tAssetPath = assets[tIdx];
				var tImp = AssetImporter.GetAtPath(tAssetPath);
				if(null != tImp)
				{
					EditorUtility.DisplayProgressBar($"移动到{bundleName}", tAssetPath, (tIdx + 1.0f) / tCount);
					tImp.assetBundleName = bundleName;
					//tImp.SaveAndReimport();
					EditorUtility.SetDirty(tImp);
				}
				m_AssetMap.Remove(tAssetPath);
			}
		}
		EditorUtility.ClearProgressBar();
		AssetDatabase.SaveAssets();
		AssetDatabase.Refresh();
		m_Tree.SetSelection(new List<int>());
		m_Tree.Reload();
	}
}


internal class AssetListTree : TreeView
{
	private System.Action<string, List<string>> m_MoveCallback = null;
	private Dictionary<string, List<string>> m_AssetMap = null;
	private HashSet<string> m_ABNames = null;
	private Dictionary<int, string> m_IdMap = new Dictionary<int, string>();

	private static MultiColumnHeaderState CreateHeaderState()
	{
		return new MultiColumnHeaderState(GetColumns());
	}

	private static MultiColumnHeaderState.Column[] GetColumns()
	{
		var retVal = new MultiColumnHeaderState.Column[] {
				new MultiColumnHeaderState.Column(),
				new MultiColumnHeaderState.Column()
			};
		retVal[0].headerContent = new GUIContent("Asset", "path of asset.");
		retVal[0].minWidth = 100;
		retVal[0].width = 600;
		retVal[0].maxWidth = 800;
		retVal[0].headerTextAlignment = TextAlignment.Left;
		retVal[0].canSort = false;
		retVal[0].autoResize = true;

		retVal[1].headerContent = new GUIContent("Bundles", "Bundles include this asset.");
		retVal[1].minWidth = 100;
		retVal[1].width = 600;
		retVal[1].maxWidth = 800;
		retVal[1].headerTextAlignment = TextAlignment.Left;
		retVal[1].canSort = false;
		retVal[1].autoResize = true;

		return retVal;
	}

	internal AssetListTree(TreeViewState state, Dictionary<string, List<string>> data, HashSet<string> abNames, System.Action<string, List<string>> moveCallback) : base(state, new MultiColumnHeader(CreateHeaderState()))
	{
		m_AssetMap = data;
		m_ABNames = abNames;
		m_MoveCallback = moveCallback;
		showBorder = true;
		showAlternatingRowBackgrounds = true;
	}

	protected override TreeViewItem BuildRoot()
	{
		var root = new TreeViewItem(-1, -1);
		return root;
	}

	protected override IList<TreeViewItem> BuildRows(TreeViewItem root)
	{
		var tRows = new List<TreeViewItem>();
		var tId = 0;
		m_IdMap.Clear();
		foreach (var tPair in m_AssetMap)
		{
			//0是新
			if(tPair.Value.Count > 0)
			{
				m_IdMap.Add(tId, tPair.Key);
				tRows.Add(new TreeViewItem(tId++, 0, tPair.Key));
			}
		}
		return tRows;
	}

	protected override void RowGUI(RowGUIArgs args)
	{
		for (int i = 0; i < args.GetNumVisibleColumns(); ++i)
			CellGUI(args.GetCellRect(i), args.item, args.GetColumn(i), ref args);
	}

	private void CellGUI(Rect cellRect, TreeViewItem item, int column, ref RowGUIArgs args)
	{
		CenterRectUsingSingleLineHeight(ref cellRect);
		switch (column)
		{
			case 0:
				{
					var iconRect = new Rect(cellRect.x + 1, cellRect.y + 1, cellRect.height - 2, cellRect.height - 2);
					if (item.icon != null)
						GUI.DrawTexture(iconRect, item.icon, ScaleMode.ScaleToFit);
					DefaultGUI.Label(
						new Rect(cellRect.x + iconRect.xMax + 1, cellRect.y, cellRect.width - iconRect.width, cellRect.height),
						item.displayName,
						args.selected,
						args.focused);
				}
				break;
			case 1:
				if(m_AssetMap.ContainsKey(item.displayName))
					DefaultGUI.Label(cellRect, string.Join(", ", m_AssetMap[item.displayName]), args.selected, args.focused);
				break;
		}
	}

	protected override void SingleClickedItem(int id)
	{
		base.SingleClickedItem(id);
		string tAssetName = null;
		if (m_IdMap.TryGetValue(id, out tAssetName))
			Selection.activeObject = AssetDatabase.LoadAssetAtPath<UnityEngine.Object>(tAssetName);
	}

	protected override bool CanMultiSelect(TreeViewItem item)
	{
		return true;
	}

	protected override void ContextClickedItem(int id)
	{
		GenericMenu tMenu = new GenericMenu();
		tMenu.AddDisabledItem(new GUIContent("MoveToBundle"));
		tMenu.AddSeparator("");
		foreach (var tABName in m_ABNames)
			tMenu.AddItem(new GUIContent(tABName), false, MoveToBundle, tABName);
		tMenu.ShowAsContext();
	}

	private void MoveToBundle(object context)
	{
		List<string> tSelectedAssets = new List<string>();
		string tAsset;
		foreach (var tNodeID in GetSelection())
		{
			if (m_IdMap.TryGetValue(tNodeID, out tAsset))
				tSelectedAssets.Add(tAsset);
		}
		m_MoveCallback?.Invoke(context as string, tSelectedAssets);
	}
}
