package io.bisq.core.dao.blockchain.p2p;

import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.SettableFuture;
import io.bisq.common.Timer;
import io.bisq.common.UserThread;
import io.bisq.common.app.Log;
import io.bisq.common.proto.network.NetworkEnvelope;
import io.bisq.core.dao.blockchain.p2p.messages.GetBsqBlocksRequest;
import io.bisq.core.dao.blockchain.p2p.messages.GetBsqBlocksResponse;
import io.bisq.network.p2p.NodeAddress;
import io.bisq.network.p2p.network.CloseConnectionReason;
import io.bisq.network.p2p.network.Connection;
import io.bisq.network.p2p.network.MessageListener;
import io.bisq.network.p2p.network.NetworkNode;
import io.bisq.network.p2p.peers.PeerManager;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.Random;

import static com.google.common.base.Preconditions.checkArgument;

@Slf4j
public class RequestBlocksHandler implements MessageListener {

    private static final long TIME_OUT_SEC = 40;


    ///////////////////////////////////////////////////////////////////////////////////////////
    // Listener
    ///////////////////////////////////////////////////////////////////////////////////////////

    public interface Listener {
        void onComplete(GetBsqBlocksResponse getBsqBlocksResponse);

        @SuppressWarnings("UnusedParameters")
        void onFault(String errorMessage, @SuppressWarnings("SameParameterValue") @Nullable Connection connection);
    }


    ///////////////////////////////////////////////////////////////////////////////////////////
    // Class fields
    ///////////////////////////////////////////////////////////////////////////////////////////

    private final NetworkNode networkNode;
    private final PeerManager peerManager;
    private final Listener listener;
    private Timer timeoutTimer;
    private final int nonce = new Random().nextInt();
    private boolean stopped;
    private Connection connection;
    private NodeAddress peersNodeAddress;


    ///////////////////////////////////////////////////////////////////////////////////////////
    // Constructor
    ///////////////////////////////////////////////////////////////////////////////////////////

    public RequestBlocksHandler(NetworkNode networkNode,
                                PeerManager peerManager,
                                Listener listener) {
        this.networkNode = networkNode;
        this.peerManager = peerManager;
        this.listener = listener;
    }

    public void cancel() {
        cleanup();
    }


    ///////////////////////////////////////////////////////////////////////////////////////////
    // API
    ///////////////////////////////////////////////////////////////////////////////////////////

    public void requestBlocks(NodeAddress nodeAddress, int startBlockHeight) {
        Log.traceCall("nodeAddress=" + nodeAddress);
        this.peersNodeAddress = nodeAddress;

        if (!stopped) {
            GetBsqBlocksRequest getBsqBlocksRequest = new GetBsqBlocksRequest(startBlockHeight, nonce);

            if (timeoutTimer == null) {
                timeoutTimer = UserThread.runAfter(() -> {  // setup before sending to avoid race conditions
                            if (!stopped) {
                                String errorMessage = "A timeout occurred at sending getBsqBlocksRequest:" + getBsqBlocksRequest +
                                        " on peersNodeAddress:" + peersNodeAddress;
                                log.debug(errorMessage + " / RequestDataHandler=" + RequestBlocksHandler.this);
                                handleFault(errorMessage, peersNodeAddress, CloseConnectionReason.SEND_MSG_TIMEOUT);
                            } else {
                                log.trace("We have stopped already. We ignore that timeoutTimer.run call. " +
                                        "Might be caused by an previous networkNode.sendMessage.onFailure.");
                            }
                        },
                        TIME_OUT_SEC);
            }

            log.debug("We send a {} to peer {}. ", getBsqBlocksRequest.getClass().getSimpleName(), peersNodeAddress);
            networkNode.addMessageListener(this);
            SettableFuture<Connection> future = networkNode.sendMessage(peersNodeAddress, getBsqBlocksRequest);
            Futures.addCallback(future, new FutureCallback<Connection>() {
                @Override
                public void onSuccess(Connection connection) {
                    if (!stopped) {
                        RequestBlocksHandler.this.connection = connection;
                        log.trace("Send " + getBsqBlocksRequest + " to " + peersNodeAddress + " succeeded.");
                    } else {
                        log.trace("We have stopped already. We ignore that networkNode.sendMessage.onSuccess call." +
                                "Might be caused by an previous timeout.");
                    }
                }

                @Override
                public void onFailure(@NotNull Throwable throwable) {
                    if (!stopped) {
                        String errorMessage = "Sending getBsqBlocksRequest to " + peersNodeAddress +
                                " failed. That is expected if the peer is offline.\n\t" +
                                "getBsqBlocksRequest=" + getBsqBlocksRequest + "." +
                                "\n\tException=" + throwable.getMessage();
                        log.error(errorMessage);
                        handleFault(errorMessage, peersNodeAddress, CloseConnectionReason.SEND_MSG_FAILURE);
                    } else {
                        log.trace("We have stopped already. We ignore that networkNode.sendMessage.onFailure call. " +
                                "Might be caused by an previous timeout.");
                    }
                }
            });
        } else {
            log.warn("We have stopped already. We ignore that requestData call.");
        }
    }


