package org.dromara.mqtt;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.dromara.app.mqtt.MqttMessageDispatcher;
import org.dromara.mqtt.config.properties.MqttProperties;
import org.eclipse.paho.client.mqttv3.IMqttDeliveryToken;
import org.eclipse.paho.client.mqttv3.MqttAsyncClient;
import org.eclipse.paho.client.mqttv3.MqttCallbackExtended;
import org.eclipse.paho.client.mqttv3.MqttConnectOptions;
import org.eclipse.paho.client.mqttv3.MqttException;
import org.eclipse.paho.client.mqttv3.MqttMessage;
import org.eclipse.paho.client.mqttv3.persist.MemoryPersistence;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.boot.autoconfigure.AutoConfiguration;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.scheduling.annotation.EnableScheduling;

import javax.net.ssl.SSLContext;
import javax.net.ssl.TrustManager;
import javax.net.ssl.X509TrustManager;
import java.nio.charset.StandardCharsets;
import java.security.GeneralSecurityException;
import java.security.SecureRandom;
import java.security.cert.X509Certificate;
import java.util.List;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

@Slf4j
@AutoConfiguration
@EnableScheduling
@EnableConfigurationProperties(MqttProperties.class)
@RequiredArgsConstructor
public class MqttClientManager implements DisposableBean {

    private final MqttProperties props;
    private final MqttMessageDispatcher dispatcher;
    private MqttAsyncClient client;
    private ThreadPoolExecutor consumerExecutor;
    private BlockingQueue<InboundMessage> messageQueue;
    private volatile boolean running;

    @Bean
    @ConditionalOnProperty(prefix = "mqtt", name = "enabled", havingValue = "true")
    public MqttAsyncClient mqttConnect() throws MqttException {
        messageQueue = new ArrayBlockingQueue<>(Math.max(1, props.getAsync().getQueueCapacity()));
        consumerExecutor = createConsumerExecutor();
        running = true;
        startConsumers();
        client = new MqttAsyncClient(props.getBrokerUrl(), props.getClientId(), new MemoryPersistence());

        MqttConnectOptions options = new MqttConnectOptions();
        options.setUserName(props.getUsername());
        if (props.getPassword() != null) {
            options.setPassword(props.getPassword().toCharArray());
        }
        options.setKeepAliveInterval(props.getKeepAlive());
        options.setConnectionTimeout(props.getConnectionTimeout());
        options.setAutomaticReconnect(props.isAutomaticReconnect());
        options.setCleanSession(props.isCleanSession());
        options.setMaxInflight(props.getMaxInflight());
        configureTls(options);

        client.setCallback(new MqttCallbackExtended() {
            @Override
            public void connectionLost(Throwable cause) {
                log.warn("[MQTT] connection lost: {}", cause == null ? "" : cause.getMessage());
            }

            @Override
            public void connectComplete(boolean reconnect, String serverURI) {
                if (reconnect) {
                    log.info("[MQTT] reconnected: {}", serverURI);
                    subscribeTopics();
                }
            }

            @Override
            public void messageArrived(String topic, MqttMessage message) {
                String payload = new String(message.getPayload(), StandardCharsets.UTF_8);
                enqueue(topic, payload);
            }

            @Override
            public void deliveryComplete(IMqttDeliveryToken token) {
            }
        });

        client.connect(options).waitForCompletion();
        log.info("[MQTT] connected: {}", props.getBrokerUrl());
        subscribeTopics();
        return client;
    }

    private void configureTls(MqttConnectOptions options) throws MqttException {
        if (!isTlsEnabled()) {
            return;
        }

        options.setHttpsHostnameVerificationEnabled(!props.getTls().isSkipVerify());
        if (props.getTls().isSkipVerify()) {
            try {
                options.setSocketFactory(createTrustAllSslContext().getSocketFactory());
                options.setSSLHostnameVerifier((hostname, session) -> true);
            } catch (GeneralSecurityException e) {
                throw new MqttException(e);
            }
        }
    }

    private boolean isTlsEnabled() {
        return props.getTls() != null
            && props.getTls().isEnabled()
            && props.getBrokerUrl() != null
            && props.getBrokerUrl().startsWith("ssl://");
    }

    private SSLContext createTrustAllSslContext() throws GeneralSecurityException {
        TrustManager[] trustManagers = new TrustManager[] {
            new X509TrustManager() {
                @Override
                public void checkClientTrusted(X509Certificate[] chain, String authType) {
                }

                @Override
                public void checkServerTrusted(X509Certificate[] chain, String authType) {
                }

                @Override
                public X509Certificate[] getAcceptedIssuers() {
                    return new X509Certificate[0];
                }
            }
        };
        SSLContext sslContext = SSLContext.getInstance("TLS");
        sslContext.init(null, trustManagers, new SecureRandom());
        return sslContext;
    }

