using System.Diagnostics.CodeAnalysis;
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";

    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;

    /// <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>
    public SessionManager(
        ISessionRegistry registry,
        ISessionWorkerClientFactory workerClientFactory,
        IOptions<GatewayOptions> options,
        GatewayMetrics metrics,
        TimeProvider? timeProvider = null,
        ILogger<SessionManager>? logger = 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;
        _options = options.Value;
        _sessionSlots = new SemaphoreSlim(_options.Sessions.MaxSessions, _options.Sessions.MaxSessions);
    }

    /// <summary>
    /// Opens a new gateway session and connects to the worker.
    /// </summary>
    /// <param name="request">Session open request.</param>
    /// <param name="clientIdentity">Client authentication identity.</param>
    /// <param name="ownerKeyId">API key identifier of the caller creating the session.</param>
    /// <param name="cancellationToken">Cancellation token.</param>
    /// <returns>Opened gateway session.</returns>
    public async Task<GatewaySession> OpenSessionAsync(
        SessionOpenRequest request,
        string? clientIdentity,
        string? ownerKeyId,
        CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(request);
        EnsureSessionCapacity();

        GatewaySession? session = null;
        bool sessionOpenedRecorded = false;
        try
        {
            session = CreateSession(request, clientIdentity, ownerKeyId);
            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);
        }
    }

    /// <summary>
    /// Attempts to retrieve a session by ID.
    /// </summary>
    /// <param name="sessionId">Session identifier.</param>
    /// <param name="session">The session if found.</param>
    /// <returns>True if session found; otherwise false.</returns>
    public bool TryGetSession(
        string sessionId,
        [MaybeNullWhen(false)] out GatewaySession session)
    {
        return _registry.TryGet(sessionId, out session);
    }

    /// <summary>
    /// Invokes a worker command on a session asynchronously.
    /// </summary>
    /// <param name="sessionId">Session identifier.</param>
    /// <param name="command">Worker command.</param>
    /// <param name="cancellationToken">Cancellation token.</param>
    /// <returns>Command reply.</returns>
    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;
        }
    }

    /// <summary>
    /// Reads events from a session's event stream asynchronously.
    /// </summary>
    /// <param name="sessionId">Session identifier.</param>
    /// <param name="cancellationToken">Cancellation token.</param>
    /// <returns>Async enumerable of worker events.</returns>
    public IAsyncEnumerable<WorkerEvent> ReadEventsAsync(
        string sessionId,
        CancellationToken cancellationToken)
    {
        GatewaySession session = GetRequiredSession(sessionId);

        return session.ReadEventsAsync(cancellationToken);
    }

    /// <summary>
    /// Closes a gateway session asynchronously.
    /// </summary>
    /// <param name="sessionId">Session identifier.</param>
    /// <param name="cancellationToken">Cancellation token.</param>
    /// <returns>Session close result.</returns>
    public async Task<SessionCloseResult> CloseSessionAsync(
        string sessionId,
        CancellationToken cancellationToken)
    {
        GatewaySession session = GetRequiredSession(sessionId);
        SessionCloseResult result = await CloseSessionCoreAsync(
            session,
            DefaultCloseReason,
            cancellationToken).ConfigureAwait(false);

        return result;
    }

    /// <summary>
    /// Forcefully terminates a session's worker without attempting graceful shutdown.
    /// Mirrors the registry/metrics cleanup that <see cref="CloseSessionCoreAsync"/>
    /// performs after a successful close, but skips the <c>WorkerClient.ShutdownAsync</c>
    /// step that <see cref="GatewaySession.CloseAsync"/> would otherwise attempt.
    /// </summary>
    /// <param name="sessionId">Session identifier.</param>
    /// <param name="reason">Reason recorded for the kill.</param>
    /// <param name="cancellationToken">Cancellation token.</param>
    /// <returns>Session close result.</returns>
    public async Task<SessionCloseResult> KillWorkerAsync(
        string sessionId,
        string reason,
        CancellationToken cancellationToken)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(reason);
        cancellationToken.ThrowIfCancellationRequested();

        GatewaySession session = GetRequiredSession(sessionId);

        // Serialize concurrent kill/close attempts on this session by routing through the
        // per-session close lock (Server-045). Returns whether the session was already in
        // Closed state when the lock was acquired so the metric counter is incremented at
        // most once across concurrent callers.
        bool wasClosed;
        try
        {
            wasClosed = await session.KillWorkerWithCloseGateAsync(reason, cancellationToken).ConfigureAwait(false);
        }
        catch (Exception exception)
        {
            session.MarkFaulted(exception.Message);
            _metrics.Fault(SessionManagerErrorCode.CloseFailed.ToString());

            // Server-044: 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 Server-006 fix on 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);
    }

    /// <summary>
    /// Closes all sessions with expired leases asynchronously.
    /// </summary>
    /// <param name="now">Current time for lease expiration check.</param>
    /// <param name="cancellationToken">Cancellation token.</param>
    /// <returns>Count of sessions closed.</returns>
    public async Task<int> CloseExpiredLeasesAsync(
        DateTimeOffset now,
        CancellationToken cancellationToken)
    {
        int closedCount = 0;
        foreach (GatewaySession session in _registry.Snapshot())
        {
            if (!session.IsLeaseExpired(now))
            {
                continue;
            }

            await CloseSessionCoreAsync(session, LeaseExpiredReason, cancellationToken).ConfigureAwait(false);
            closedCount++;
        }

        return closedCount;
    }

    /// <summary>
    /// Shuts down all active sessions gracefully asynchronously.
    /// </summary>
    /// <param name="cancellationToken">Cancellation token.</param>
    /// <returns>Completed task.</returns>
    public async Task ShutdownAsync(CancellationToken cancellationToken)
    {
        foreach (GatewaySession session in _registry.Snapshot())
        {
            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 (Server-046). The
                // 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 _))
                {
                    try
                    {
                        await KillWorkerAsync(session.SessionId, GatewayShutdownReason, cancellationToken).ConfigureAwait(false);
                    }
                    catch (SessionManagerException killException)
                    {
                        _logger.LogWarning(
                            killException,
                            "Worker kill fallback failed for session {SessionId}.",
                            session.SessionId);
                    }
                }
            }
        }
    }

    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)
            {
                // Server-046: 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 under-counted 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)
    {
        string sessionId = CreateSessionId();
        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);
        string pipeName = $"mxaccess-gateway-{Environment.ProcessId}-{sessionId}";
        string nonce = CreateNonce();
        DateTimeOffset openedAt = _timeProvider.GetUtcNow();
        string clientCorrelationId = CreateClientCorrelationId(request.ClientSessionName, sessionId);

        return new GatewaySession(
            sessionId,
            backendName,
            pipeName,
            nonce,
            clientIdentity,
            ownerKeyId,
            request.ClientSessionName,
            clientCorrelationId,
            commandTimeout,
            startupTimeout,
            shutdownTimeout,
            leaseDuration,
            openedAt);
    }

    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 CreateSessionId()
    {
        return $"session-{Guid.NewGuid():N}";
    }

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

        return Convert.ToBase64String(bytes);
    }

}