class QueryService extends QueryServiceGrpc.QueryServiceImplBase implements AxonServerClientService {

    @Override
    public StreamObserver<QueryProviderOutbound> openStream(
            StreamObserver<QueryProviderInbound> inboundStreamObserver) {
        String context = contextProvider.getContext();
        SendingStreamObserver<QueryProviderInbound> wrappedQueryProviderInboundObserver = new SendingStreamObserver<>(
                inboundStreamObserver);

        return new ReceivingStreamObserver<QueryProviderOutbound>(logger) {
            private final AtomicReference<String> clientIdRef = new AtomicReference<>();
            private final AtomicReference<GrpcQueryDispatcherListener> listener = new AtomicReference<>();
            private final AtomicReference<ClientStreamIdentification> clientRef = new AtomicReference<>();
            private final AtomicReference<QueryHandler<QueryProviderInbound>> queryHandler = new AtomicReference<>();

            @Override
            protected void consume(QueryProviderOutbound queryProviderOutbound) {
                switch (queryProviderOutbound.getRequestCase()) {
                    case SUBSCRIBE:
                        instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(),
                                                               wrappedQueryProviderInboundObserver);
                        QuerySubscription subscription = queryProviderOutbound.getSubscribe();
                        checkInitClient(subscription.getClientId(), subscription.getComponentName());
                        String clientStreamId = clientRef.get().getClientStreamId();
                        logger.debug("{}: Subscribe Query {} for {}",
                                     context,
                                     subscription.getQuery(),
                                     subscription.getClientId());
                        SubscribeQuery subscribeQuery = new SubscribeQuery(context,
                                                                           clientStreamId,
                                                                           subscription,
                                                                           queryHandler.get());
                        eventPublisher.publishEvent(
                                subscribeQuery);
                        break;
                    case UNSUBSCRIBE:
                        instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(),
                                                               wrappedQueryProviderInboundObserver);
                        if (clientRef.get() != null) {
                            QuerySubscription unsubscribe = queryProviderOutbound.getUnsubscribe();
                            UnsubscribeQuery unsubscribeQuery = new UnsubscribeQuery(context,
                                                                                     clientRef.get()
                                                                                              .getClientStreamId(),
                                                                                     unsubscribe,
                                                                                     false);
                            eventPublisher.publishEvent(unsubscribeQuery);
                        }
                        break;
                    case FLOW_CONTROL:
                        instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(),
                                                               wrappedQueryProviderInboundObserver);
                        flowControl(queryProviderOutbound.getFlowControl());
                        break;
                    case QUERY_RESPONSE:
                        instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(),
                                                               wrappedQueryProviderInboundObserver);
                        queryDispatcher.handleResponse(queryProviderOutbound.getQueryResponse(),
                                                       clientRef.get().getClientStreamId(),
                                                       clientIdRef.get(),
                                                       false);
                        break;
                    case QUERY_COMPLETE:
                        instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(),
                                                               wrappedQueryProviderInboundObserver);
                        queryDispatcher.handleComplete(queryProviderOutbound.getQueryComplete().getRequestId(),
                                                       clientRef.get().getClientStreamId(),
                                                       clientIdRef.get(),
                                                       false);
                        break;
                    case SUBSCRIPTION_QUERY_RESPONSE:
                        instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(),
                                                               wrappedQueryProviderInboundObserver);
                        SubscriptionQueryResponse response = queryProviderOutbound.getSubscriptionQueryResponse();
                        eventPublisher.publishEvent(new SubscriptionQueryResponseReceived(response, () ->
                                wrappedQueryProviderInboundObserver
                                        .onNext(unsubscribeMessage(response.getSubscriptionIdentifier()))));
                        break;
                    case ACK:
                        InstructionAck ack = queryProviderOutbound.getAck();
                        if (isUnsupportedInstructionErrorResult(ack)) {
                            logger.warn("Unsupported instruction sent to the client {} of context {}.",
                                        clientRef.get().getClientStreamId(),
                                        context);
                        } else {
                            logger.trace("Received instruction ack from the client {} of context {}. Result {}.",
                                         clientRef.get().getClientStreamId(),
                                         context,
                                         ack);
                        }
                        break;
                    default:
                        instructionAckSource.sendUnsupportedInstruction(queryProviderOutbound.getInstructionId(),
                                                                        topology.getMe().getName(),
                                                                        wrappedQueryProviderInboundObserver);
                        break;
                }
            }

            private void flowControl(FlowControl flowControl) {
                initClientReference(flowControl.getClientId());
                if (listener.compareAndSet(null, new GrpcQueryDispatcherListener(queryDispatcher,
                                                                                 clientRef.get().toString(),
                                                                                 wrappedQueryProviderInboundObserver,
                                                                                 processingThreads))) {
                    dispatcherListeners.put(clientRef.get(), listener.get());
                }
                listener.get().addPermits(flowControl.getPermits());
            }

            private void initClientReference(String clientId) {
                String clientStreamId = clientId + "." + UUID.randomUUID().toString();
                if (clientRef.compareAndSet(null, new ClientStreamIdentification(context, clientStreamId))) {
                    clientIdRegistry.register(clientStreamId, clientId, ClientIdRegistry.ConnectionType.QUERY);
                }
                clientIdRef.compareAndSet(null, clientId);
            }

            private void checkInitClient(String clientId, String componentName) {
                initClientReference(clientId);
                queryHandler.compareAndSet(null,
                                           new DirectQueryHandler(wrappedQueryProviderInboundObserver, clientRef.get(),
                                                                  componentName, clientId));
            }

            @Override
            protected String sender() {
                return clientRef.toString();
            }

            @Override
            public void onError(Throwable cause) {
                if (!ExceptionUtils.isCancelled(cause)) {
                    logger.warn("{}: Error on connection from subscriber - {}", clientRef, cause.getMessage());
                }

                cleanup();
            }

            private void cleanup() {
                if (clientRef.get() != null) {
                    String clientStreamId = clientRef.get().getClientStreamId();
                    String clientId = this.clientIdRef.get();
                    clientIdRegistry.unregister(clientStreamId, ClientIdRegistry.ConnectionType.QUERY);
                    eventPublisher.publishEvent(new QueryHandlerDisconnected(context,
                                                                             clientId,
                                                                             clientStreamId));
                }
                if (listener.get() != null) {
                    dispatcherListeners.remove(clientRef.get());
                    listener.get().cancel();
                }
                StreamObserverUtils.complete(inboundStreamObserver);
            }

            @Override
            public void onCompleted() {
                cleanup();
            }
        };
    }

    public  QueryService(Topology topology, QueryDispatcher queryDispatcher, ContextProvider contextProvider,
                        ClientIdRegistry clientIdRegistry,
                        ApplicationEventPublisher eventPublisher,
                        @Qualifier("queryInstructionAckSource")
                                InstructionAckSource<QueryProviderInbound> instructionAckSource);

}