using System.Collections.Concurrent;
using System.Threading.Channels;
using ZB.MOM.WW.MxGateway.Contracts.Proto;

namespace ZB.MOM.WW.MxGateway.Server.Sessions;

/// <summary>
///     Per-session event pump and fan-out. A single background task drains the
///     session's event source <em>exactly once</em> and fans each event out to
///     every currently-registered subscriber's own bounded channel.
/// </summary>
/// <remarks>
///     <para>
///         This is the skeleton introduced by Task 2 of the Session Resilience epic.
///         It is a standalone class — it is NOT yet wired into <c>GatewaySession</c> or
///         <c>EventStreamService</c> (Task 4), it has no replay ring buffer (Task 3),
///         no per-subscriber backpressure-isolation policy (Task 5), and it does not
///         remove the single-subscriber guard (Tasks 7/8).
///     </para>
///     <para>
///         <b>Source seam.</b> The event source is injected as a
///         <see cref="Func{T, TResult}"/> producing an
///         <see cref="IAsyncEnumerable{T}"/> of already-mapped public
///         <see cref="MxEvent"/>s, given a <see cref="CancellationToken"/>. This is the
///         cleanest seam for Task 4: it can pass
///         <c>ct =&gt; session.ReadEventsAsync(ct).Select(mapper.MapEvent)</c> (or a
///         channel reader's <c>ReadAllAsync</c>), while unit tests pass a plain
///         channel reader's <c>ReadAllAsync</c> with no real session. The pump owns the
///         single consumption of this enumerable; fan-out happens on the public
///         <see cref="MxEvent"/> after mapping, mirroring today's
///         <c>EventStreamService.ProduceEventsAsync</c> ordering.
///     </para>
///     <para>
///         <b>Concurrency.</b> The subscriber set is a
///         <see cref="ConcurrentDictionary{TKey, TValue}"/> keyed by a monotonic id.
///         The pump iterates it with a snapshot-free enumerator (which never throws on
///         concurrent add/remove), and <see cref="Register"/> / lease disposal mutate it
///         without any lock held across an <c>await</c>. Each subscriber channel has a
///         single writer — the pump — so per-channel writes never race. MXAccess parity:
///         events are fanned in the order received; the pump never reorders or
///         synthesizes events.
///     </para>
/// </remarks>
public sealed class SessionEventDistributor : IAsyncDisposable
{
    /// <summary>
    ///     Bounded wait for the pump to stop during disposal. A source factory that
    ///     ignores cancellation must not hang dispose forever; after this window the
    ///     pump is abandoned and subscribers are completed anyway.
    /// </summary>
    private static readonly TimeSpan DefaultShutdownTimeout = TimeSpan.FromSeconds(5);

    private readonly string _sessionId;
    private readonly Func<CancellationToken, IAsyncEnumerable<MxEvent>> _eventSourceFactory;
    private readonly int _subscriberQueueCapacity;
    private readonly TimeSpan _shutdownTimeout;
    private readonly ILogger<SessionEventDistributor> _logger;
    private readonly ConcurrentDictionary<long, Subscriber> _subscribers = new();
    private readonly CancellationTokenSource _shutdownCts = new();
    private readonly object _lifecycleLock = new();

    private long _nextSubscriberId;
    private Task? _pumpTask;
    private bool _started;
    private bool _disposed;

    /// <summary>
    /// Initializes a per-session event distributor.
    /// </summary>
    /// <param name="sessionId">Owning session id, used only for logging context.</param>
    /// <param name="eventSourceFactory">
    ///     Factory producing the session's event stream given a cancellation token.
    ///     The pump consumes this exactly once. See the type remarks for the seam Task 4
    ///     plugs into.
    /// </param>
    /// <param name="subscriberQueueCapacity">
    ///     Bounded capacity of each per-subscriber channel. Mirrors the gRPC event-stream
    ///     queue capacity shape used today.
    /// </param>
    /// <param name="logger">Logger for pump lifecycle diagnostics.</param>
    public SessionEventDistributor(
        string sessionId,
        Func<CancellationToken, IAsyncEnumerable<MxEvent>> eventSourceFactory,
        int subscriberQueueCapacity,
        ILogger<SessionEventDistributor> logger)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(sessionId);
        ArgumentNullException.ThrowIfNull(eventSourceFactory);
        ArgumentOutOfRangeException.ThrowIfLessThan(subscriberQueueCapacity, 1);
        ArgumentNullException.ThrowIfNull(logger);