    ///////////////////////////////////////////////////////////////////////////////////////////
    // MessageListener implementation
    ///////////////////////////////////////////////////////////////////////////////////////////

    @Override
    public void onMessage(NetworkEnvelope networkEnvelop, Connection connection) {
        if (connection.getPeersNodeAddressOptional().isPresent()) {
            if (connection.getPeersNodeAddressOptional().get().equals(peersNodeAddress)) {
                if (networkEnvelop instanceof GetBsqBlocksResponse) {
                    Log.traceCall(networkEnvelop.toString() + "\n\tconnection=" + connection);
                    if (!stopped) {
                        GetBsqBlocksResponse getBsqBlocksResponse = (GetBsqBlocksResponse) networkEnvelop;
                        if (getBsqBlocksResponse.getRequestNonce() == nonce) {
                            stopTimeoutTimer();
                            checkArgument(connection.getPeersNodeAddressOptional().isPresent(),
                                    "RequestDataHandler.onMessage: connection.getPeersNodeAddressOptional() must be present " +
                                            "at that moment");
                            cleanup();
                            listener.onComplete(getBsqBlocksResponse);
                        } else {
                            log.warn("Nonce not matching. That can happen rarely if we get a response after a canceled " +
                                            "handshake (timeout causes connection close but peer might have sent a msg before " +
                                            "connection was closed).\n\t" +
                                            "We drop that message. nonce={} / requestNonce={}",
                                    nonce, getBsqBlocksResponse.getRequestNonce());
                        }
                    } else {
                        log.warn("We have stopped already. We ignore that onDataRequest call.");
                    }
                }
            } else {
                log.warn("We got a message from another connection and ignore it. That should never happen.");
            }
        } else {
            log.debug("Peers node address is not set yet.");
        }
    }

    public void stop() {
        cleanup();
    }

    ///////////////////////////////////////////////////////////////////////////////////////////
    // Private
    ///////////////////////////////////////////////////////////////////////////////////////////


    @SuppressWarnings("UnusedParameters")
    private void handleFault(String errorMessage, NodeAddress nodeAddress, CloseConnectionReason closeConnectionReason) {
        cleanup();
        peerManager.handleConnectionFault(nodeAddress);
        listener.onFault(errorMessage, null);
    }

    private void cleanup() {
        Log.traceCall();
        stopped = true;
        networkNode.removeMessageListener(this);
        stopTimeoutTimer();
    }

    private void stopTimeoutTimer() {
        if (timeoutTimer != null) {
            timeoutTimer.stop();
            timeoutTimer = null;
        }
    }
}

@@ -6551,24 +6551,94 @@
 resent()) %7B%0A
+            if (networkEnvelop instanceof GetBsqBlocksResponse) %7B%0A    
             
@@ -6720,78 +6720,8 @@
 ) %7B%0A
-                if (networkEnvelop instanceof GetBsqBlocksResponse) %7B%0A
     
@@ -8215,30 +8215,16 @@
    %7D
-%0A            %7D
  else %7B%0A
     
@@ -8211,32 +8211,36 @@
        %7D else %7B%0A
+    
                 
@@ -8330,24 +8330,42 @@
  happen.%22);%0A
+                %7D%0A
             

Fix check if onMessage at connection is matching requested peer address