using Google.Protobuf.WellKnownTypes;
using Microsoft.Extensions.Options;
using MxGateway.Contracts.Proto;
using MxGateway.Server.Configuration;
using MxGateway.Server.Metrics;
using MxGateway.Server.Sessions;
using MxGateway.Server.Workers;

namespace MxGateway.Tests.Gateway.Sessions;

public sealed class SessionManagerTests
{
    [Fact]
    public async Task OpenSessionAsync_WithWorkerReady_RegistersReadySession()
    {
        FakeWorkerClient workerClient = new();
        FakeSessionWorkerClientFactory factory = new(workerClient)
        {
            ApplyLifecycleTransitions = true,
        };
        using GatewayMetrics metrics = new();
        SessionManager manager = CreateManager(factory, metrics: metrics);

        GatewaySession session = await manager.OpenSessionAsync(CreateOpenRequest(), "client-1", CancellationToken.None);

        Assert.True(manager.TryGetSession(session.SessionId, out GatewaySession registered));
        Assert.Same(session, registered);
        Assert.Equal(SessionState.Ready, session.State);
        Assert.Equal("client-1", session.ClientIdentity);
        Assert.Equal(["StartingWorker", "WaitingForPipe", "Handshaking", "InitializingWorker"], factory.ObservedStates);
        Assert.Equal(1, metrics.GetSnapshot().OpenSessions);
        Assert.Equal(1, metrics.GetSnapshot().SessionsOpened);
    }

    [Fact]
    public async Task OpenSessionAsync_GeneratesClientCorrelationIdFromClientNameAndSessionId()
    {
        SessionOpenRequest request = CreateOpenRequest() with
        {
            ClientSessionName = "rust-load-client",
            ClientCorrelationId = "caller-provided-correlation",
        };
        SessionManager manager = CreateManager(new FakeSessionWorkerClientFactory(new FakeWorkerClient()));

        GatewaySession session = await manager.OpenSessionAsync(request, "client-1", CancellationToken.None);

        Assert.Equal($"rust-load-client-{session.SessionId}", session.ClientCorrelationId);
    }

    [Fact]
    public async Task OpenSessionAsync_WhenClientSessionNameMissing_UsesClientCorrelationPrefix()
    {
        SessionOpenRequest request = CreateOpenRequest() with
        {
            ClientSessionName = "",
        };
        SessionManager manager = CreateManager(new FakeSessionWorkerClientFactory(new FakeWorkerClient()));

        GatewaySession session = await manager.OpenSessionAsync(request, "client-1", CancellationToken.None);

        Assert.Equal($"client-{session.SessionId}", session.ClientCorrelationId);
    }

    [Fact]
    public async Task InvokeAsync_WhenSessionReady_ForwardsCommandToWorker()
    {
        FakeWorkerClient workerClient = new();
        SessionManager manager = CreateManager(new FakeSessionWorkerClientFactory(workerClient));
        GatewaySession session = await manager.OpenSessionAsync(CreateOpenRequest(), "client-1", CancellationToken.None);

        WorkerCommandReply reply = await manager.InvokeAsync(
            session.SessionId,
            CreateCommand(MxCommandKind.Ping),
            CancellationToken.None);

        Assert.Equal(1, workerClient.InvokeCount);
        Assert.Equal(MxCommandKind.Ping, reply.Reply.Kind);
    }

    [Fact]
    public async Task GatewaySessionSubscribeBulkAsync_ForwardsOneBulkCommandAndReturnsResults()
    {
        FakeWorkerClient workerClient = new()
        {
            InvokeReply = new WorkerCommandReply
            {
                Reply = new MxCommandReply
                {
                    SessionId = "session-1",
                    CorrelationId = "correlation-1",
                    Kind = MxCommandKind.SubscribeBulk,
                    ProtocolStatus = new ProtocolStatus { Code = ProtocolStatusCode.Ok },
                    SubscribeBulk = new BulkSubscribeReply
                    {
                        Results =
                        {
                            new SubscribeResult
                            {
                                ServerHandle = 12,
                                TagAddress = "Galaxy.Tag.Value",
                                ItemHandle = 512,
                                WasSuccessful = true,
                            },
                        },
                    },
                },
            },
        };
        SessionManager manager = CreateManager(new FakeSessionWorkerClientFactory(workerClient));
        GatewaySession session = await manager.OpenSessionAsync(CreateOpenRequest(), "client-1", CancellationToken.None);

        IReadOnlyList<SubscribeResult> results = await session.SubscribeBulkAsync(
            12,
            ["Galaxy.Tag.Value"],
            CancellationToken.None);

        SubscribeResult result = Assert.Single(results);
        Assert.Equal(512, result.ItemHandle);
        Assert.Equal(1, workerClient.InvokeCount);
        Assert.Equal(MxCommandKind.SubscribeBulk, workerClient.LastCommand?.Command.Kind);
        Assert.Equal(["Galaxy.Tag.Value"], workerClient.LastCommand?.Command.SubscribeBulk.TagAddresses);
    }

