using System.Diagnostics.CodeAnalysis;
using System.Runtime.ExceptionServices;
using System.Security.Cryptography;
using Google.Protobuf.WellKnownTypes;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Logging.Abstractions;
using Microsoft.Extensions.Options;
using ZB.MOM.WW.MxGateway.Contracts;
using ZB.MOM.WW.MxGateway.Contracts.Proto;
using ZB.MOM.WW.MxGateway.Server.Configuration;
using ZB.MOM.WW.MxGateway.Server.Metrics;
using ZB.MOM.WW.MxGateway.Server.Workers;

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

public sealed class SessionManager : ISessionManager
{
    public const string DefaultCloseReason = "client-close";
    public const string GatewayShutdownReason = "gateway-shutdown";
    public const string LeaseExpiredReason = "lease-expired";
    public const string DetachGraceExpiredReason = "detach-grace-expired";
    public const string FaultedReason = "faulted-reaped";

    // Bounded so a mass expiry (or a host stop with a full registry) cannot stampede
    // worker-process teardown: every concurrent close is one x86 worker being shut down or
    // killed, and the point of the fan-out is to hide a few hung workers, not to tear the whole
    // registry down at once.
    private const int MaxParallelSessionCloses = 4;

    private readonly ISessionRegistry _registry;
    private readonly ISessionWorkerClientFactory _workerClientFactory;
    private readonly GatewayMetrics _metrics;
    private readonly TimeProvider _timeProvider;
    private readonly ILogger<SessionManager> _logger;
    private readonly GatewayOptions _options;
    private readonly SemaphoreSlim _sessionSlots;
    private readonly Grpc.MxAccessGrpcMapper _eventMapper;
    private readonly ILogger<SessionEventDistributor> _distributorLogger;
    private readonly Dashboard.Hubs.IDashboardEventBroadcaster? _dashboardEventBroadcaster;
    private readonly ArrayAddressNormalizer? _addressNormalizer;

    /// <summary>
    /// Initializes a new instance of <see cref="SessionManager"/>.
    /// </summary>
    /// <param name="registry">Session registry.</param>
    /// <param name="workerClientFactory">Worker client factory.</param>
    /// <param name="options">Gateway options.</param>
    /// <param name="metrics">Gateway metrics.</param>
    /// <param name="timeProvider">Time provider for timestamps.</param>
    /// <param name="logger">Logger.</param>
    /// <param name="eventMapper">Mapper used by each session's event distributor to map worker events to public events.</param>
    /// <param name="distributorLogger">Logger passed to each session's event distributor pump.</param>
    /// <param name="dashboardEventBroadcaster">
    ///     Dashboard SignalR fan-out sink. Each session registers an internal distributor
    ///     subscriber that mirrors raw session events to this broadcaster, so the
    ///     dashboard receives events regardless of whether a gRPC client is streaming. Null in
    ///     unit tests that do not exercise the dashboard mirror.
    /// </param>
    /// <param name="addressNormalizer">
    ///     Rewrites bare array AddItem addresses to their writable <c>[]</c> form using Galaxy
    ///     metadata; handed to each session so the normalization runs at the outbound choke point.
    ///     Null in unit tests that do not exercise array-write ergonomics.
    /// </param>
    public SessionManager(
        ISessionRegistry registry,
        ISessionWorkerClientFactory workerClientFactory,
        IOptions<GatewayOptions> options,
        GatewayMetrics metrics,
        TimeProvider? timeProvider = null,
        ILogger<SessionManager>? logger = null,
        Grpc.MxAccessGrpcMapper? eventMapper = null,
        ILogger<SessionEventDistributor>? distributorLogger = null,
        Dashboard.Hubs.IDashboardEventBroadcaster? dashboardEventBroadcaster = null,
        ArrayAddressNormalizer? addressNormalizer = null)
    {
        _registry = registry ?? throw new ArgumentNullException(nameof(registry));
        _workerClientFactory = workerClientFactory ?? throw new ArgumentNullException(nameof(workerClientFactory));
        ArgumentNullException.ThrowIfNull(options);
        _metrics = metrics ?? throw new ArgumentNullException(nameof(metrics));
        _timeProvider = timeProvider ?? TimeProvider.System;
        _logger = logger ?? NullLogger<SessionManager>.Instance;
        _eventMapper = eventMapper ?? new Grpc.MxAccessGrpcMapper();
        _distributorLogger = distributorLogger ?? NullLogger<SessionEventDistributor>.Instance;
        _dashboardEventBroadcaster = dashboardEventBroadcaster;
        _addressNormalizer = addressNormalizer;
        _options = options.Value;
        _sessionSlots = new SemaphoreSlim(_options.Sessions.MaxSessions, _options.Sessions.MaxSessions);
    }

