package com.zb.mom.ww.mxgateway.client;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;

import io.grpc.Context;
import io.grpc.ManagedChannel;
import io.grpc.Metadata;
import io.grpc.Server;
import io.grpc.ServerCall;
import io.grpc.ServerCallHandler;
import io.grpc.ServerInterceptor;
import io.grpc.inprocess.InProcessChannelBuilder;
import io.grpc.inprocess.InProcessServerBuilder;
import io.grpc.stub.ServerCallStreamObserver;
import io.grpc.stub.StreamObserver;
import java.time.Duration;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import mxaccess_gateway.v1.MxAccessGatewayGrpc;
import mxaccess_gateway.v1.MxaccessGateway.ActiveAlarmSnapshot;
import mxaccess_gateway.v1.MxaccessGateway.AddItemReply;
import mxaccess_gateway.v1.MxaccessGateway.AlarmConditionState;
import mxaccess_gateway.v1.MxaccessGateway.BulkSubscribeReply;
import mxaccess_gateway.v1.MxaccessGateway.CloseSessionReply;
import mxaccess_gateway.v1.MxaccessGateway.CloseSessionRequest;
import mxaccess_gateway.v1.MxaccessGateway.MxCommandKind;
import mxaccess_gateway.v1.MxaccessGateway.MxCommandReply;
import mxaccess_gateway.v1.MxaccessGateway.MxCommandRequest;
import mxaccess_gateway.v1.MxaccessGateway.MxEvent;
import mxaccess_gateway.v1.MxaccessGateway.OpenSessionReply;
import mxaccess_gateway.v1.MxaccessGateway.OpenSessionRequest;
import mxaccess_gateway.v1.MxaccessGateway.ProtocolStatus;
import mxaccess_gateway.v1.MxaccessGateway.ProtocolStatusCode;
import mxaccess_gateway.v1.MxaccessGateway.QueryActiveAlarmsRequest;
import mxaccess_gateway.v1.MxaccessGateway.RegisterReply;
import mxaccess_gateway.v1.MxaccessGateway.SessionState;
import mxaccess_gateway.v1.MxaccessGateway.StreamEventsRequest;
import mxaccess_gateway.v1.MxaccessGateway.SubscribeResult;
import org.junit.jupiter.api.Test;

final class MxGatewayClientSessionTests {
    @Test
    void unaryCallsCarryAuthMetadataAndDeadline() throws Exception {
        AtomicReference<String> authorization = new AtomicReference<>();
        AtomicReference<MxCommandRequest> commandRequest = new AtomicReference<>();
        AtomicReference<Boolean> deadlineSeen = new AtomicReference<>(false);

        TestGatewayService service = new TestGatewayService() {
            @Override
            public void openSession(OpenSessionRequest request, StreamObserver<OpenSessionReply> responseObserver) {
                deadlineSeen.set(Context.current().getDeadline() != null);
                responseObserver.onNext(OpenSessionReply.newBuilder()
                        .setSessionId("session-java")
                        .setProtocolStatus(ok())
                        .build());
                responseObserver.onCompleted();
            }

            @Override
            public void invoke(MxCommandRequest request, StreamObserver<MxCommandReply> responseObserver) {
                commandRequest.set(request);
                responseObserver.onNext(MxCommandReply.newBuilder()
                        .setSessionId(request.getSessionId())
                        .setKind(request.getCommand().getKind())
                        .setProtocolStatus(ok())
                        .setRegister(RegisterReply.newBuilder().setServerHandle(42))
                        .build());
                responseObserver.onCompleted();
            }
        };

        try (InProcessGateway gateway = InProcessGateway.start(service, authorization);
                MxGatewayClient client = gateway.client("mxgw_visible_secret", Duration.ofSeconds(5))) {
            MxGatewaySession session = client.openSession("junit-session");

            int serverHandle = session.register("java-test-client");

            assertEquals(42, serverHandle);
            assertEquals("Bearer mxgw_visible_secret", authorization.get());
            assertEquals("session-java", commandRequest.get().getSessionId());
            assertEquals(MxCommandKind.MX_COMMAND_KIND_REGISTER, commandRequest.get().getCommand().getKind());
            assertTrue(deadlineSeen.get());
        }
    }

