package org.dromara.mqtt;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.dromara.app.mqtt.MqttMessageDispatcher;
import org.dromara.common.core.exception.ServiceException;
import org.dromara.common.core.utils.StringUtils;
import org.dromara.mqtt.config.properties.MqttProperties;
import org.eclipse.paho.client.mqttv3.*;
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.*;
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 List<BlockingQueue<InboundMessage>> messageQueues;
    private volatile boolean running;
    private int consumerCount;

    static int deviceIdentityHash(String topic) {
        if (topic == null) {
            return 0;
        }
        int start = topic.startsWith("/") ? 1 : 0;
        int end = topic.indexOf('/', start);
        String identity = end < 0 ? topic.substring(start) : topic.substring(start, end);
        return identity.hashCode();
    }

    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;
        };
    }

    @Bean
    @ConditionalOnProperty(prefix = "mqtt", name = "enabled", havingValue = "true")
    public MqttAsyncClient mqttConnect() throws MqttException {
        if (StringUtils.isBlank(props.getBrokerUrl())) {
            throw new ServiceException("启用 MQTT 时必须配置 broker-url");
        }
        if (StringUtils.isBlank(props.getClientId())) {
            throw new ServiceException("启用 MQTT 时必须配置 client-id");
        }
        consumerCount = Math.max(1, props.getAsync().getConsumerCount());
        messageQueues = createMessageQueues();
        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] 连接已断开：{}", cause == null ? "" : cause.getMessage());
            }

            @Override
            public void connectComplete(boolean reconnect, String serverURI) {
                if (reconnect) {
                    log.info("[MQTT] 已重新连接：{}", serverURI);
                    subscribeTopics();
                }
            }

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

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

        client.connect(options).waitForCompletion();
        log.info("[MQTT] 已连接：{}", props.getBrokerUrl());
        subscribeTopics();
        return client;
    }

    private void startConsumers() {
        for (int i = 0; i < consumerCount; i++) {
            final int consumerIndex = i;
            consumerExecutor.execute(() -> consumeLoop(consumerIndex));
        }
    }

    private void enqueue(String topic, String payload, boolean retained) {
        InboundMessage inboundMessage = new InboundMessage(topic, payload, retained);
        try {
            messageQueues.get(queueIndex(topic)).put(inboundMessage);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            log.warn("[MQTT] 上行消息入队被中断 主题={}", topic);
        }
    }

    private void consumeLoop(int consumerIndex) {
        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);
        BlockingQueue<InboundMessage> queue = messageQueues.get(consumerIndex);

        while (running || !queue.isEmpty()) {
            try {
                InboundMessage first = queue.poll(pollTimeoutMs, TimeUnit.MILLISECONDS);
                if (first == null) {
                    continue;
                }
                batch.add(first);
                queue.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(), message.retained());
        } catch (Exception e) {
            log.error("[MQTT] 消息分发失败 主题={}", message.topic(), e);
        }
    }

    private List<BlockingQueue<InboundMessage>> createMessageQueues() {
        int capacity = Math.max(1, props.getAsync().getQueueCapacity());
        int queueCapacity = Math.max(1, (capacity + Math.max(1, consumerCount) - 1) / Math.max(1, consumerCount));
        List<BlockingQueue<InboundMessage>> queues = new java.util.ArrayList<>(consumerCount);
        for (int i = 0; i < consumerCount; i++) {
            queues.add(new ArrayBlockingQueue<>(queueCapacity));
        }
        return queues;
    }

    private int queueIndex(String topic) {
        return Math.floorMod(deviceIdentityHash(topic), Math.max(1, consumerCount));
    }

    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] 未配置订阅主题");
                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] 已订阅 主题={} 服务质量等级={}", topic, props.getQos()));
        } catch (MqttException e) {
            log.error("[MQTT] 订阅失败", e);
        }
    }

    public void publish(String topic, String payload) {
        try {
            if (client == null || !client.isConnected()) {
                throw new ServiceException("MQTT客户端未连接");
            }
            MqttMessage message = new MqttMessage(payload.getBytes(StandardCharsets.UTF_8));
            message.setQos(props.getQos());
            message.setRetained(false);
            IMqttDeliveryToken token = client.publish(topic, message);
            token.waitForCompletion(Math.max(1000L, props.getConnectionTimeout() * 1000L));
            log.info("[MQTT] 消息已发布 主题={} 消息体={}", topic, payload);
        } catch (MqttException e) {
            log.error("[MQTT] 消息发布失败 主题={}", topic, e);
            throw new ServiceException("MQTT 消息发布失败：{}", e.getMessage());
        }
    }

    @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, boolean retained) {
    }
}