    /// <inheritdoc />
    public Task<GatewaySession> OpenSessionAsync(
        SessionOpenRequest request,
        string? clientIdentity,
        string? ownerKeyId,
        CancellationToken cancellationToken)
        => OpenSessionAsync(request, clientIdentity, ownerKeyId, ownerDashboardTags: null, cancellationToken);

    /// <inheritdoc />
    public async Task<GatewaySession> OpenSessionAsync(
        SessionOpenRequest request,
        string? clientIdentity,
        string? ownerKeyId,
        IReadOnlyList<string>? ownerDashboardTags,
        CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(request);
        EnsureSessionCapacity();

        GatewaySession? session = null;
        bool sessionOpenedRecorded = false;
        try
        {
            session = CreateSession(request, clientIdentity, ownerKeyId, ownerDashboardTags);
            if (!_registry.TryAdd(session))
            {
                throw new SessionManagerException(
                    SessionManagerErrorCode.OpenFailed,
                    $"Session id collision while opening session {session.SessionId}.");
            }

            session.TransitionTo(SessionState.StartingWorker);
            IWorkerClient workerClient = await _workerClientFactory
                .CreateAsync(session, cancellationToken)
                .ConfigureAwait(false);

            session.AttachWorkerClient(workerClient);
            session.MarkReady();
            _metrics.SessionOpened();
            sessionOpenedRecorded = true;

            return session;
        }
        catch (Exception exception)
        {
            session?.MarkFaulted(exception.Message);
            if (session is not null)
            {
                _registry.TryRemove(session.SessionId, out _);
                await session.DisposeAsync().ConfigureAwait(false);
            }

            // If SessionOpened() already incremented the open-session gauge,
            // a failure after that point (e.g. auto-subscribe rejection) must
            // decrement it again so mxgateway.sessions.open does not leak.
            if (sessionOpenedRecorded)
            {
                _metrics.SessionRemoved();
            }

            ReleaseSessionSlot();
            _metrics.Fault(SessionManagerErrorCode.OpenFailed.ToString());
            _logger.LogWarning(
                exception,
                "Failed to open gateway session {SessionId}.",
                session?.SessionId ?? "<not-created>");

            throw new SessionManagerException(
                SessionManagerErrorCode.OpenFailed,
                session is null ? "Failed to create session." : $"Failed to open session {session.SessionId}.",
                exception);
        }
    }

    /// <inheritdoc />
    public bool TryGetSession(
        string sessionId,
        [MaybeNullWhen(false)] out GatewaySession session)
    {
        return _registry.TryGet(sessionId, out session);
    }

    /// <inheritdoc />
    public async Task<WorkerCommandReply> InvokeAsync(
        string sessionId,
        WorkerCommand command,
        CancellationToken cancellationToken)
    {
        GatewaySession session = GetRequiredSession(sessionId);

        try
        {
            return await session.InvokeAsync(command, cancellationToken).ConfigureAwait(false);
        }
        catch (SessionManagerException)
        {
            throw;
        }
        catch (Exception exception)
        {
            if (session.WorkerClient?.State == WorkerClientState.Faulted)
            {
                session.MarkFaulted(exception.Message);
            }

            throw;
        }
    }

