1

Src/Unity/Scripts/Entity/Core/StateBasedComponent.cs

@@ -1,6 +1,7 @@
 using System;
 using System.Collections;
 using System.Collections.Generic;
+using System.Linq;
 using BITKit.Entities;
 using BITKit.StateMachine;
 using UnityEngine;
@@ -27,20 +28,27 @@ namespace BITKit.Entities
 			add => stateMachine.OnStateChanged += value;
 			remove => stateMachine.OnStateChanged -= value;
 		}
+
+		public event Action<T> OnStateRegistered;
+		public event Action<T> OnStateUnRegistered;
 		public IDictionary<Type, T> StateDictionary => stateMachine.StateDictionary;
-		public override void Initialize(IEntity entity)
+		public override void OnAwake()
 		{
-			base.Initialize(entity);
+			base.OnAwake();
 			if (stateMachine is null)
 			{
 				Debug.LogWarning(GetType().Name);
 			}
 			else
 			{
+				foreach (var x in stateMachine.states)
+				{
+					Entity.Inject(x);
+				}
 				stateMachine.Initialize();
 			}
 		}
-        void IStateMachine<T>.Initialize()
+		void IStateMachine<T>.Initialize()
 		{
 			stateMachine.Initialize();
 		}
@@ -62,6 +70,23 @@ namespace BITKit.Entities
 		{
 			stateMachine.TransitionState(state);
 		}
+		
+		public virtual void Register(T newState)
+		{
+			Entity.Inject(newState);
+			StateMachineUtils.Register(this, newState);
+		}
+
+		public virtual void UnRegister(T newState)=>StateMachineUtils.UnRegister(this,newState);
+		public void InvokeOnStateRegistered(T state)
+		{
+			OnStateRegistered?.Invoke(state);
+		}
+
+		public void InvokeOnStateUnRegistered(T state)
+		{
+			OnStateUnRegistered?.Invoke(state);
+		}
 	}
 
 }