﻿using System.Collections;
using System.Collections.Generic;
using UnityEngine;
using UnityEditor;

/*
 * Instructions:
 * 1. Ensure this script is under "Editor" folder.
 * 2. After you copied this script to your project, you will notice a new Menu Option on the bar
 *    named "Atlas Tool".
 * 3. The first 2 are toggles.
 *    3.1 First toggle is delete packing tag after a non-dry run. 
 *        This will ensure Sprite Packer will not be use since packing tags are removed.
 *    3.2 Second toggle is dry run. With it turned on, the run will just report the result,
 *        instead of executing it.
 * 4. 3rd option is execute the convertion operation. 
 * 5. If "dry run" is disabled, you should notice new Sprite Atlas asset are created at Assets/ folder.
 * 6. If "Toggle Remove Packing Tag" is turned on, all textures will lost their packing tag.
 *    Thus, when you pack, Sprite Atlas will be use instead. 
 */ 

[InitializeOnLoad]
public class PackingTagToSpriteAtlas
{
	// Edit this to get all sprite atlas asset to your desire location. "Assets/" is compulsory.
	private const string kDestinationPath = "Assets/";

	private const string kMainMenu = "Atlas Tool/";
	private const string kToggleRemovePackingTag = "Toggle Remove Packing Tag";
	private const string kToggleDryRun = "Toggle Dry Run";
	private const string kCreateAtlasFromPackingTag = "Create Atlas from Packing Tag";

	private static bool m_RemovePackingTagAfterPacked = true;
	private static bool m_DryRun = true;

	static PackingTagToSpriteAtlas ()
	{
		EditorApplication.delayCall += () =>
		{
			Menu.SetChecked (kMainMenu + kToggleRemovePackingTag, m_RemovePackingTagAfterPacked);
			Menu.SetChecked (kMainMenu + kToggleDryRun, m_DryRun);
		};
	}

	[MenuItem (kMainMenu + kToggleRemovePackingTag, false, 0)]
	private static void ToggleRemovalPackingTag ()
	{
		m_RemovePackingTagAfterPacked = !m_RemovePackingTagAfterPacked;
		Menu.SetChecked (kMainMenu + kToggleRemovePackingTag, m_RemovePackingTagAfterPacked);
	}

	[MenuItem (kMainMenu + kToggleDryRun, false, 1)]
	private static void ToggleDryRun ()
	{
		m_DryRun = !m_DryRun;
		Menu.SetChecked (kMainMenu + kToggleDryRun, m_DryRun);
	}

	[MenuItem (kMainMenu + kCreateAtlasFromPackingTag, false, 21)]
	private static void GenerateAtlasAssetFromPackingTag ()
	{
		CreateAtlasFromPackingTag ();
	}

	// For Sprite Atlas Creation
	private class TextureWithOldTag
	{
		public string oldTag = "";
		public TextureImporter textureImporter;
	}

	private class TreeNode
	{
		public string path = "";
		public TreeNode parent = null;
		public List<string> tags = new List<string> ();
		public List<TextureWithOldTag> textures = new List<TextureWithOldTag> ();
		public Dictionary<string, TreeNode> children = new Dictionary<string, TreeNode> ();
	}

	private static void CreateAtlasFromPackingTag ()
	{
		TreeNode root = new TreeNode ();
		GenerateTreeFromProject (root);

		Dictionary<string, SpriteAtlas> spriteAtlases = new Dictionary<string, SpriteAtlas> ();
		CollectExisitngSpriteAtlas (spriteAtlases);
		SpawnSpriteAtlasFromTree (root, spriteAtlases);

        if (!m_DryRun)
            EditorSettings.spritePackerMode = SpritePackerMode.Disabled;

        DumpLog(spriteAtlases);
	}

	private static void GenerateTreeFromProject (TreeNode root)
	{
		string[] guids = AssetDatabase.FindAssets ("t:texture2d");
		foreach (string guid in guids)
		{
			string path = AssetDatabase.GUIDToAssetPath (guid);

			TreeNode node = FindOrCreateNodeFromPath (root, path);

			TextureImporter t = AssetImporter.GetAtPath (path) as TextureImporter;
			var tag = t.spritePackingTag;
			if (tag != null)
			{
				// Remove the packing tag (used by legacy sprite packer)
				if (m_RemovePackingTagAfterPacked && !m_DryRun)
					t.spritePackingTag = "";

				if (tag != "")
				{
					TextureWithOldTag tot = new TextureWithOldTag ();
					tot.textureImporter = t;
					tot.oldTag = tag;
					node.textures.Add (tot);
				}

				// Register this tag to all the parent folders
				while (node != null)
				{
					if (node.tags.Contains (tag) == false)
						node.tags.Add (tag);
					node = node.parent;
				};

				// Needed for removing of packing tag
				if (m_RemovePackingTagAfterPacked && !m_DryRun)
					t.SaveAndReimport ();
			}
		}
	}