    /// <inheritdoc />
    public async Task<SessionCloseResult> CloseSessionAsync(
        string sessionId,
        CancellationToken cancellationToken)
    {
        GatewaySession session = GetRequiredSession(sessionId);
        SessionCloseResult result = await CloseSessionCoreAsync(
            session,
            DefaultCloseReason,
            cancellationToken).ConfigureAwait(false);

        return result;
    }

    /// <inheritdoc />
    public async Task<SessionCloseResult> KillWorkerAsync(
        string sessionId,
        string reason,
        CancellationToken cancellationToken)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(reason);
        cancellationToken.ThrowIfCancellationRequested();

        GatewaySession session = GetRequiredSession(sessionId);

        bool wasClosed;
        try
        {
            wasClosed = await session.KillWorkerWithCloseGateAsync(reason, cancellationToken).ConfigureAwait(false);
        }
        catch (Exception exception)
        {
            session.MarkFaulted(exception.Message);
            _metrics.Fault(SessionManagerErrorCode.CloseFailed.ToString());

            // The open-session gauge was incremented in OpenSessionAsync; every
            // session reaching KillWorkerAsync had SessionOpened recorded. If the
            // kill path throws, decrement the gauge here so mxgateway.sessions.open
            // does not leak — mirroring the equivalent guard in OpenSessionAsync.
            _metrics.SessionRemoved();
            await RemoveSessionAsync(session).ConfigureAwait(false);
            throw new SessionManagerException(
                SessionManagerErrorCode.CloseFailed,
                $"Failed to kill worker for session {sessionId}.",
                exception);
        }

        if (!wasClosed)
        {
            _metrics.SessionClosed();
        }

        await RemoveSessionAsync(session).ConfigureAwait(false);
        _logger.LogInformation(
            "Worker for session {SessionId} killed; reason={Reason}.",
            sessionId,
            reason);

