zcbot/windows-node/Zcbot.WindowsNode/Application/NodeApplicationController.cs

namespace Zcbot.WindowsNode;

internal enum NodeShutdownMode
{
    WaitForJobs,
    CancelJobs,
}

internal interface INodeConnection
{
    Task RunAsync(CancellationToken cancellationToken);
    void CancelActiveJobsForExit();
}

internal sealed class NodeApplicationController : IDisposable
{
    private readonly INodeConfigStore store;
    private readonly NodePaths paths;
    private readonly Func<NodeConfig, Action<NodeStatus>, INodeConnection> connectionFactory;
    private readonly Func<EnrollOptions, INodeConfigStore, CancellationToken, Task> enroll;
    private readonly Func<string, string, CancellationToken, Task<NodeDataMigrationResult>> migrate;
    private readonly Action<string> saveDataRoot;
    private readonly SemaphoreSlim lifecycleGate = new(1, 1);
    private readonly object stopSync = new();
    private CancellationTokenSource? connectionStop;
    private Task? connectionTask;
    private INodeConnection? connection;
    private Task? stopTask;
    private bool disposed;

    internal NodeApplicationController(
        INodeConfigStore store,
        NodePaths paths,
        Func<NodeConfig, Action<NodeStatus>, INodeConnection> connectionFactory,
        Func<EnrollOptions, INodeConfigStore, CancellationToken, Task> enroll,
        Func<string, string, CancellationToken, Task<NodeDataMigrationResult>> migrate,
        Action<string> saveDataRoot)
    {
        this.store = store;
        this.paths = paths;
        this.connectionFactory = connectionFactory;
        this.enroll = enroll;
        this.migrate = migrate;
        this.saveDataRoot = saveDataRoot;
    }

    internal event Action<NodeStatus>? StatusChanged;

    internal NodeConfig? CurrentConfig { get; private set; }

    internal NodeStatus CurrentStatus { get; private set; } =
        NodeStatus.Create(NodeState.NotRegistered, "尚未注册");

    internal int ActiveJobCount => new JobMonitorService(paths.JobsDirectory)
        .ReadSnapshots()
        .Count(item => item.IsActive);

    internal NodePaths Paths => paths;

    internal void Start()
    {
        ObjectDisposedException.ThrowIf(disposed, this);
        if (!store.Exists)
        {
            Publish(CurrentStatus);
            return;
        }

        try
        {
            CurrentConfig = store.Load();
            Publish(NodeStatus.Create(NodeState.Connecting, "正在连接 zcbot…"));
            StartConnection();
        }
        catch (NodeConfigurationException exception)
        {
            Publish(NodeStatus.Create(NodeState.AuthenticationRequired, exception.Message));
        }
    }

    internal async Task RegisterAsync(
        EnrollOptions options,
        CancellationToken cancellationToken = default)
    {
        ObjectDisposedException.ThrowIf(disposed, this);
        await lifecycleGate.WaitAsync(cancellationToken);
        try
        {
            Publish(NodeStatus.Create(NodeState.Connecting, "正在注册节点…"));
            try
            {
                await enroll(options, store, cancellationToken);
                CurrentConfig = store.Load();
                Publish(NodeStatus.Create(NodeState.Connecting, "注册成功，正在连接…"));
                StartConnection();
            }
            catch
            {
                Publish(NodeStatus.Create(NodeState.NotRegistered, "注册失败，请检查配置"));
                throw;
            }
        }
        finally
        {
            lifecycleGate.Release();
        }
    }

    internal async Task<bool> ReconnectAsync(CancellationToken cancellationToken = default)
    {
        ObjectDisposedException.ThrowIf(disposed, this);
        await lifecycleGate.WaitAsync(cancellationToken);
        try
        {
            if (CurrentConfig is null)
            {
                return false;
            }

            Publish(NodeStatus.Create(NodeState.Connecting, "等待本机任务收尾后重连…"));
            await StopConnectionAsync();
            cancellationToken.ThrowIfCancellationRequested();
            StartConnection();
            return true;
        }
        finally
        {
            lifecycleGate.Release();
        }
    }

