package net.xmeter.samplers.mqtt.hivemq;

import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.logging.Logger;

import com.hivemq.client.mqtt.MqttClientSslConfig;
import com.hivemq.client.mqtt.MqttWebSocketConfig;
import com.hivemq.client.mqtt.MqttWebSocketConfigBuilder;
import com.hivemq.client.mqtt.lifecycle.MqttClientAutoReconnect;
import com.hivemq.client.mqtt.mqtt3.Mqtt3BlockingClient;
import com.hivemq.client.mqtt.mqtt3.Mqtt3Client;
import com.hivemq.client.mqtt.mqtt3.Mqtt3ClientBuilder;
import com.hivemq.client.mqtt.mqtt3.message.auth.Mqtt3SimpleAuth;
import com.hivemq.client.mqtt.mqtt3.message.auth.Mqtt3SimpleAuthBuilder;
import com.hivemq.client.mqtt.mqtt3.message.connect.Mqtt3ConnectBuilder;
import com.hivemq.client.mqtt.mqtt3.message.connect.connack.Mqtt3ConnAck;

import net.xmeter.samplers.mqtt.ConnectionParameters;
import net.xmeter.samplers.mqtt.MQTTClient;
import net.xmeter.samplers.mqtt.MQTTConnection;
import net.xmeter.samplers.mqtt.MQTTClientException;

class HiveMQTTClient implements MQTTClient {
    private static final Logger logger = Logger.getLogger(HiveMQTTClient.class.getCanonicalName());
    private final ConnectionParameters parameters;
    private final Mqtt3BlockingClient client;

    HiveMQTTClient(ConnectionParameters parameters) throws Exception {
        this.parameters = parameters;
        Mqtt3ClientBuilder mqtt3ClientBuilder = Mqtt3Client.builder()
                .identifier(parameters.getClientId())
                .serverHost(parameters.getHost())
                .serverPort(parameters.getPort());
        mqtt3ClientBuilder = applyAdditionalConfig(mqtt3ClientBuilder, parameters);
        client = mqtt3ClientBuilder
                .buildBlocking();
    }

    private Mqtt3ClientBuilder applyAdditionalConfig(Mqtt3ClientBuilder builder, ConnectionParameters parameters)
            throws Exception {
        if (parameters.getReconnectMaxAttempts() > 0) {
            builder = builder.automaticReconnect(MqttClientAutoReconnect.builder().build());
        }
        if (parameters.isSecureProtocol()) {
            MqttClientSslConfig sslConfig = ((HiveMQTTSsl) parameters.getSsl()).getSslConfig();
            builder = builder.sslConfig(sslConfig);
        }
        if (parameters.isWebSocketProtocol()) {
            MqttWebSocketConfigBuilder wsConfigBuilder = MqttWebSocketConfig.builder();
            if (parameters.getPath() != null) {
                wsConfigBuilder.serverPath(parameters.getPath());
            }
            builder = builder.webSocketConfig(wsConfigBuilder.build());
        }
        return builder;
    }

    @Override
    public String getClientId() {
        return parameters.getClientId();
    }

    @Override
    public MQTTConnection connect() throws Exception {
        Mqtt3ConnectBuilder.Send<CompletableFuture<Mqtt3ConnAck>> connectSend = client.toAsync().connectWith()
                .cleanSession(parameters.isCleanSession())
                .keepAlive(parameters.getKeepAlive());
        Mqtt3SimpleAuth auth = createAuth();
        if (auth != null) {
            connectSend = connectSend.simpleAuth(auth);
        }
        logger.info(() -> "Connect client: " + parameters.getClientId());
        CompletableFuture<Mqtt3ConnAck> connectFuture = connectSend.send();
        try {
            Mqtt3ConnAck connAck = connectFuture.get(parameters.getConnectTimeout(), TimeUnit.SECONDS);
            logger.info(() -> "Connected client: " + parameters.getClientId());
            return new HiveMQTTConnection(client, parameters.getClientId(), connAck);
        } catch (TimeoutException e) {
            client.disconnect();
            throw new MQTTClientException("Connection timeout " + client, e);
        }
    }

    private Mqtt3SimpleAuth createAuth() {
        if (parameters.getUsername() != null) {
            Mqtt3SimpleAuthBuilder.Complete simpleAuth = Mqtt3SimpleAuth.builder()
                    .username(parameters.getUsername());
            if (parameters.getPassword() != null) {
                simpleAuth.password(parameters.getPassword().getBytes());
            }
            return simpleAuth.build();
        }
        return null;
    }
}

@@ -162,24 +162,56 @@
 tException;%0A
+import java.util.logging.Level;%0A
 import java.
@@ -1037,29 +1037,34 @@
 s.mqtt.MQTTC
-onnec
+lientExcep
 tion;%0Aimport
@@ -1090,34 +1090,29 @@
 s.mqtt.MQTTC
-lientExcep
+onnec
 tion;%0A%0Aclass
@@ -3729,45 +3729,217 @@
 ion 
-e) %7B%0A            client.disconnect();
+timeoutException) %7B%0A            try %7B%0A                client.disconnect();%0A            %7D catch (Exception e) %7B%0A                logger.log(Level.FINE, %22Disconnect on timeout failed %22 + client, e);%0A            %7D
 %0A   
@@ -4001,33 +4001,48 @@
 out %22 + client, 
-e
+timeoutException
 );%0A        %7D%0A   

Fix hivemq timeout disconnect handling