        _sessionId = sessionId;
        _eventSourceFactory = eventSourceFactory;
        _subscriberQueueCapacity = subscriberQueueCapacity;
        _shutdownTimeout = DefaultShutdownTimeout;
        _logger = logger;
    }

    /// <summary>
    /// Gets the count of currently-registered subscribers.
    /// </summary>
    public int SubscriberCount => _subscribers.Count;

    /// <summary>
    /// Starts the background pump. Idempotent — a second call is a no-op.
    /// </summary>
    /// <param name="cancellationToken">Token observed only while starting.</param>
    public Task StartAsync(CancellationToken cancellationToken)
    {
        cancellationToken.ThrowIfCancellationRequested();

        lock (_lifecycleLock)
        {
            ObjectDisposedException.ThrowIf(_disposed, this);
            if (_started)
            {
                return Task.CompletedTask;
            }

            _started = true;
            _pumpTask = Task.Run(() => PumpAsync(_shutdownCts.Token), CancellationToken.None);
        }

        return Task.CompletedTask;
    }

    /// <summary>
    ///     Registers a new subscriber and returns its lease. The lease exposes the
    ///     subscriber's <see cref="ChannelReader{T}"/> and, when disposed, unregisters the
    ///     subscriber and completes its channel without disturbing the pump or other
    ///     subscribers.
    /// </summary>
    public IEventSubscriberLease Register()
    {
        // The pump is the single writer for this channel; readers are single-consumer
        // (one gRPC stream / dashboard subscriber). Synchronous continuations are
        // disabled so a slow reader can never stall the pump on its completion.
        //
        // FullMode is Wait but the pump currently writes with TryWrite (drop-on-full):
        // these are deliberately opposite policies and only a placeholder. Task 5 owns
        // the overflow policy and will reconcile them by either switching the pump to
        // WriteAsync (true backpressure, honouring Wait) or changing this to a Drop mode.
        // Do not "fix" the mismatch here — leave the decision to Task 5.
        Channel<MxEvent> channel = Channel.CreateBounded<MxEvent>(
            new BoundedChannelOptions(_subscriberQueueCapacity)
            {
                SingleReader = true,
                SingleWriter = true,
                FullMode = BoundedChannelFullMode.Wait,
                AllowSynchronousContinuations = false,
            });

        long id = Interlocked.Increment(ref _nextSubscriberId);
        Subscriber subscriber = new(id, channel);

        // The disposed check AND the map add happen under the same lock with no await
        // in between. DisposeAsync sets _disposed=true under this same lock before it
        // calls CompleteAllSubscribers, so once disposal has begun no further subscriber
        // can be added — closing the Register-after-DisposeAsync window that would
        // otherwise leave a subscriber's channel never completed.
        lock (_lifecycleLock)
        {
            ObjectDisposedException.ThrowIf(_disposed, this);
            _subscribers[id] = subscriber;
        }

        return new SubscriberLease(this, subscriber);
    }

    /// <summary>
    /// Stops the pump and completes all subscriber channels. Idempotent.
    /// </summary>
    public async ValueTask DisposeAsync()
    {
        Task? pumpTask;
        lock (_lifecycleLock)
        {
            if (_disposed)
            {
                return;
            }

            _disposed = true;
            pumpTask = _pumpTask;
        }

        // Signal the pump to stop. It must not block on a non-reading subscriber:
        // it writes with non-blocking TryWrite, so cancellation tears it down promptly.
        await _shutdownCts.CancelAsync().ConfigureAwait(false);

        if (pumpTask is not null)
        {
            // Bound the wait: a source factory that ignores cancellation would otherwise
            // hang dispose forever. If the pump does not stop in time we log and proceed
            // to complete subscribers anyway; DisposeAsync must not throw on this path.
            Task completed = await Task.WhenAny(pumpTask, Task.Delay(_shutdownTimeout)).ConfigureAwait(false);
            if (!ReferenceEquals(completed, pumpTask))
            {
                _logger.LogWarning(
                    "Event distributor pump did not stop within {ShutdownTimeoutSeconds}s for session {SessionId}; completing subscribers and abandoning the pump.",
                    _shutdownTimeout.TotalSeconds,
                    _sessionId);
            }
            else
            {
                try
                {
                    await pumpTask.ConfigureAwait(false);
                }
                catch (OperationCanceledException)
                {
                }
                catch (Exception exception)
                {
                    _logger.LogDebug(
                        exception,
                        "Event distributor pump faulted during shutdown for session {SessionId}.",
                        _sessionId);
                }
            }
        }

        CompleteAllSubscribers(error: null);
        _shutdownCts.Dispose();
    }

    private async Task PumpAsync(CancellationToken cancellationToken)
    {
        try
        {
            await foreach (MxEvent mxEvent in _eventSourceFactory(cancellationToken)
                .WithCancellation(cancellationToken)
                .ConfigureAwait(false))
            {
                // Enumerating a ConcurrentDictionary's Values never throws on concurrent
                // add/remove; a subscriber registered mid-iteration may miss this event,
                // which matches "late subscribers see events after they register".
                foreach (Subscriber subscriber in _subscribers.Values)
                {
                    // TODO(Task 5): define overflow policy (per-subscriber isolation —
                    // drop / disconnect / fault that one subscriber). For the Task 2
                    // skeleton, a non-blocking TryWrite that silently drops on a full
                    // channel is the placeholder so one slow reader never stalls the pump.
                    if (!subscriber.Channel.Writer.TryWrite(mxEvent))
                    {
                        // Visibility only — Task 5 owns the actual drop/backpressure policy.
                        // Logs identifiers (worker sequence, subscriber id, session) only,
                        // never the event payload or tag values.
                        _logger.LogDebug(
                            "Event distributor dropped event (worker sequence {WorkerSequence}) for subscriber {SubscriberId} in session {SessionId}: channel full.",
                            mxEvent.WorkerSequence,
                            subscriber.Id,
                            _sessionId);
                    }
                }
            }

            CompleteAllSubscribers(error: null);
        }
        catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
        {
            // Shutdown path: DisposeAsync completes subscribers.
        }
        catch (Exception exception)
        {
            // Unexpected source fault (not the shutdown-cancellation path above) — visible
            // by default so an event stream silently dying is not lost in Debug noise.
            _logger.LogError(
                exception,
                "Event distributor source faulted for session {SessionId}.",
                _sessionId);
            CompleteAllSubscribers(exception);
        }
    }

    private void CompleteAllSubscribers(Exception? error)
    {
        foreach (Subscriber subscriber in _subscribers.Values)
        {
            subscriber.Channel.Writer.TryComplete(error);
        }
    }

    private void Unregister(Subscriber subscriber)
    {
        if (_subscribers.TryRemove(subscriber.Id, out _))
        {
            subscriber.Channel.Writer.TryComplete();
        }
    }

    private sealed class Subscriber(long id, Channel<MxEvent> channel)
    {
        public long Id { get; } = id;

        public Channel<MxEvent> Channel { get; } = channel;
    }

    private sealed class SubscriberLease(SessionEventDistributor distributor, Subscriber subscriber)
        : IEventSubscriberLease
    {
        private int _leaseDisposed;

        public ChannelReader<MxEvent> Reader => subscriber.Channel.Reader;

        public void Dispose()
        {
            // Atomic check-and-set so concurrent Dispose calls unregister at most once.
            if (Interlocked.Exchange(ref _leaseDisposed, 1) == 0)
            {
                distributor.Unregister(subscriber);
            }
        }
    }
}