    internal async Task ResetIdentityAsync(CancellationToken cancellationToken = default)
    {
        ObjectDisposedException.ThrowIf(disposed, this);
        await lifecycleGate.WaitAsync(cancellationToken);
        try
        {
            await StopConnectionAsync();
            store.DeleteLocalIdentity();
            CurrentConfig = null;
            Publish(NodeStatus.Create(
                NodeState.NotRegistered,
                "本机身份已清除，请使用新注册码注册"));
        }
        finally
        {
            lifecycleGate.Release();
        }
    }

    internal async Task<NodeDataMigrationResult> MigrateDataRootAsync(
        string targetRoot,
        CancellationToken cancellationToken = default)
    {
        ObjectDisposedException.ThrowIf(disposed, this);
        await lifecycleGate.WaitAsync(cancellationToken);
        try
        {
            await StopConnectionAsync();
            try
            {
                var result = await migrate(paths.RootDirectory, targetRoot, cancellationToken);
                saveDataRoot(result.TargetDirectory);
                return result;
            }
            catch
            {
                StartConnection();
                throw;
            }
        }
        finally
        {
            lifecycleGate.Release();
        }
    }

    internal Task StopAsync(NodeShutdownMode mode)
    {
        lock (stopSync)
        {
            ObjectDisposedException.ThrowIf(disposed, this);
            return stopTask ??= StopOnceAsync(mode);
        }
    }

    private async Task StopOnceAsync(NodeShutdownMode mode)
    {
        await lifecycleGate.WaitAsync();
        try
        {
            var activeJobs = ActiveJobCount;
            Publish(NodeStatus.Create(
                NodeState.Stopped,
                mode == NodeShutdownMode.CancelJobs
                    ? "正在取消本机任务并退出…"
                    : activeJobs > 0 ? "等待本机任务完成后退出…" : "正在退出…"));
            if (mode == NodeShutdownMode.CancelJobs)
            {
                connection?.CancelActiveJobsForExit();
            }
            await StopConnectionAsync();
        }
        finally
        {
            lifecycleGate.Release();
        }
    }

    private void StartConnection()
    {
        if (CurrentConfig is null || connectionTask is { IsCompleted: false })
        {
            return;
        }

        connectionStop?.Dispose();
        connectionStop = new CancellationTokenSource();
        connection = connectionFactory(CurrentConfig, Publish);
        connectionTask = RunConnectionAsync(connection, connectionStop.Token);
    }

    private async Task RunConnectionAsync(
        INodeConnection activeConnection,
        CancellationToken cancellationToken)
    {
        try
        {
            await activeConnection.RunAsync(cancellationToken);
        }
        catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
        {
            Publish(NodeStatus.Create(NodeState.Stopped, "连接已停止"));
        }
        catch (NodeConfigurationException exception)
        {
            Publish(NodeStatus.Create(NodeState.AuthenticationRequired, exception.Message));
        }
        catch (Exception exception)
        {
            Publish(NodeStatus.Create(NodeState.Offline, $"节点异常：{exception.Message}"));
        }
        finally
        {
            if (ReferenceEquals(connection, activeConnection))
            {
                connection = null;
            }
        }
    }

    private async Task StopConnectionAsync()
    {
        connectionStop?.Cancel();
        if (connectionTask is not null)
        {
            await connectionTask;
        }
        connectionTask = null;
        connection = null;
    }

    private void Publish(NodeStatus status)
    {
        CurrentStatus = status;
        StatusChanged?.Invoke(status);
    }

    public void Dispose()
    {
        if (disposed)
        {
            return;
        }
        disposed = true;
        connectionStop?.Cancel();
        connectionStop?.Dispose();
        lifecycleGate.Dispose();
    }
}