BITKit/Src/Core/StateMachine/AsyncStateMachine.cs

using System;
using System.Collections.Generic;
using System.Threading;
using System.Threading.Tasks;
using Cysharp.Threading.Tasks;
using Microsoft.Extensions.DependencyInjection;

namespace BITKit.StateMachine
{
    /// <summary>
    /// 动态异步状态机
    /// </summary>
    /// <typeparam name="T">异步状态</typeparam>
    public class AsyncStateMachine<T>:IStateMachine<T> ,IDisposable where T  : class, IStateAsync
    {
        private readonly IServiceProvider _serviceProvider;
        public AsyncStateMachine(IServiceProvider serviceProvider)
        {
            _serviceProvider = serviceProvider;
        }
        public AsyncStateMachine()
        {
            _serviceProvider = new ServiceCollection().BuildServiceProvider();
        }

        public bool Enabled { get; set; } = true;
        public T CurrentState { get;private set; }
        public event Action<T, T> OnStateChanging;
        public event Func<T, T, UniTask> OnStateChangeAsync; 
        public event Action<T, T> OnStateChanged;
        public event Action<T> OnStateRegistered;
        public IReadOnlyDictionary<int, T> Dictionary => _dictionary;
        public event Action<T> OnStateUnRegistered;
        public IDictionary<Type, T> StateDictionary { get; } = new Dictionary<Type, T>();
        public readonly ValidHandle IsBusy=new();
        private readonly CancellationTokenSource _cancellationTokenSource=new();
        private readonly Dictionary<int, T> _dictionary = new();
        
        private readonly HashSet<int> _isRegistered = new();

        private CancellationTokenSource _transitionCts;

        private T _entryCompletedState;
        
        public async void Initialize()
        {
            await IsBusy;
            if(_cancellationTokenSource.IsCancellationRequested)return;
            using var _ = IsBusy.GetHandle();
            foreach (var (_,value) in StateDictionary)
            {
                if (_isRegistered.Add(value.Identifier))
                {
                    await value.InitializeAsync();
                    value.Initialize();
                }
            }
        }


        public async void UpdateState(float deltaTime)
        {
            if (_entryCompletedState is null) return;
            using var _ = IsBusy.GetHandle();
            _entryCompletedState.OnStateUpdate(deltaTime);
            await _entryCompletedState.OnStateUpdateAsync(deltaTime);
        }

        public void DisposeState()
        {
            TransitionState(null);
        }
        public T TransitionState<TState>() where TState : T
        {
            T nextState;
            foreach (var (type, value) in Dictionary)
            {
                if ((typeof(TState) == value.GetType()) is false) continue;
                
                nextState = value;
                
                TransitionState(nextState);
                return nextState;
            }
            nextState = _serviceProvider.GetRequiredService<TState>();

            return  TransitionState(nextState);;
        }

        public async UniTask TransitionStateAsync(T nextState)
        {
            if (nextState is not null)
            {
                if (nextState.Identifier == 0)
                {
                    nextState.Identifier = nextState.GetHashCode();
                }
            }
            
            if (Equals(nextState, CurrentState)) return;
            
            var tempState = CurrentState;
            
            CurrentState = nextState;
            
            _transitionCts?.Cancel();
            _transitionCts = new CancellationTokenSource();

            var ct = _transitionCts.Token;
            
            await IsBusy;
            using var _ = IsBusy.GetHandle();
            if(ct.IsCancellationRequested||_cancellationTokenSource.IsCancellationRequested)return;
            
            OnStateChanging?.Invoke(tempState,nextState);

            if (tempState is not null)
            {
                if (_entryCompletedState == tempState)
                {
                    _entryCompletedState = null;
                }
                
                tempState.Enabled = false;
                await tempState.OnStateExitAsync(tempState, nextState);
                tempState.OnStateExit(tempState,nextState);
                if(_cancellationTokenSource.IsCancellationRequested)return;
            }
            
            if(ct.IsCancellationRequested)return;
            
            await OnStateChangeAsync.UniTaskFunc(CurrentState,nextState);
    
            if(ct.IsCancellationRequested)return;
            
            if (nextState is not null)
            {
                if (_isRegistered.Add(nextState.Identifier))
                {
                    await RegisterAsync(nextState);
                   if(ct.IsCancellationRequested || _cancellationTokenSource.IsCancellationRequested)return;
                }
                
                nextState.Enabled = true;
                await nextState.OnStateEntryAsync(CurrentState);
                nextState.OnStateEntry(CurrentState);
                if(ct.IsCancellationRequested || _cancellationTokenSource.IsCancellationRequested)return;

                _entryCompletedState = nextState;
            }
       
            OnStateChanged?.Invoke(tempState, nextState);
    
        }
        public T TransitionState(T nextState)
        {
            TransitionStateAsync(nextState).Forget();
            return nextState;
        }

        private async UniTask RegisterAsync(T newState)
        {
            StateDictionary.TryAdd(newState.GetType(),newState);

            _dictionary.TryAdd(newState.Identifier, newState);
   

            newState.Initialize();
            await newState.InitializeAsync();
        }

        public async void Register(T newState)
        {
            await IsBusy;
            using var _ = IsBusy.GetHandle();
            await RegisterAsync(newState);
        }
        public async void UnRegister(T newState)
        {
            if (newState is null) return;
            if (Dictionary.ContainsKey(newState.Identifier) is false) return;
            _dictionary.Remove(newState.Identifier);

            await IsBusy;
            using var _ = IsBusy.GetHandle();

            if (Equals(CurrentState, newState))
            {
                await CurrentState.OnStateExitAsync(CurrentState, null);
                CurrentState.OnStateExit(CurrentState, null);
         
                CurrentState = null;
            }
            
            if (newState is IAsyncDisposable asyncDisposable)
            {
                await asyncDisposable.DisposeAsync();
            }
            if (newState is IDisposable disposable)
            {
                disposable.Dispose();
            }
            
            OnStateUnRegistered?.Invoke(newState);
        }
        public void Dispose()
        {
            if(_isDisposed)return;
            _cancellationTokenSource.Cancel();
            _cancellationTokenSource.Dispose();
            _isDisposed = true;
        }

        private bool _isDisposed;
    }

}