        return new SessionCloseResult(sessionId, SessionState.Closed, AlreadyClosed: wasClosed);
    }

    /// <inheritdoc />
    public async Task<int> CloseExpiredLeasesAsync(
        DateTimeOffset now,
        CancellationToken cancellationToken)
    {
        // Selection phase — deliberately sequential. Only the close calls below run in parallel:
        // deciding WHICH sessions to close must stay a single ordered pass so the sweep-precedence
        // rule and the TOCTOU re-check keep their meaning.
        List<(GatewaySession Session, string Reason)> selected = [];
        foreach (GatewaySession session in _registry.Snapshot())
        {
            // A session is swept when its normal lease has expired, it has FAULTED (a faulted
            // session is permanently unusable yet otherwise pins a session slot and a live x86
            // worker until its DefaultLeaseSeconds lease expires), OR its detach-grace retention
            // window has elapsed (last external subscriber dropped and no client reconnected
            // within DetachGraceSeconds). All three are the same teardown; only the reason differs
            // so operators can tell a fault reap from a short reconnect-window expiry from a long
            // idle-lease expiry in logs/metrics. Precedence when several fire simultaneously:
            // lease-expired, then faulted, then detach-grace.
            //
            // TOCTOU note: eligibility is re-verified atomically inside TryBeginCloseIfExpired
            // under _syncRoot, so a client that reattaches a subscriber between the check above
            // and the close call wins the race and the session is left open and usable.
            string? reason = session.IsLeaseExpired(now)
                ? LeaseExpiredReason
                : session.IsFaultedReapable(now)
                    ? FaultedReason
                    : session.IsDetachGraceExpired(now)
                        ? DetachGraceExpiredReason
                        : null;
            if (reason is null)
            {
                continue;
            }

            // Re-verify eligibility atomically and begin the Closing transition before
            // delegating to CloseSessionCoreAsync. If a subscriber reattached between the
            // IsLeaseExpired/IsDetachGraceExpired check above and here, TryBeginCloseIfExpired
            // returns false and we skip this session (it is no longer expired).
            if (!session.TryBeginCloseIfExpired(now, out bool alreadyClosing) && !alreadyClosing)
            {
                continue;
            }

            selected.Add((session, reason));
        }

        if (selected.Count == 0)
        {
            return 0;
        }

        int closedCount = 0;
        object failureSyncRoot = new();
        ExceptionDispatchInfo? firstFailure = null;

        // Close phase. Each close is bounded by the worker shutdown timeout (default 10 s), so a
        // mass expiry with a few hung workers would serialize reaping and starve session slots.
        // Parallel close is safe because TryBeginCloseIfExpired above already flipped every
        // selected session to Closing under its own lock — that idempotent begin-close is the
        // per-session exclusivity invariant, so no two teardowns can run against one session and
        // a session selected here cannot be re-selected by a concurrent sweep.
        await Parallel.ForEachAsync(
            selected,
            new ParallelOptions
            {
                MaxDegreeOfParallelism = MaxParallelSessionCloses,
                CancellationToken = cancellationToken,
            },
            async (candidate, closeToken) =>
            {
                try
                {
                    await CloseSessionCoreAsync(candidate.Session, candidate.Reason, closeToken).ConfigureAwait(false);
                    Interlocked.Increment(ref closedCount);
                }
                catch (Exception exception)
                {
                    // The sequential sweep let a close failure propagate to the lease monitor,
                    // which logs it; that signal is preserved by rethrowing the first failure
                    // below. It is captured rather than thrown here so one failed (or hung)
                    // teardown does not abandon the rest of the already-selected set.
                    lock (failureSyncRoot)
                    {
                        firstFailure ??= ExceptionDispatchInfo.Capture(exception);
                    }
                }
            }).ConfigureAwait(false);

        firstFailure?.Throw();

        return closedCount;
    }

    /// <inheritdoc />
    public async Task ShutdownAsync(CancellationToken cancellationToken)
    {
        // Sessions are drained in parallel: at a worst-case worker shutdown timeout each, a
        // one-at-a-time drain of a full registry outruns any host stop-timeout and leaves the
        // tail to the orphan killer. Per-session exclusivity comes from GatewaySession.CloseAsync's
        // own close gate, and each iteration touches only its own session plus thread-safe
        // registry/metrics state.
        //
        // The body must be exception-TOTAL. Parallel.ForEachAsync cancels the token it hands the
        // sibling bodies as soon as one body throws, so a single escaping exception would abort up
        // to MaxParallelSessionCloses - 1 in-flight graceful shutdowns AND make their kill fallback
        // throw immediately on the freshly cancelled token — sessions neither closed nor killed,
        // i.e. leaked x86 workers that nothing reattaches to (a gateway restart terminates orphans
        // rather than adopting them).
        //
        // For the same reason the loop itself is NOT bound to cancellationToken: a cancelled
        // ParallelOptions token stops dispatching the remaining sessions entirely, so a stop
        // deadline would leave the untried tail neither closed nor killed. Note this FIXES a leak
        // the sequential drain also had rather than restoring its behavior: there the kill fallback
        // ran on the caller's cancelled token, and KillWorkerAsync's entry
        // ThrowIfCancellationRequested threw out of the loop on the first session — zero kills, not
        // "fail fast and still kill". The token is passed to the graceful close only, and the kill
        // runs on CancellationToken.None, so a host stop deadline turns the drain into a kill sweep
        // rather than into a leak.
        //
        // The asymmetry with CloseExpiredLeasesAsync (whose ParallelOptions IS token-bound) is
        // intentional: that sweep is periodic maintenance whose missed sessions are picked up by
        // the next pass and, ultimately, by this drain. This drain is terminal — nothing runs after
        // it — so it must not be abandoned partway.
        await Parallel.ForEachAsync(
            _registry.Snapshot(),
            new ParallelOptions { MaxDegreeOfParallelism = MaxParallelSessionCloses },
            async (session, _) =>
            {
                try
                {
                    await CloseSessionCoreAsync(session, GatewayShutdownReason, cancellationToken).ConfigureAwait(false);
                }
                catch (Exception exception)
                {
                    _logger.LogWarning(
                        exception,
                        "Graceful shutdown failed for session {SessionId}; killing worker.",
                        session.SessionId);

                    // Defensive fallback: CloseSessionCoreAsync's inner
                    // SessionCloseStartedException catch normally removes the session
                    // and accounts the close. This outer fallback only fires for
                    // sessions still in the registry — route through KillWorkerAsync
                    // so the bookkeeping is identical to the dashboard kill path.
                    if (_registry.TryGet(session.SessionId, out GatewaySession? registeredSession)
                        && registeredSession is not null)
                    {
                        try
                        {
                            // Deliberately NOT the caller's token: the kill is the last-resort orphan
                            // preventer, so it must still run when the host stop deadline (or a
                            // sibling body's failure) has already cancelled the drain. It is a
                            // synchronous Kill plus registry/dispose bookkeeping, not a wait on
                            // the worker, so it cannot extend the drain meaningfully.
                            await KillWorkerAsync(session.SessionId, GatewayShutdownReason, CancellationToken.None).ConfigureAwait(false);
                        }
                        catch (Exception killException)
                        {
                            _logger.LogWarning(
                                killException,
                                "Worker kill fallback failed for session {SessionId}.",
                                session.SessionId);
                        }
                    }
                }
            }).ConfigureAwait(false);
    }

    private async Task<SessionCloseResult> CloseSessionCoreAsync(
        GatewaySession session,
        string reason,
        CancellationToken cancellationToken)
    {
        bool wasClosed = session.State == SessionState.Closed;
        try
        {
            SessionCloseResult result = await session.CloseAsync(reason, cancellationToken).ConfigureAwait(false);
            if (!wasClosed && !result.AlreadyClosed)
            {
                _metrics.SessionClosed();
            }

            await RemoveSessionAsync(session).ConfigureAwait(false);
            return result;
        }
        catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
        {
            throw;
        }
        catch (SessionCloseStartedException exception)
        {
            session.MarkFaulted(exception.Message);
            if (!wasClosed)
            {
                // Account the close as a SessionClosed (decrements the open-session
                // gauge AND increments the sessions.closed counter), not just
                // SessionRemoved. The session is being removed from the registry
                // below; treating this as a half-finished close that only
                // decremented the gauge would under-count the closed counter.
                _metrics.SessionClosed();
            }

            _metrics.Fault(SessionManagerErrorCode.CloseFailed.ToString());
            await RemoveSessionAsync(session).ConfigureAwait(false);
            throw new SessionManagerException(
                SessionManagerErrorCode.CloseFailed,
                $"Failed to close session {session.SessionId}.",
                exception);
        }
    }

    private GatewaySession GetRequiredSession(string sessionId)
    {
        if (!_registry.TryGet(sessionId, out GatewaySession? session) || session is null)
        {
            throw new SessionManagerException(
                SessionManagerErrorCode.SessionNotFound,
                $"Session {sessionId} was not found.");
        }

        return session;
    }

    private void EnsureSessionCapacity()
    {
        if (!_sessionSlots.Wait(0))
        {
            throw new SessionManagerException(
                SessionManagerErrorCode.SessionLimitExceeded,
                $"Gateway session limit {_options.Sessions.MaxSessions} has been reached.");
        }
    }

    private async Task RemoveSessionAsync(GatewaySession session)
    {
        if (!_registry.TryRemove(session.SessionId, out GatewaySession? removedSession))
        {
            return;
        }

        _metrics.RemoveSessionEvents(session.SessionId);
        ReleaseSessionSlot();
        await removedSession.DisposeAsync().ConfigureAwait(false);
    }

    private void ReleaseSessionSlot()
    {
        try
        {
            _sessionSlots.Release();
        }
        catch (SemaphoreFullException)
        {
        }
    }

    private GatewaySession CreateSession(
        SessionOpenRequest request,
        string? clientIdentity,
        string? ownerKeyId,
        IReadOnlyList<string>? ownerDashboardTags)
    {
        string sessionUid = Guid.NewGuid().ToString("N");
        string sessionId = $"session-{sessionUid}";
        string backendName = string.IsNullOrWhiteSpace(request.RequestedBackend)
            ? GatewayContractInfo.DefaultBackendName
            : request.RequestedBackend!;
        TimeSpan commandTimeout = ResolveCommandTimeout(request.CommandTimeout);
        TimeSpan startupTimeout = TimeSpan.FromSeconds(_options.Worker.StartupTimeoutSeconds);
        TimeSpan shutdownTimeout = TimeSpan.FromSeconds(_options.Worker.ShutdownTimeoutSeconds);
        TimeSpan leaseDuration = TimeSpan.FromSeconds(_options.Sessions.DefaultLeaseSeconds);
        // The short prefix and bare guid keep the pipe's Unix-domain-socket path
        // (TMPDIR + "CoreFxPipe_" + name) inside the 104-byte sun_path limit on
        // macOS, whose default per-user TMPDIR is ~49 chars; the gateway PID keeps
        // the name collision-free across gateway restarts.
        string pipeName = $"mxgw-{Environment.ProcessId}-{sessionUid}";
        string nonce = CreateNonce();
        DateTimeOffset openedAt = _timeProvider.GetUtcNow();
        string clientCorrelationId = CreateClientCorrelationId(request.ClientSessionName, sessionId);

        SessionEventStreaming eventStreaming = new(
            _eventMapper,
            _options.Events,
            _distributorLogger,
            _timeProvider,
            _metrics,
            _dashboardEventBroadcaster,
            _options.Sessions.AllowMultipleEventSubscribers);

        return new GatewaySession(
            sessionId,
            backendName,
            pipeName,
            nonce,
            clientIdentity,
            ownerKeyId,
            request.ClientSessionName,
            clientCorrelationId,
            commandTimeout,
            startupTimeout,
            shutdownTimeout,
            leaseDuration,
            openedAt,
            eventStreaming,
            TimeSpan.FromSeconds(Math.Max(0, _options.Sessions.DetachGraceSeconds)),
            TimeSpan.FromMilliseconds(Math.Max(0, _options.Sessions.WorkerReadyWaitTimeoutMs)),
            _addressNormalizer,
            TimeSpan.FromSeconds(Math.Max(0, _options.Sessions.FaultedGraceSeconds)),
            ownerDashboardTags);
    }

    private static string CreateClientCorrelationId(
        string? clientSessionName,
        string sessionId)
    {
        string clientName = string.IsNullOrWhiteSpace(clientSessionName)
            ? "client"
            : clientSessionName!;

        return $"{clientName}-{sessionId}";
    }

    private TimeSpan ResolveCommandTimeout(Duration? requestedTimeout)
    {
        if (requestedTimeout is null)
        {
            return TimeSpan.FromSeconds(_options.Sessions.DefaultCommandTimeoutSeconds);
        }

        TimeSpan timeout = requestedTimeout.ToTimeSpan();
        return timeout <= TimeSpan.Zero
            ? TimeSpan.FromSeconds(_options.Sessions.DefaultCommandTimeoutSeconds)
            : timeout;
    }

    private static string CreateNonce()
    {
        Span<byte> bytes = stackalloc byte[32];
        RandomNumberGenerator.Fill(bytes);

        return Convert.ToBase64String(bytes);
    }

}