    [Fact]
    public async Task InvokeAsync_WhenSessionFaulted_RejectsCommand()
    {
        FakeWorkerClient workerClient = new();
        SessionManager manager = CreateManager(new FakeSessionWorkerClientFactory(workerClient));
        GatewaySession session = await manager.OpenSessionAsync(CreateOpenRequest(), "client-1", CancellationToken.None);
        session.MarkFaulted("test fault");

        SessionManagerException exception = await Assert.ThrowsAsync<SessionManagerException>(
            async () => await manager.InvokeAsync(
                session.SessionId,
                CreateCommand(MxCommandKind.Ping),
                CancellationToken.None));

        Assert.Equal(SessionManagerErrorCode.SessionNotReady, exception.ErrorCode);
        Assert.Equal(0, workerClient.InvokeCount);
    }

    [Fact]
    public async Task CloseSessionAsync_RemovesClosedSession()
    {
        FakeWorkerClient workerClient = new();
        using GatewayMetrics metrics = new();
        SessionManager manager = CreateManager(new FakeSessionWorkerClientFactory(workerClient), metrics: metrics);
        GatewaySession session = await manager.OpenSessionAsync(CreateOpenRequest(), "client-1", CancellationToken.None);

        SessionCloseResult firstClose = await manager.CloseSessionAsync(session.SessionId, CancellationToken.None);
        SessionManagerException secondClose = await Assert.ThrowsAsync<SessionManagerException>(
            async () => await manager.CloseSessionAsync(session.SessionId, CancellationToken.None));

        Assert.False(firstClose.AlreadyClosed);
        Assert.Equal(SessionState.Closed, firstClose.FinalState);
        Assert.Equal(SessionManagerErrorCode.SessionNotFound, secondClose.ErrorCode);
        Assert.Equal(1, workerClient.ShutdownCount);
        Assert.Equal(1, metrics.GetSnapshot().SessionsClosed);
        Assert.Equal(0, metrics.GetSnapshot().OpenSessions);
    }

    [Fact]
    public async Task CloseSessionAsync_WhenWorkerShutdownFails_KillsWorker()
    {
        FakeWorkerClient workerClient = new()
        {
            ShutdownException = new WorkerClientException(
                WorkerClientErrorCode.ShutdownTimeout,
                "Worker shutdown timed out."),
        };
        SessionManager manager = CreateManager(new FakeSessionWorkerClientFactory(workerClient));
        GatewaySession session = await manager.OpenSessionAsync(CreateOpenRequest(), "client-1", CancellationToken.None);

        SessionManagerException exception = await Assert.ThrowsAsync<SessionManagerException>(
            async () => await manager.CloseSessionAsync(session.SessionId, CancellationToken.None));

        Assert.Equal(SessionManagerErrorCode.CloseFailed, exception.ErrorCode);
        Assert.Equal(1, workerClient.ShutdownCount);
        Assert.Equal(1, workerClient.KillCount);
    }

    [Fact]
    public async Task CloseSessionAsync_WhenWorkerShutdownFails_RemovesSessionAndReleasesSlot()
    {
        FakeWorkerClient failingWorkerClient = new()
        {
            ShutdownException = new WorkerClientException(
                WorkerClientErrorCode.ShutdownTimeout,
                "Worker shutdown timed out."),
        };
        FakeWorkerClient replacementWorkerClient = new();
        SessionRegistry registry = new();
        using GatewayMetrics metrics = new();
        SessionManager manager = CreateManager(
            new QueueingSessionWorkerClientFactory(failingWorkerClient, replacementWorkerClient),
            registry,
            metrics,
            CreateOptions(maxSessions: 1));
        GatewaySession firstSession = await manager.OpenSessionAsync(
            CreateOpenRequest(),
            "client-1",
            CancellationToken.None);
        metrics.EventReceived(firstSession.SessionId, MxEventFamily.OnDataChange.ToString());

        SessionManagerException exception = await Assert.ThrowsAsync<SessionManagerException>(
            async () => await manager.CloseSessionAsync(firstSession.SessionId, CancellationToken.None));
        GatewaySession secondSession = await manager.OpenSessionAsync(
            CreateOpenRequest(),
            "client-2",
            CancellationToken.None);

        Assert.Equal(SessionManagerErrorCode.CloseFailed, exception.ErrorCode);
        Assert.False(manager.TryGetSession(firstSession.SessionId, out _));
        Assert.True(manager.TryGetSession(secondSession.SessionId, out _));
        Assert.Equal(1, registry.Count);
        Assert.Equal(1, failingWorkerClient.KillCount);
        Assert.Equal(1, failingWorkerClient.DisposeCount);
        GatewayMetricsSnapshot snapshot = metrics.GetSnapshot();
        Assert.Equal(0, snapshot.SessionsClosed);
        Assert.False(snapshot.EventsBySession.ContainsKey(firstSession.SessionId));
        Assert.Equal(1, snapshot.OpenSessions);
    }

