﻿/*
*  **Based on the Global Selection Base Attribute behaviour script**
*
* Similar to Unity's built-in, but doesn't automatically assume a prefab should
* be a SelectionBase as this is an issue for nested prefabs.
*
* Behaviour (for selections made in SceneView only):
*
*   1. If selected GO was selected in the last selection
*      > the selection remains unchanged
*    
*   2. If the selected GO has [PrimaryObject]
*      > the selection remains unchanged
*
*   3. If GO without [PrimaryObject] object selected
*        AND it has a GO with [PrimaryObject] in its parents
*        AND that parent was NOT in the previous selection
*      
*      > the selection is switched to the parent with [PrimaryObject]
*
* (1) Is straight-forward
* (2) Ensures once an object has been adjusted once, it won't be again
* (3) Allows the user to drill-down by re-clicking the same object.  Sequence is:
*       > first click: ImmediateParent with [PrimaryObject] selected instead
*       > second click: selection remains on clicked object
*
*/

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

[InitializeOnLoad]
public class PrimaryObjectEvents : Editor
{
	private static List<UnityEngine.Object> newSelection = null;
	private static UnityEngine.Object[] lastSelection = new UnityEngine.Object[]{};

	static PrimaryObjectEvents()
	{
		// Ensure we're told when the selection changes
		Selection.selectionChanged += OnSelectionChanged;

		// For some reason I can't be bothered investigating, you can't modify selections
		// while in OnSelectionChanged() so... hack to do it in Update() instead
		EditorApplication.update += OnSceneUpdate;
	}


	public static void OnSelectionChanged()
	{
		// Only modify user selection if selected from the SceneView
		System.Type windowOver = SceneView.mouseOverWindow.GetType();
		System.Type sceneView = typeof(SceneView);
		if (!windowOver.Equals(sceneView)) return;

		//  Look through them all, adjusting as needed
		newSelection = new List<UnityEngine.Object>();
		bool changed = false;
		foreach(GameObject go in Selection.GetFiltered<GameObject>(SelectionMode.Unfiltered))
		{
			changed = changed | AdjustIfNeeded(go, lastSelection, newSelection);
		}

		// If nothing has changed, give the update nothing to reselect
		if (!changed)
		{
			newSelection = null;
		}

		// Remember this selection so we can compare the next selection to it
		lastSelection = Selection.objects;
	}


	private static bool AdjustIfNeeded(GameObject go, object[] lastSelection,List<UnityEngine.Object> newSelection)
	{
		// If it was in the last selection set, leave it be
		if (Array.IndexOf(lastSelection, go) < 0)
		{
			GameObject parentWithPrimaryObject = null;
			bool goHasPrimaryObject = ObjectHasPrimaryObject(go);
			parentWithPrimaryObject = ParentWithPrimaryObject(go);

			if
				(
					(!goHasPrimaryObject && parentWithPrimaryObject != null) &&
					(Array.IndexOf(lastSelection, parentWithPrimaryObject) < 0)
				)
			{
				// User NOT drilling down - replace selection with PrimaryObject parent
				newSelection.Add(parentWithPrimaryObject.gameObject);
				return true;
			}
		}


		newSelection.Add(go);   // original go
		return false;
	}

	public static void OnSceneUpdate()
	{
		if (newSelection != null)
		{
			Selection.objects = newSelection.ToArray();
			newSelection = null;
		}
	}

	public static bool ObjectHasPrimaryObject(GameObject go)
	{
		foreach(Component component in go.GetComponents<MonoBehaviour>())
		{
			if (component.GetType().GetCustomAttributes(typeof(PrimaryObject), true).Length > 0)
			{
				return true;
			}
		}

		return false;
	}


	public static GameObject ParentWithPrimaryObject(GameObject go)
	{
		if (go.transform.parent == null)
			return null;

		foreach (Component component in go.transform.parent.GetComponentsInParent<MonoBehaviour>(false)) {
			if (component.GetType ().GetCustomAttributes (typeof(PrimaryObject), true).Length > 0) {
				return component.gameObject;
			}
		}

		return null;
	}
}
