package com.lingniu.ingest.sink.kafka;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.context.SmartLifecycle;

import java.time.Duration;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Supplier;

public final class KafkaEnvelopeConsumerRunner implements SmartLifecycle, AutoCloseable {

    private static final Logger log = LoggerFactory.getLogger(KafkaEnvelopeConsumerRunner.class);

    private final Supplier<List<KafkaEnvelopeConsumerWorker>> workersSupplier;
    private final Duration pollTimeout;
    private final Duration loopBackoff;
    private final boolean autoStartup;
    private final AtomicBoolean running = new AtomicBoolean(false);
    private final AtomicBoolean closed = new AtomicBoolean(false);
    private volatile List<KafkaEnvelopeConsumerWorker> workers;
    private ExecutorService executor;

    public KafkaEnvelopeConsumerRunner(List<KafkaEnvelopeConsumerWorker> workers,
                                       Duration pollTimeout,
                                       Duration loopBackoff,
                                       boolean autoStartup) {
        if (workers == null || workers.isEmpty()) {
            throw new IllegalArgumentException("workers must not be empty");
        }
        this.workersSupplier = () -> List.copyOf(workers);
        this.workers = List.copyOf(workers);
        this.pollTimeout = pollTimeout == null ? Duration.ofSeconds(1) : pollTimeout;
        this.loopBackoff = loopBackoff == null ? Duration.ofSeconds(1) : loopBackoff;
        this.autoStartup = autoStartup;
    }

    public KafkaEnvelopeConsumerRunner(Supplier<List<KafkaEnvelopeConsumerWorker>> workersSupplier,
                                       Duration pollTimeout,
                                       Duration loopBackoff,
                                       boolean autoStartup) {
        if (workersSupplier == null) {
            throw new IllegalArgumentException("workersSupplier must not be null");
        }
        this.workersSupplier = workersSupplier;
        this.pollTimeout = pollTimeout == null ? Duration.ofSeconds(1) : pollTimeout;
        this.loopBackoff = loopBackoff == null ? Duration.ofSeconds(1) : loopBackoff;
        this.autoStartup = autoStartup;
    }

    public List<KafkaEnvelopeConsumerWorker> workers() {
        return workersOrCreate();
    }

    @Override
    public void start() {
        if (!running.compareAndSet(false, true)) {
            return;
        }
        List<KafkaEnvelopeConsumerWorker> activeWorkers = workersOrCreate();
        // 每个 worker 一条后台线程，避免某个处理器阻塞时拖慢其他消费组。
        executor = Executors.newFixedThreadPool(activeWorkers.size(), r -> {
            Thread thread = new Thread(r, "kafka-envelope-consumer");
            thread.setDaemon(true);
            return thread;
        });
        for (KafkaEnvelopeConsumerWorker worker : activeWorkers) {
            executor.submit(() -> pollLoop(worker));
        }
    }

    private List<KafkaEnvelopeConsumerWorker> workersOrCreate() {
        List<KafkaEnvelopeConsumerWorker> current = workers;
        if (current != null) {
            return current;
        }
        synchronized (this) {
            if (workers == null) {
                List<KafkaEnvelopeConsumerWorker> created = workersSupplier.get();
                if (created == null || created.isEmpty()) {
                    throw new IllegalStateException("no kafka envelope consumer workers created; check consumer bindings");
                }
                workers = List.copyOf(created);
            }
            return workers;
        }
    }

    private void pollLoop(KafkaEnvelopeConsumerWorker worker) {
        while (running.get()) {
            try {
                worker.pollOnce(pollTimeout);
            } catch (RuntimeException ex) {
                if (!running.get()) {
                    break;
                }
                // Kafka/处理器异常不让 Spring 生命周期退出；退避后继续消费，
                // 具体坏消息由 EnvelopeConsumerProcessor 写入 DLQ。
                log.warn("Kafka envelope consumer poll failed; the worker will retry after backoff", ex);
                sleepBackoff();
            }
        }
    }

    private void sleepBackoff() {
        try {
            Thread.sleep(loopBackoff.toMillis());
        } catch (InterruptedException ex) {
            Thread.currentThread().interrupt();
        }
    }

    @Override
    public void stop() {
        if (!running.compareAndSet(true, false)) {
            return;
        }
        wakeupWorkers();
        if (executor != null) {
            executor.shutdown();
            try {
                if (!executor.awaitTermination(5, TimeUnit.SECONDS)) {
                    executor.shutdownNow();
                    executor.awaitTermination(5, TimeUnit.SECONDS);
                }
            } catch (InterruptedException ex) {
                executor.shutdownNow();
                Thread.currentThread().interrupt();
            }
        }
        closeWorkers();
    }

    @Override
    public void stop(Runnable callback) {
        stop();
        callback.run();
    }

    @Override
    public boolean isRunning() {
        return running.get();
    }

    @Override
    public boolean isAutoStartup() {
        return autoStartup;
    }

    @Override
    public void close() {
        stop();
        closeWorkers();
    }

    private void closeWorkers() {
        if (!closed.compareAndSet(false, true)) {
            return;
        }
        List<KafkaEnvelopeConsumerWorker> current = workers;
        if (current == null) {
            return;
        }
        for (KafkaEnvelopeConsumerWorker worker : current) {
            worker.close();
        }
    }

    private void wakeupWorkers() {
        List<KafkaEnvelopeConsumerWorker> current = workers;
        if (current == null) {
            return;
        }
        for (KafkaEnvelopeConsumerWorker worker : current) {
            worker.wakeup();
        }
    }
}
