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.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;

/**
 * Shared base for the cancellable subscription handles returned by the
 * async-style server-streaming RPCs ({@code streamEvents}, {@code streamAlarms},
 * {@code queryActiveAlarms}, {@code watchDeployEvents}).
 *
 * <p>All four subscription classes share the same lifecycle and cancellation
 * contract:
 *
 * <ul>
 *   <li>{@link #wrap(StreamObserver)} returns a {@link ClientResponseObserver}
 *       that captures the underlying {@link ClientCallStreamObserver} in
 *       {@code beforeStart}. If {@link #cancel()} was called before the gRPC
 *       call attached, the stream is cancelled eagerly inside
 *       {@code beforeStart} (the Client.Java-014 close-before-beforeStart
 *       fix).</li>
 *   <li>{@link #cancel()} is idempotent. It records the cancellation flag and
 *       forwards {@code cancel(message, cause)} to the underlying stream when
 *       one is attached; otherwise the flag is checked in {@code beforeStart}
 *       once the stream attaches.</li>
 *   <li>{@link #close()} delegates to {@link #cancel()} so the handle can be
 *       used with try-with-resources.</li>
 * </ul>
 *
 * <p>Subclasses supply only the cancel-message string used by {@code cancel()}.
 * Refactor introduced for Client.Java-036 — the four prior subscription
 * classes were structural near-clones (~60 lines each).
 */
abstract class MxGatewayStreamSubscription<TRequest, TResponse> implements AutoCloseable {
    private final AtomicReference<ClientCallStreamObserver<TRequest>> requestStream = new AtomicReference<>();
    private final AtomicBoolean cancelled = new AtomicBoolean();
    private final String cancelMessage;

    MxGatewayStreamSubscription(String cancelMessage) {
        this.cancelMessage = cancelMessage;
    }

    final ClientResponseObserver<TRequest, TResponse> wrap(StreamObserver<TResponse> observer) {
        return new ClientResponseObserver<>() {
            @Override
            public void beforeStart(ClientCallStreamObserver<TRequest> stream) {
                requestStream.set(stream);
                if (cancelled.get()) {
                    stream.cancel(cancelMessage, null);
                }
            }

            @Override
            public void onNext(TResponse 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 final void cancel() {
        cancelled.set(true);
        ClientCallStreamObserver<TRequest> stream = requestStream.get();
        if (stream != null) {
            stream.cancel(cancelMessage, null);
        }
    }

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