    @Test
    void methodHelpersReturnTypedHandlesAndRawReplies() throws Exception {
        TestGatewayService service = new TestGatewayService() {
            @Override
            public void invoke(MxCommandRequest request, StreamObserver<MxCommandReply> responseObserver) {
                MxCommandReply.Builder reply = MxCommandReply.newBuilder()
                        .setSessionId(request.getSessionId())
                        .setKind(request.getCommand().getKind())
                        .setProtocolStatus(ok());
                if (request.getCommand().getKind() == MxCommandKind.MX_COMMAND_KIND_ADD_ITEM) {
                    reply.setAddItem(AddItemReply.newBuilder().setItemHandle(7));
                }
                responseObserver.onNext(reply.build());
                responseObserver.onCompleted();
            }
        };

        try (InProcessGateway gateway = InProcessGateway.start(service, new AtomicReference<>());
                MxGatewayClient client = gateway.client("", Duration.ofSeconds(5))) {
            MxGatewaySession session = MxGatewaySession.forSessionId(client, "existing-session");

            int itemHandle = session.addItem(12, "TestObject.TestInt");
            MxCommandReply raw = session.adviseRaw(12, itemHandle);

            assertEquals(7, itemHandle);
            assertEquals(MxCommandKind.MX_COMMAND_KIND_ADVISE, raw.getKind());
        }
    }

    @Test
    void subscribeBulkBuildsOneBulkCommandAndReturnsResults() throws Exception {
        AtomicReference<MxCommandRequest> commandRequest = new AtomicReference<>();
        TestGatewayService service = new TestGatewayService() {
            @Override
            public void invoke(MxCommandRequest request, StreamObserver<MxCommandReply> responseObserver) {
                commandRequest.set(request);
                responseObserver.onNext(MxCommandReply.newBuilder()
                        .setSessionId(request.getSessionId())
                        .setKind(request.getCommand().getKind())
                        .setProtocolStatus(ok())
                        .setSubscribeBulk(BulkSubscribeReply.newBuilder()
                                .addResults(SubscribeResult.newBuilder()
                                        .setServerHandle(12)
                                        .setTagAddress("Area001.Pump001.Speed")
                                        .setItemHandle(34)
                                        .setWasSuccessful(true)))
                        .build());
                responseObserver.onCompleted();
            }
        };

        try (InProcessGateway gateway = InProcessGateway.start(service, new AtomicReference<>());
                MxGatewayClient client = gateway.client("", Duration.ofSeconds(5))) {
            MxGatewaySession session = MxGatewaySession.forSessionId(client, "existing-session");

            List<SubscribeResult> results = session.subscribeBulk(12, List.of("Area001.Pump001.Speed"));

            assertEquals(34, results.get(0).getItemHandle());
            assertEquals(MxCommandKind.MX_COMMAND_KIND_SUBSCRIBE_BULK, commandRequest.get().getCommand().getKind());
            assertEquals(
                    List.of("Area001.Pump001.Speed"),
                    commandRequest.get().getCommand().getSubscribeBulk().getTagAddressesList());
        }
    }

    @Test
    void streamCancellationCancelsServerCall() throws Exception {
        CountDownLatch cancelled = new CountDownLatch(1);
        TestGatewayService service = new TestGatewayService() {
            @Override
            public void streamEvents(StreamEventsRequest request, StreamObserver<MxEvent> responseObserver) {
                ServerCallStreamObserver<MxEvent> serverObserver =
                        (ServerCallStreamObserver<MxEvent>) responseObserver;
                serverObserver.setOnCancelHandler(cancelled::countDown);
                responseObserver.onNext(MxEvent.newBuilder()
                        .setSessionId(request.getSessionId())
                        .setWorkerSequence(1)
                        .build());
            }
        };

        try (InProcessGateway gateway = InProcessGateway.start(service, new AtomicReference<>());
                MxGatewayClient client = gateway.client("", Duration.ofSeconds(5))) {
            MxEventStream events = MxGatewaySession.forSessionId(client, "stream-session").streamEvents();

            assertTrue(events.hasNext());
            assertEquals(1, events.next().getWorkerSequence());
            events.close();

            assertTrue(cancelled.await(5, TimeUnit.SECONDS));
        }
    }

