package com.lingniu.ingest.core.concurrency;

import com.lingniu.ingest.api.annotation.AsyncBatch;
import com.lingniu.ingest.api.event.VehicleEvent;
import com.lingniu.ingest.core.dispatcher.HandlerDefinition;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;
import java.util.function.UnaryOperator;

/**
 * {@link AsyncBatch} 注解的执行器：把 Handler 方法从"单条调用"变成"批量调用"。
 *
 * <p>每个 Handler 方法对应一个 {@link Batcher}，背后是一个固定容量的 {@link BlockingQueue}
 * 和 {@link AsyncBatch#poolSize()} 条守护虚拟线程。批量触发条件：
 * <ul>
 *   <li>累积 {@link AsyncBatch#size()} 条立即 flush
 *   <li>从第一条入队起等待 {@link AsyncBatch#waitMs()} 毫秒后强制 flush
 * </ul>
 *
 * <p>目标 Handler 方法签名约定：
 * <pre>{@code
 *   @MessageMapping(...)
 *   @AsyncBatch(size = 4000, waitMs = 1000, poolSize = 2)
 *   public List<VehicleEvent> onBatch(List<PayloadType> batch) { ... }
 * }</pre>
 *
 * <p>flush 产出的事件由构造器注入的 {@code eventPublisher} 异步投递（通常是
 * {@link DisruptorEventBus#publish}），不阻塞 Netty EventLoop。
 */
public final class AsyncBatchExecutor implements AutoCloseable {

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

    private final Consumer<VehicleEvent> eventPublisher;
    private final ConcurrentMap<Method, Batcher> batchers = new ConcurrentHashMap<>();

    public AsyncBatchExecutor(Consumer<VehicleEvent> eventPublisher) {
        this.eventPublisher = eventPublisher;
    }

    /** 供 Dispatcher 调用：把单条消息交给目标 Handler 的 batcher。 */
    public void submit(HandlerDefinition def, Object message) {
        Batcher b = batchers.computeIfAbsent(def.method(), m -> new Batcher(def, eventPublisher));
        b.offer(new BatchItem(message, UnaryOperator.identity(), false));
    }

    public void submit(HandlerDefinition def, Object message, UnaryOperator<VehicleEvent> eventTransformer) {
        Batcher b = batchers.computeIfAbsent(def.method(), m -> new Batcher(def, eventPublisher));
        b.offer(new BatchItem(message, eventTransformer == null ? UnaryOperator.identity() : eventTransformer, true));
    }

    @Override
    public void close() {
        batchers.values().forEach(Batcher::close);
        batchers.clear();
        log.info("AsyncBatchExecutor closed");
    }

    // ===== internals =====

    private static final class Batcher implements AutoCloseable {
        private final HandlerDefinition def;
        private final Consumer<VehicleEvent> publisher;
        private final int batchSize;
        private final long waitMs;
        private final BlockingQueue<BatchItem> queue;
        private final Thread[] workers;
        private volatile boolean running = true;

        Batcher(HandlerDefinition def, Consumer<VehicleEvent> publisher) {
            AsyncBatch cfg = def.asyncBatch();
            this.def = def;
            this.publisher = publisher;
            this.batchSize = Math.max(1, cfg.size());
            this.waitMs = Math.max(1, cfg.waitMs());
            this.queue = new ArrayBlockingQueue<>(Math.max(batchSize * 4, 16));
            int pool = Math.max(1, cfg.poolSize());
            this.workers = new Thread[pool];
            for (int i = 0; i < pool; i++) {
                Thread t = Thread.ofVirtual()
                        .name("batcher-" + def.method().getName() + "-" + i)
                        .unstarted(this::loop);
                workers[i] = t;
                t.start();
            }
            log.info("batcher started method={} size={} waitMs={} pool={}",
                    def.method().getName(), batchSize, waitMs, pool);
        }

        void offer(BatchItem item) {
            try {
                if (!queue.offer(item, 100, TimeUnit.MILLISECONDS)) {
                    log.warn("batcher queue full, dropping item for {}", def.method().getName());
                }
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        }

        private void loop() {
            while (running) {
                try {
                    BatchItem first = queue.poll(200, TimeUnit.MILLISECONDS);
                    if (first == null) continue;
                    List<BatchItem> buf = new ArrayList<>(batchSize);
                    buf.add(first);
                    long deadlineNanos = System.nanoTime() + waitMs * 1_000_000L;
                    while (buf.size() < batchSize) {
                        long remain = deadlineNanos - System.nanoTime();
                        if (remain <= 0) break;
                        BatchItem next = queue.poll(remain, TimeUnit.NANOSECONDS);
                        if (next == null) break;
                        buf.add(next);
                    }
                    flush(buf);
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                    return;
                } catch (Exception e) {
                    log.error("batcher loop error method={}", def.method().getName(), e);
                }
            }
        }

        private void flush(List<BatchItem> buf) {
            if (requiresPerItemTransform(buf)) {
                // rawArchiveUri/rawSize 等元数据是每帧唯一的。为了不把第一条帧的元数据错误复制给整批事件，
                // 一旦存在逐条 transformer，就按单条调用 Handler，再逐条发布。
                for (BatchItem item : buf) {
                    flushTransformedItem(item);
                }
                return;
            }

            publishEvents(invoke(buf.stream().map(BatchItem::message).toList()), UnaryOperator.identity());
        }

        private boolean requiresPerItemTransform(List<BatchItem> buf) {
            for (BatchItem item : buf) {
                if (item.transformRequired()) {
                    return true;
                }
            }
            return false;
        }

        private void flushTransformedItem(BatchItem item) {
            publishEvents(invoke(List.of(item.message())), item.eventTransformer());
        }

        @SuppressWarnings("unchecked")
        private List<VehicleEvent> invoke(List<Object> messages) {
            List<VehicleEvent> events;
            try {
                // Handler 仍然只暴露一个批量签名；是否批满或超时由 Batcher 统一处理。
                Object result = def.method().invoke(def.bean(), messages);
                events = switch (result) {
                    case null -> List.of();
                    case List<?> list -> (List<VehicleEvent>) list;
                    case VehicleEvent e -> List.of(e);
                    default -> throw new IllegalStateException(
                            "@AsyncBatch method must return List<VehicleEvent>: " + def.method());
                };
            } catch (InvocationTargetException e) {
                log.error("batcher invoke failed method={}", def.method().getName(), e.getCause());
                return List.of();
            } catch (IllegalAccessException e) {
                log.error("batcher access denied method={}", def.method().getName(), e);
                return List.of();
            }
            return events;
        }

        private void publishEvents(List<VehicleEvent> events, UnaryOperator<VehicleEvent> eventTransformer) {
            for (VehicleEvent e : events) {
                try {
                    publisher.accept(eventTransformer.apply(e));
                } catch (Exception ex) {
                    log.warn("batcher publish failed eventId={}", e.eventId(), ex);
                }
            }
        }

        @Override
        public void close() {
            running = false;
            for (Thread t : workers) t.interrupt();
        }
    }

    private record BatchItem(Object message,
                             UnaryOperator<VehicleEvent> eventTransformer,
                             boolean transformRequired) {
    }
}
