# HG changeset patch
# User IBBoard <dev@ibboard.co.uk>
# Date 1334520051 -3600
# Node ID d0454aae02416bdacf863b73b5bfb4da88381d9c
# Parent  2c3474d652e109de5552e880b9fc14b5db09a00e
Re #360: Add WinForms UI for adding contained units
 * Fix army tree for nesting

diff -r 2c3474d652e1 -r d0454aae0241 FrmArmyTree.cs
--- a/FrmArmyTree.cs	Sat Apr 14 20:48:13 2012 +0100
+++ b/FrmArmyTree.cs	Sun Apr 15 21:00:51 2012 +0100
@@ -80,23 +80,14 @@
 			{
 				ArmyCategory[] cats = army.Categories;
 				TreeNode[] catNodes = new TreeNode[cats.Length];
-				Unit[] units;
 
 				for (int i = 0; i < cats.Length; i++)
 				{
 					ArmyCategory cat = cats[i];
-					units = army.GetUnits(cat);
 					cat.NameChanged += TreeNameChangedMethod;
 					cat.PointsValueChanged += NodePointsValueChanged;
-					TreeNode[] unitNodes = new TreeNode[units.Length];
-					TreeNode temp;
-
-					for (int j = 0; j < units.Length; j++)
-					{
-						unitNodes[j] = CreateTreeNode(units[j]);
-					}
-
-					temp = new TreeNode(CreateNodeName(cat), unitNodes);
+					TreeNode[] unitNodes = CreateChildNodes(army.GetMainUnits(cat));
+					TreeNode temp = new TreeNode(CreateNodeName(cat), unitNodes);
 					temp.Tag = cat;
 					catNodes[i] = temp;
 					nodes[cat.ID] = temp;
@@ -136,7 +127,8 @@
 
 		private TreeNode CreateTreeNode(Unit unit)
 		{
-			TreeNode temp = new TreeNode(CreateNodeName(unit));
+			TreeNode[] unitNodes = CreateChildNodes(unit.ContainedUnits);
+			TreeNode temp = new TreeNode(CreateNodeName(unit), unitNodes);
 			temp.Tag = unit;
 			unit.NameChanged += UnitNameChangedMethod;
 			unit.PointsValueChanged += NodePointsValueChanged;
@@ -144,6 +136,27 @@
 			return temp;
 		}
 
+		private TreeNode[] CreateChildNodes(Unit[] containedUnits)
+		{
+			TreeNode[] unitNodes;
+
+			if (containedUnits.Length > 0)
+			{
+				unitNodes = new TreeNode[containedUnits.Length];
+
+				for (int j = 0; j < containedUnits.Length; j++)
+				{
+					unitNodes[j] = CreateTreeNode(containedUnits[j]);
+				}
+			}
+			else
+			{
+				unitNodes = new TreeNode[0];
+			}
+
+			return unitNodes;
+		}
+
 		public void FrmArmyTree_ArmyChanged(Army oldArmy, Army newArmy)
 		{
 			if (oldArmy != null)
@@ -170,8 +183,18 @@
 			if (obj is Unit)
 			{
 				Unit unit = (Unit)obj;
-				ArmyCategory cat = unit.Category;
-				TreeNode parent = nodes[cat.ID];
+				TreeNode parent;
+
+				if (unit.ParentUnit == null)
+				{
+					ArmyCategory cat = unit.Category;
+					parent = nodes[cat.ID];
+				}
+				else
+				{
+					parent = nodes[unit.ParentUnit.ID];
+				}
+
 				TreeNode unitNode = CreateTreeNode(unit);
 				parent.Nodes.Add(unitNode);
 				parent.Expand(); //make sure it's expanded
@@ -391,7 +414,7 @@
 				if (dr == DialogResult.OK)
 				{
 					UnitType newUnit = newUnitFrm.SelectedUnit;
-					CreateAndAddUnitCommand cmd = new CreateAndAddUnitCommand(newUnitFrm.SelectedUnit, WarFoundryCore.CurrentArmy.GetCategory(newUnit.MainCategory));
+					CreateAndAddUnitCommand cmd = new CreateAndAddUnitCommand(newUnitFrm.SelectedUnit, WarFoundryCore.CurrentArmy.GetCategory(newUnit.MainCategory), unit);
 					commandStack.Execute(cmd);
 					MdiParent.OpenUnitDialog(cmd.Unit);
 				}