using System;
using System.Collections.Generic;
using System.Threading.Tasks;
using ZB.MOM.WW.MxGateway.Contracts.Proto;
using ZB.MOM.WW.MxGateway.Worker.Conversion;

namespace ZB.MOM.WW.MxGateway.Worker.Sta;

public sealed class StaCommandDispatcher
{
    public const int DefaultMaxPendingCommands = 128;

    private readonly HResultConverter hresultConverter;
    private readonly IStaCommandExecutor commandExecutor;
    private readonly Queue<QueuedStaCommand> commandQueue = new();
    private readonly StaRuntime staRuntime;
    private readonly int maxPendingCommands;
    private readonly object gate = new();
    private bool drainActive;
    private bool shutdownRequested;
    private string currentCommandCorrelationId = string.Empty;

    /// <summary>
    /// Initializes a new instance of <see cref="StaCommandDispatcher"/> with default converter.
    /// </summary>
    /// <param name="staRuntime">STA thread runtime.</param>
    /// <param name="commandExecutor">Command executor.</param>
    public StaCommandDispatcher(
        StaRuntime staRuntime,
        IStaCommandExecutor commandExecutor)
        : this(staRuntime, commandExecutor, new HResultConverter())
    {
    }

    /// <summary>
    /// Initializes a new instance of <see cref="StaCommandDispatcher"/> with custom converter.
    /// </summary>
    /// <param name="staRuntime">STA thread runtime.</param>
    /// <param name="commandExecutor">Command executor.</param>
    /// <param name="hresultConverter">HResult converter.</param>
    public StaCommandDispatcher(
        StaRuntime staRuntime,
        IStaCommandExecutor commandExecutor,
        HResultConverter hresultConverter)
        : this(staRuntime, commandExecutor, hresultConverter, DefaultMaxPendingCommands)
    {
    }

    /// <summary>
    /// Initializes a new instance of <see cref="StaCommandDispatcher"/> with all parameters.
    /// </summary>
    /// <param name="staRuntime">STA thread runtime.</param>
    /// <param name="commandExecutor">Command executor.</param>
    /// <param name="hresultConverter">HResult converter.</param>
    /// <param name="maxPendingCommands">Maximum pending commands allowed.</param>
    public StaCommandDispatcher(
        StaRuntime staRuntime,
        IStaCommandExecutor commandExecutor,
        HResultConverter hresultConverter,
        int maxPendingCommands)
    {
        if (maxPendingCommands <= 0)
        {
            throw new ArgumentOutOfRangeException(
                nameof(maxPendingCommands),
                "Max pending STA commands must be greater than zero.");
        }

        this.staRuntime = staRuntime ?? throw new ArgumentNullException(nameof(staRuntime));
        this.commandExecutor = commandExecutor ?? throw new ArgumentNullException(nameof(commandExecutor));
        this.hresultConverter = hresultConverter ?? throw new ArgumentNullException(nameof(hresultConverter));
        this.maxPendingCommands = maxPendingCommands;
    }

    /// <summary>
    /// Gets the count of pending commands in the queue.
    /// </summary>
    public int PendingCommandCount
    {
        get
        {
            lock (gate)
            {
                return commandQueue.Count;
            }
        }
    }

    /// <summary>
    /// Gets the correlation ID of the currently executing command.
    /// </summary>
    public string CurrentCommandCorrelationId
    {
        get
        {
            lock (gate)
            {
                return currentCommandCorrelationId;
            }
        }
    }

    /// <summary>
    /// Dispatches a command to the queue for asynchronous STA execution.
    /// </summary>
    /// <param name="command">The command to dispatch.</param>
    /// <returns>Task for the command reply.</returns>
    public Task<MxCommandReply> DispatchAsync(StaCommand command)
    {
        if (command is null)
        {
            throw new ArgumentNullException(nameof(command));
        }

        lock (gate)
        {
            if (shutdownRequested)
            {
                return Task.FromResult(CreateRejectedReply(
                    command,
                    ProtocolStatusCode.WorkerUnavailable,
                    "The STA command dispatcher is shutting down."));
            }

            if (commandQueue.Count >= maxPendingCommands)
            {
                return Task.FromResult(CreateRejectedReply(
                    command,
                    ProtocolStatusCode.WorkerUnavailable,
                    $"The STA command dispatcher already has {maxPendingCommands} pending command(s)."));
            }

            QueuedStaCommand queuedCommand = new(command);
            commandQueue.Enqueue(queuedCommand);

            if (!drainActive)
            {
                drainActive = true;
                _ = DrainAsync();
            }

            return queuedCommand.Task;
        }
    }

    /// <summary>
    /// Cancels a queued command by its correlation ID.
    /// </summary>
    /// <param name="correlationId">Correlation ID of the command to cancel.</param>
    /// <returns>True if the command was canceled; otherwise false.</returns>
    public bool CancelQueuedCommand(string correlationId)
    {
        if (string.IsNullOrWhiteSpace(correlationId))
        {
            return false;
        }

        lock (gate)
        {
            if (commandQueue.Count == 0)
            {
                return false;
            }

            bool canceled = false;
            Queue<QueuedStaCommand> retainedCommands = new(commandQueue.Count);
            while (commandQueue.Count > 0)
            {
                QueuedStaCommand queuedCommand = commandQueue.Dequeue();
                if (!canceled
                    && string.Equals(
                        queuedCommand.Command.CorrelationId,
                        correlationId,
                        StringComparison.Ordinal))
                {
                    queuedCommand.Complete(CreateRejectedReply(
                        queuedCommand.Command,
                        ProtocolStatusCode.Canceled,
                        "The STA command was canceled before execution."));
                    canceled = true;
                    continue;
                }

                retainedCommands.Enqueue(queuedCommand);
            }

            while (retainedCommands.Count > 0)
            {
                commandQueue.Enqueue(retainedCommands.Dequeue());
            }

            return canceled;
        }
    }