	private static TreeNode FindOrCreateNodeFromPath (TreeNode root, string path)
	{
		TreeNode node = root;

		string[] tokens = path.Split ('/');
		string currentPath = tokens[0];
		for (int i = 1; i < tokens.Length - 1; ++i)
		{
			currentPath += "/" + tokens[i];

			TreeNode childNode;
			node.children.TryGetValue (tokens[i], out childNode);
			if (childNode == null)
			{
				childNode = new TreeNode ();
				childNode.path = currentPath;
				childNode.parent = node;
				node.children.Add (tokens[i], childNode);
			}

			node = childNode;
		}

		return node;
	}

	private static void CollectExisitngSpriteAtlas (Dictionary<string, SpriteAtlas> spriteAtlases)
	{
		string[] guids = AssetDatabase.FindAssets ("t:spriteAtlas");
		foreach (string guid in guids)
		{
			SpriteAtlas sa = AssetDatabase.LoadAssetAtPath<SpriteAtlas> (AssetDatabase.GUIDToAssetPath (guid));
			spriteAtlases.Add (sa.name, sa);
		}
	}

	private static void SpawnSpriteAtlasFromTree (TreeNode root, Dictionary<string, SpriteAtlas> spriteAtlases)
	{
		// If the folder is...
		// 1, Not root aka Assets/
		// 2, Contains only 1 tag. Which indicate all the textures and subfolder underneath this folder all use the same packing tag.
		// Add this folder to the sprite atlas and stop here, since we assume everything underneath this folder belong to same atlas.
		if (root.parent != null
			&& root.tags.Count == 1)
		{
			string tag = root.tags[0];
			if (tag != "")
			{
				SpriteAtlas sa = FindOrCreateSpriteAtlasWithTag (tag, spriteAtlases);
				DefaultAsset folderAsset = AssetDatabase.LoadAssetAtPath<DefaultAsset> (root.path);
				sa.AddPackableObjects (new Object[] { folderAsset });
			}
		}
		// If the folder is...
		// 1, root.
		// 2, Contains more than 1 tag. Empty tag (texture without packing tag) does consider as one kind of tag.
		else
		{
			// If there are texture under this folder, add each of them to their respected sprite atlas.
			foreach (var t in root.textures)
			{
				SpriteAtlas sa = FindOrCreateSpriteAtlasWithTag (t.oldTag, spriteAtlases);
				Texture2D tex = AssetDatabase.LoadAssetAtPath<Texture2D> (t.textureImporter.assetPath);
				sa.AddPackableObjects (new Object[] { tex });
			}

			// Traverse sub-folders if there are.
			foreach (var node in root.children.Values)
				SpawnSpriteAtlasFromTree (node, spriteAtlases);
		}
	}

	private static SpriteAtlas FindOrCreateSpriteAtlasWithTag (string tag, Dictionary<string, SpriteAtlas> spriteAtlases)
	{
		SpriteAtlas sa;
		spriteAtlases.TryGetValue (tag, out sa);
		if (!sa)
		{
			sa = new SpriteAtlas ();
			sa.name = tag;
			spriteAtlases.Add (tag, sa);

			// Do not create the asset during dry run.
			if (!m_DryRun)
				AssetDatabase.CreateAsset (sa, kDestinationPath + tag + ".spriteatlas");
		}

		return sa;
	}
	
	private static void DumpLog (Dictionary<string, SpriteAtlas> spriteAtlases)
	{
		if (m_DryRun)
			Debug.LogWarning ("This is a dry run. Please uncheck the option and run the \"" + kCreateAtlasFromPackingTag + "\" option again to get a real run.");
		else if (!m_RemovePackingTagAfterPacked)
			Debug.LogWarning ("A real run was done without removing packing tag from textures. This will cause textures with packing tag still pack by legacy Sprite Packer instead of the atlas." +
			                  "Sprite Atlas Asset will only pack textures without a packing tag.");

		bool canBeOptimized = false;
		foreach (var item in spriteAtlases)
		{
			SpriteAtlas sa = item.Value;
			string message = "Atlas <b>" + sa.name + "</b> contains... (click for full list)\n";
			foreach (var obj in sa.packables)
			{
				bool isAFolder = obj is DefaultAsset;
				if (isAFolder)
					message += "Folder : " + AssetDatabase.GetAssetPath (obj) + "\n";
				else
					message += "*Texture : " + AssetDatabase.GetAssetPath (obj) + "\n";

				if (!isAFolder)
					canBeOptimized = true;
			}
			Debug.Log (message);
		}

		if (canBeOptimized)
			Debug.LogWarning ("* indicate that this texture was placed in a folder with mixed sub-folder or texture of different packing tag." +
							  "This can be optimized by putting textures of same packing tag into one common folder and do not mix other texture of different packing tag." +
							  "Once this is done, add this new folder into the atlas instead of individual texture.");
	}
}