    [Fact]
    public async Task OpenSessionAsync_WhenWorkerCreationFails_RemovesSessionFromRegistry()
    {
        SessionRegistry registry = new();
        using GatewayMetrics metrics = new();
        SessionManager manager = CreateManager(
            new FailingSessionWorkerClientFactory(),
            registry,
            metrics);

        SessionManagerException exception = await Assert.ThrowsAsync<SessionManagerException>(
            async () => await manager.OpenSessionAsync(CreateOpenRequest(), "client-1", CancellationToken.None));

        Assert.Equal(SessionManagerErrorCode.OpenFailed, exception.ErrorCode);
        Assert.Equal(0, registry.Count);
        Assert.Equal(0, metrics.GetSnapshot().SessionsOpened);
        Assert.Equal(1, metrics.GetSnapshot().Faults);
    }

    [Fact]
    public async Task CloseExpiredLeasesAsync_ClosesExpiredSessionsOnly()
    {
        FakeWorkerClient expiredClient = new();
        FakeWorkerClient activeClient = new();
        QueueingSessionWorkerClientFactory factory = new(expiredClient, activeClient);
        SessionManager manager = CreateManager(factory);
        GatewaySession expiredSession = await manager.OpenSessionAsync(CreateOpenRequest(), "client-1", CancellationToken.None);
        GatewaySession activeSession = await manager.OpenSessionAsync(CreateOpenRequest(), "client-2", CancellationToken.None);
        DateTimeOffset now = DateTimeOffset.UtcNow;
        expiredSession.ExtendLease(now.AddSeconds(-1));
        activeSession.ExtendLease(now.AddMinutes(5));

        int closedCount = await manager.CloseExpiredLeasesAsync(now, CancellationToken.None);

        Assert.Equal(1, closedCount);
        Assert.Equal(SessionState.Closed, expiredSession.State);
        Assert.Equal(SessionState.Ready, activeSession.State);
        Assert.Equal(1, expiredClient.ShutdownCount);
        Assert.Equal(0, activeClient.ShutdownCount);
    }

    [Fact]
    public async Task ShutdownAsync_ClosesAllRegisteredSessions()
    {
        FakeWorkerClient firstClient = new();
        FakeWorkerClient secondClient = new();
        QueueingSessionWorkerClientFactory factory = new(firstClient, secondClient);
        using GatewayMetrics metrics = new();
        SessionManager manager = CreateManager(factory, metrics: metrics);
        GatewaySession firstSession = await manager.OpenSessionAsync(CreateOpenRequest(), "client-1", CancellationToken.None);
        GatewaySession secondSession = await manager.OpenSessionAsync(CreateOpenRequest(), "client-2", CancellationToken.None);

        await manager.ShutdownAsync(CancellationToken.None);

        Assert.Equal(SessionState.Closed, firstSession.State);
        Assert.Equal(SessionState.Closed, secondSession.State);
        Assert.Equal(1, firstClient.ShutdownCount);
        Assert.Equal(1, secondClient.ShutdownCount);
        Assert.Equal(2, metrics.GetSnapshot().SessionsClosed);
        Assert.Equal(0, metrics.GetSnapshot().OpenSessions);
    }

    private static SessionManager CreateManager(
        ISessionWorkerClientFactory factory,
        ISessionRegistry? registry = null,
        GatewayMetrics? metrics = null,
        GatewayOptions? options = null)
    {
        return new SessionManager(
            registry ?? new SessionRegistry(),
            factory,
            Options.Create(options ?? CreateOptions()),
            metrics ?? new GatewayMetrics());
    }