    /// <summary>
    /// Requests graceful shutdown, rejecting all queued commands.
    /// </summary>
    public void RequestShutdown()
    {
        lock (gate)
        {
            shutdownRequested = true;
            while (commandQueue.Count > 0)
            {
                QueuedStaCommand queuedCommand = commandQueue.Dequeue();
                queuedCommand.Complete(CreateRejectedReply(
                    queuedCommand.Command,
                    ProtocolStatusCode.WorkerUnavailable,
                    "The STA command dispatcher is shutting down."));
            }
        }
    }

    /// <summary>
    /// Populates the given heartbeat with current dispatcher state.
    /// </summary>
    /// <param name="heartbeat">Heartbeat to populate.</param>
    public void PopulateHeartbeat(WorkerHeartbeat heartbeat)
    {
        if (heartbeat is null)
        {
            throw new ArgumentNullException(nameof(heartbeat));
        }

        lock (gate)
        {
            heartbeat.PendingCommandCount = (uint)commandQueue.Count;
            heartbeat.CurrentCommandCorrelationId = currentCommandCorrelationId;
        }
    }

    private async Task DrainAsync()
    {
        while (true)
        {
            QueuedStaCommand queuedCommand;
            lock (gate)
            {
                if (commandQueue.Count == 0)
                {
                    drainActive = false;
                    return;
                }

                queuedCommand = commandQueue.Dequeue();
            }

            await ExecuteQueuedCommandAsync(queuedCommand).ConfigureAwait(false);
        }
    }

    private async Task ExecuteQueuedCommandAsync(QueuedStaCommand queuedCommand)
    {
        StaCommand command = queuedCommand.Command;
        if (command.CancellationToken.IsCancellationRequested)
        {
            queuedCommand.Complete(CreateRejectedReply(
                command,
                ProtocolStatusCode.Canceled,
                "The STA command was canceled before execution."));
            return;
        }

        SetCurrentCommand(command.CorrelationId);
        try
        {
            MxCommandReply reply = await staRuntime
                .InvokeAsync(() => commandExecutor.Execute(command))
                .ConfigureAwait(false);

            queuedCommand.Complete(NormalizeReply(command, reply));
        }
        catch (Exception exception)
        {
            queuedCommand.Complete(CreateExceptionReply(command, exception));
        }
        finally
        {
            SetCurrentCommand(string.Empty);
        }
    }

    private void SetCurrentCommand(string correlationId)
    {
        lock (gate)
        {
            currentCommandCorrelationId = correlationId;
        }
    }

    private MxCommandReply NormalizeReply(
        StaCommand command,
        MxCommandReply reply)
    {
        if (reply is null)
        {
            return CreateRejectedReply(
                command,
                ProtocolStatusCode.ProtocolViolation,
                "STA command executor returned null.");
        }

        if (string.IsNullOrWhiteSpace(reply.SessionId))
        {
            reply.SessionId = command.SessionId;
        }

        if (string.IsNullOrWhiteSpace(reply.CorrelationId))
        {
            reply.CorrelationId = command.CorrelationId;
        }

        if (reply.Kind == MxCommandKind.Unspecified)
        {
            reply.Kind = command.Kind;
        }

        if (reply.ProtocolStatus is null)
        {
            reply.ProtocolStatus = new ProtocolStatus
            {
                Code = ProtocolStatusCode.Ok,
                Message = "OK",
            };
        }

        return reply;
    }

    private MxCommandReply CreateExceptionReply(
        StaCommand command,
        Exception exception)
    {
        HResultConversion conversion = hresultConverter.Convert(exception);
        MxCommandReply reply = CreateBaseReply(command);
        reply.ProtocolStatus = conversion.ProtocolStatus;
        reply.Hresult = conversion.HResult;
        reply.DiagnosticMessage = conversion.DiagnosticMessage;

        return reply;
    }

    private static MxCommandReply CreateRejectedReply(
        StaCommand command,
        ProtocolStatusCode statusCode,
        string message)
    {
        MxCommandReply reply = CreateBaseReply(command);
        reply.ProtocolStatus = new ProtocolStatus
        {
            Code = statusCode,
            Message = message,
        };
        reply.DiagnosticMessage = message;

        return reply;
    }

    private static MxCommandReply CreateBaseReply(StaCommand command)
    {
        return new MxCommandReply
        {
            SessionId = command.SessionId,
            CorrelationId = command.CorrelationId,
            Kind = command.Kind,
        };
    }

    private sealed class QueuedStaCommand
    {
        private readonly TaskCompletionSource<MxCommandReply> completion = new(
            TaskCreationOptions.RunContinuationsAsynchronously);

        /// <summary>
        /// Initializes a new instance of <see cref="QueuedStaCommand"/>.
        /// </summary>
        /// <param name="command">The STA command to queue.</param>
        public QueuedStaCommand(StaCommand command)
        {
            Command = command;
        }

        /// <summary>
        /// Gets the queued STA command.
        /// </summary>
        public StaCommand Command { get; }

        /// <summary>
        /// Gets the task representing the command's completion.
        /// </summary>
        public Task<MxCommandReply> Task => completion.Task;

        /// <summary>
        /// Completes the command with the given reply.
        /// </summary>
        /// <param name="reply">The command reply.</param>
        public void Complete(MxCommandReply reply)
        {
            completion.TrySetResult(reply);
        }
    }
}