    @Test
    void queryActiveAlarmsForwardsRequestAndStreamsSnapshots() throws Exception {
        AtomicReference<QueryActiveAlarmsRequest> queryRequest = new AtomicReference<>();
        TestGatewayService service = new TestGatewayService() {
            @Override
            public void queryActiveAlarms(
                    QueryActiveAlarmsRequest request, StreamObserver<ActiveAlarmSnapshot> responseObserver) {
                queryRequest.set(request);
                responseObserver.onNext(ActiveAlarmSnapshot.newBuilder()
                        .setAlarmFullReference("Galaxy!TestArea.TestMachine_001.HiTemp")
                        .setSourceObjectReference("TestMachine_001")
                        .setAlarmTypeName("HiAlarm")
                        .setSeverity(800)
                        .setCurrentState(AlarmConditionState.ALARM_CONDITION_STATE_ACTIVE)
                        .setCategory("Process")
                        .setDescription("High temperature alarm")
                        .build());
                responseObserver.onNext(ActiveAlarmSnapshot.newBuilder()
                        .setAlarmFullReference("Galaxy!TestArea.TestMachine_002.LoTemp")
                        .setSourceObjectReference("TestMachine_002")
                        .setAlarmTypeName("LoAlarm")
                        .setSeverity(500)
                        .setCurrentState(AlarmConditionState.ALARM_CONDITION_STATE_ACTIVE_ACKED)
                        .setOperatorUser("operator-1")
                        .setOperatorComment("acknowledged")
                        .build());
                responseObserver.onCompleted();
            }
        };

        try (InProcessGateway gateway = InProcessGateway.start(service, new AtomicReference<>());
                MxGatewayClient client = gateway.client("", Duration.ofSeconds(5))) {
            CountDownLatch completed = new CountDownLatch(1);
            java.util.List<ActiveAlarmSnapshot> received = new java.util.ArrayList<>();
            AtomicReference<Throwable> errorRef = new AtomicReference<>();

            QueryActiveAlarmsRequest request = QueryActiveAlarmsRequest.newBuilder()
                    .setSessionId("alarm-session")
                    .setClientCorrelationId("test-corr-1")
                    .setAlarmFilterPrefix("Galaxy!TestArea")
                    .build();

            MxGatewayActiveAlarmsSubscription subscription = client.queryActiveAlarms(
                    request,
                    new StreamObserver<>() {
                        @Override
                        public void onNext(ActiveAlarmSnapshot value) {
                            received.add(value);
                        }

                        @Override
                        public void onError(Throwable t) {
                            errorRef.set(t);
                            completed.countDown();
                        }

                        @Override
                        public void onCompleted() {
                            completed.countDown();
                        }
                    });

            assertTrue(completed.await(5, TimeUnit.SECONDS));
            subscription.close();

            assertNotNull(queryRequest.get());
            assertEquals("alarm-session", queryRequest.get().getSessionId());
            assertEquals("test-corr-1", queryRequest.get().getClientCorrelationId());
            assertEquals("Galaxy!TestArea", queryRequest.get().getAlarmFilterPrefix());

            assertEquals(2, received.size());
            assertEquals("Galaxy!TestArea.TestMachine_001.HiTemp", received.get(0).getAlarmFullReference());
            assertEquals(
                    AlarmConditionState.ALARM_CONDITION_STATE_ACTIVE,
                    received.get(0).getCurrentState());
            assertEquals(800, received.get(0).getSeverity());
            assertEquals("Galaxy!TestArea.TestMachine_002.LoTemp", received.get(1).getAlarmFullReference());
            assertEquals(
                    AlarmConditionState.ALARM_CONDITION_STATE_ACTIVE_ACKED,
                    received.get(1).getCurrentState());
            assertEquals("operator-1", received.get(1).getOperatorUser());
            assertNull(errorRef.get());
        }
    }