    private static GatewayOptions CreateOptions(int maxSessions = 64)
    {
        return new GatewayOptions
        {
            Sessions = new SessionOptions
            {
                DefaultCommandTimeoutSeconds = 30,
                MaxSessions = maxSessions,
            },
            Worker = new WorkerOptions
            {
                StartupTimeoutSeconds = 30,
                ShutdownTimeoutSeconds = 10,
            },
        };
    }

    private static SessionOpenRequest CreateOpenRequest()
    {
        return new SessionOpenRequest(
            RequestedBackend: null,
            ClientSessionName: "test-session",
            ClientCorrelationId: "client-correlation-1",
            CommandTimeout: Duration.FromTimeSpan(TimeSpan.FromSeconds(5)));
    }

    private static WorkerCommand CreateCommand(MxCommandKind kind)
    {
        return new WorkerCommand
        {
            Command = new MxCommand
            {
                Kind = kind,
            },
        };
    }

    private sealed class FakeSessionWorkerClientFactory(IWorkerClient workerClient) : ISessionWorkerClientFactory
    {
        public List<string> ObservedStates { get; } = [];

        public bool ApplyLifecycleTransitions { get; init; }

        public Task<IWorkerClient> CreateAsync(
            GatewaySession session,
            CancellationToken cancellationToken)
        {
            ObservedStates.Add(session.State.ToString());
            if (ApplyLifecycleTransitions)
            {
                session.TransitionTo(SessionState.WaitingForPipe);
                ObservedStates.Add(session.State.ToString());
                session.TransitionTo(SessionState.Handshaking);
                ObservedStates.Add(session.State.ToString());
                session.TransitionTo(SessionState.InitializingWorker);
                ObservedStates.Add(session.State.ToString());
            }

            return Task.FromResult(workerClient);
        }
    }

    private sealed class QueueingSessionWorkerClientFactory : ISessionWorkerClientFactory
    {
        private readonly Queue<IWorkerClient> _workerClients;

        public QueueingSessionWorkerClientFactory(params IWorkerClient[] workerClients)
        {
            _workerClients = new Queue<IWorkerClient>(workerClients);
        }

        public Task<IWorkerClient> CreateAsync(
            GatewaySession session,
            CancellationToken cancellationToken)
        {
            return Task.FromResult(_workerClients.Dequeue());
        }
    }

    private sealed class FailingSessionWorkerClientFactory : ISessionWorkerClientFactory
    {
        public Task<IWorkerClient> CreateAsync(
            GatewaySession session,
            CancellationToken cancellationToken)
        {
            throw new InvalidOperationException("worker startup failed");
        }
    }

    private sealed class FakeWorkerClient : IWorkerClient
    {
        public string SessionId { get; init; } = "session-1";

        public int? ProcessId { get; init; } = 1234;

        public WorkerClientState State { get; set; } = WorkerClientState.Ready;

        public DateTimeOffset LastHeartbeatAt { get; init; } = DateTimeOffset.UtcNow;

        public int InvokeCount { get; private set; }

        public int ShutdownCount { get; private set; }

        public int KillCount { get; private set; }

        public int DisposeCount { get; private set; }

        public Exception? ShutdownException { get; init; }

        public WorkerCommand? LastCommand { get; private set; }

        public WorkerCommandReply? InvokeReply { get; init; }

        public Task StartAsync(CancellationToken cancellationToken)
        {
            return Task.CompletedTask;
        }

        public Task<WorkerCommandReply> InvokeAsync(
            WorkerCommand command,
            TimeSpan timeout,
            CancellationToken cancellationToken)
        {
            InvokeCount++;
            LastCommand = command;
            if (InvokeReply is not null)
            {
                return Task.FromResult(InvokeReply);
            }

            MxCommandKind kind = command.Command?.Kind ?? MxCommandKind.Unspecified;

            return Task.FromResult(new WorkerCommandReply
            {
                Reply = new MxCommandReply
                {
                    SessionId = SessionId,
                    CorrelationId = "correlation-1",
                    Kind = kind,
                },
            });
        }

        public async IAsyncEnumerable<WorkerEvent> ReadEventsAsync(
            [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken cancellationToken)
        {
            await Task.CompletedTask;
            yield break;
        }

        public Task ShutdownAsync(
            TimeSpan timeout,
            CancellationToken cancellationToken)
        {
            ShutdownCount++;
            if (ShutdownException is not null)
            {
                throw ShutdownException;
            }

            State = WorkerClientState.Closed;
            return Task.CompletedTask;
        }

        public void Kill(string reason)
        {
            KillCount++;
            State = WorkerClientState.Faulted;
        }

        public ValueTask DisposeAsync()
        {
            DisposeCount++;
            return ValueTask.CompletedTask;
        }
    }
}