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

import io.grpc.stub.ClientCallStreamObserver;
import io.grpc.stub.ClientResponseObserver;
import io.grpc.stub.StreamObserver;
import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.atomic.AtomicBoolean;
import mxaccess_gateway.v1.MxaccessGateway.MxEvent;
import mxaccess_gateway.v1.MxaccessGateway.StreamEventsRequest;

/**
 * Cancellable handle returned by the async {@code streamEvents} variant.
 *
 * <p>Wraps a caller-supplied {@link StreamObserver} and exposes a
 * {@link #cancel()} entry point that aborts the underlying gRPC call. The
 * subscription also implements {@link AutoCloseable} so it can participate in
 * try-with-resources blocks.
 */
public final class MxGatewayEventSubscription implements AutoCloseable {
    private final AtomicReference<ClientCallStreamObserver<StreamEventsRequest>> requestStream = new AtomicReference<>();
    private final AtomicBoolean cancelled = new AtomicBoolean();

    ClientResponseObserver<StreamEventsRequest, MxEvent> wrap(StreamObserver<MxEvent> observer) {
        return new ClientResponseObserver<>() {
            @Override
            public void beforeStart(ClientCallStreamObserver<StreamEventsRequest> stream) {
                requestStream.set(stream);
                if (cancelled.get()) {
                    stream.cancel("client cancelled event stream", null);
                }
            }

            @Override
            public void onNext(MxEvent value) {
                observer.onNext(value);
            }

            @Override
            public void onError(Throwable error) {
                observer.onError(error);
            }

            @Override
            public void onCompleted() {
                observer.onCompleted();
            }
        };
    }

    /**
     * Cancels the underlying gRPC call. Safe to invoke before the call has
     * started; cancellation is recorded and applied as soon as the stream
     * attaches.
     */
    public void cancel() {
        cancelled.set(true);
        ClientCallStreamObserver<StreamEventsRequest> stream = requestStream.get();
        if (stream != null) {
            stream.cancel("client cancelled event stream", null);
        }
    }

    @Override
    public void close() {
        cancel();
    }
}