    @Test
    void commandFailureKeepsRawReply() throws Exception {
        TestGatewayService service = new TestGatewayService() {
            @Override
            public void invoke(MxCommandRequest request, StreamObserver<MxCommandReply> responseObserver) {
                responseObserver.onNext(MxCommandReply.newBuilder()
                        .setSessionId(request.getSessionId())
                        .setKind(MxCommandKind.MX_COMMAND_KIND_WRITE)
                        .setProtocolStatus(ProtocolStatus.newBuilder()
                                .setCode(ProtocolStatusCode.PROTOCOL_STATUS_CODE_MXACCESS_FAILURE)
                                .setMessage("MXAccess rejected the write."))
                        .setHresult(-2147220992)
                        .build());
                responseObserver.onCompleted();
            }
        };

        try (InProcessGateway gateway = InProcessGateway.start(service, new AtomicReference<>());
                MxGatewayClient client = gateway.client("", Duration.ofSeconds(5))) {
            MxGatewaySession session = MxGatewaySession.forSessionId(client, "failure-session");

            MxAccessException error = assertThrows(
                    MxAccessException.class,
                    () -> session.write(1, 2, MxValues.int32Value(123), 0));

            assertNotNull(error.reply());
            assertEquals(-2147220992, error.reply().getHresult());
        }
    }

    private static ProtocolStatus ok() {
        return ProtocolStatus.newBuilder()
                .setCode(ProtocolStatusCode.PROTOCOL_STATUS_CODE_OK)
                .build();
    }

    private abstract static class TestGatewayService extends MxAccessGatewayGrpc.MxAccessGatewayImplBase {
        @Override
        public void openSession(OpenSessionRequest request, StreamObserver<OpenSessionReply> responseObserver) {
            responseObserver.onNext(OpenSessionReply.newBuilder()
                    .setSessionId("session-java")
                    .setProtocolStatus(ok())
                    .build());
            responseObserver.onCompleted();
        }

        @Override
        public void closeSession(CloseSessionRequest request, StreamObserver<CloseSessionReply> responseObserver) {
            responseObserver.onNext(CloseSessionReply.newBuilder()
                    .setSessionId(request.getSessionId())
                    .setFinalState(SessionState.SESSION_STATE_CLOSED)
                    .setProtocolStatus(ok())
                    .build());
            responseObserver.onCompleted();
        }
    }

    private record InProcessGateway(Server server, ManagedChannel channel) implements AutoCloseable {
        static InProcessGateway start(
                MxAccessGatewayGrpc.MxAccessGatewayImplBase service, AtomicReference<String> authorization)
                throws Exception {
            String serverName = "mxgw-java-" + UUID.randomUUID();
            ServerInterceptor interceptor = new ServerInterceptor() {
                @Override
                public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
                        ServerCall<ReqT, RespT> call,
                        Metadata headers,
                        ServerCallHandler<ReqT, RespT> next) {
                    authorization.set(headers.get(MxGatewayAuthInterceptor.AUTHORIZATION_HEADER));
                    return next.startCall(call, headers);
                }
            };
            Server server = InProcessServerBuilder.forName(serverName)
                    .directExecutor()
                    .addService(io.grpc.ServerInterceptors.intercept(service, interceptor))
                    .build()
                    .start();
            ManagedChannel channel = InProcessChannelBuilder.forName(serverName)
                    .directExecutor()
                    .build();
            return new InProcessGateway(server, channel);
        }

        MxGatewayClient client(String apiKey, Duration callTimeout) {
            return new MxGatewayClient(
                    channel,
                    MxGatewayClientOptions.builder()
                            .endpoint("in-process")
                            .apiKey(apiKey)
                            .plaintext(true)
                            .callTimeout(callTimeout)
                            .build());
        }

        @Override
        public void close() {
            channel.shutdownNow();
            server.shutdownNow();
        }
    }
}