    private ThreadPoolExecutor createConsumerExecutor() {
        MqttProperties.Async async = props.getAsync();
        int corePoolSize = Math.max(1, async.getCorePoolSize());
        int maxPoolSize = Math.max(corePoolSize, async.getMaxPoolSize());

        ThreadPoolExecutor executor = new ThreadPoolExecutor(
            corePoolSize,
            maxPoolSize,
            60L,
            TimeUnit.SECONDS,
            new ArrayBlockingQueue<>(maxPoolSize),
            mqttThreadFactory(),
            new ThreadPoolExecutor.CallerRunsPolicy()
        );
        executor.allowCoreThreadTimeOut(true);
        return executor;
    }

    private ThreadFactory mqttThreadFactory() {
        AtomicLong threadNo = new AtomicLong();
        return runnable -> {
            Thread thread = new Thread(runnable);
            thread.setName("mqtt-consumer-" + threadNo.incrementAndGet());
            thread.setDaemon(true);
            return thread;
        };
    }

    private void startConsumers() {
        int consumerCount = Math.max(1, props.getAsync().getConsumerCount());
        for (int i = 0; i < consumerCount; i++) {
            consumerExecutor.execute(this::consumeLoop);
        }
    }

    private void enqueue(String topic, String payload) {
        InboundMessage inboundMessage = new InboundMessage(topic, payload);
        try {
            if (!messageQueue.offer(inboundMessage, props.getAsync().getOfferTimeoutMs(), TimeUnit.MILLISECONDS)) {
                log.warn("[MQTT] inbound queue full, dropped topic={}", topic);
            }
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            log.warn("[MQTT] enqueue interrupted topic={}", topic);
        }
    }

    private void consumeLoop() {
        MqttProperties.Async async = props.getAsync();
        int batchSize = Math.max(1, async.getBatchSize());
        long pollTimeoutMs = Math.max(1, async.getPollTimeoutMs());
        List<InboundMessage> batch = new java.util.ArrayList<>(batchSize);

        while (running || !messageQueue.isEmpty()) {
            try {
                InboundMessage first = messageQueue.poll(pollTimeoutMs, TimeUnit.MILLISECONDS);
                if (first == null) {
                    continue;
                }
                batch.add(first);
                messageQueue.drainTo(batch, batchSize - 1);
                for (InboundMessage message : batch) {
                    dispatch(message);
                }
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                break;
            } finally {
                batch.clear();
            }
        }
    }

    private void dispatch(InboundMessage message) {
        try {
            dispatcher.dispatch(message.topic(), message.payload());
        } catch (Exception e) {
            log.error("[MQTT] dispatch failed topic={}", message.topic(), e);
        }
    }

    private void subscribeTopics() {
        try {
            if (client == null || !client.isConnected()) {
                return;
            }
            List<String> topics = props.getTopics() == null ? List.of() : props.getTopics().getSubscribe();
            if (topics == null || topics.isEmpty()) {
                log.warn("[MQTT] no subscribe topics configured");
                return;
            }
            int[] qos = topics.stream().mapToInt(topic -> props.getQos()).toArray();
            client.subscribe(topics.toArray(new String[0]), qos).waitForCompletion();
            topics.forEach(topic -> log.info("[MQTT] subscribed: {} qos={}", topic, props.getQos()));
        } catch (MqttException e) {
            log.error("[MQTT] subscribe failed", e);
        }
    }

    public void publish(String topic, String payload) {
        try {
            if (client == null || !client.isConnected()) {
                log.warn("[MQTT] publish skipped, client not connected topic={}", topic);
                return;
            }
            MqttMessage message = new MqttMessage(payload.getBytes(StandardCharsets.UTF_8));
            message.setQos(props.getQos());
            message.setRetained(false);
            client.publish(topic, message);
            log.info("[MQTT] published topic={} payload={}", topic, payload);
        } catch (MqttException e) {
            log.error("[MQTT] publish failed topic={}", topic, e);
        }
    }

    @Override
    public void destroy() throws Exception {
        if (client != null) {
            if (client.isConnected()) {
                client.disconnect().waitForCompletion(5000);
            }
            client.close();
        }
        if (consumerExecutor != null) {
            running = false;
            consumerExecutor.shutdown();
            if (!consumerExecutor.awaitTermination(10, TimeUnit.SECONDS)) {
                consumerExecutor.shutdownNow();
            }
        }
    }

    private record InboundMessage(String topic, String payload) {
    }
}
