package org.kne.cloud.network.congestion;

import org.jctools.queues.MessagePassingQueue;
import org.jctools.queues.varhandle.MpscVarHandleArrayQueue;
import org.kne.cloud.network.ThreadTool;
import org.kne.opencl64.Releaser;

import java.lang.ref.Cleaner;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.locks.LockSupport;
import java.util.function.Consumer;

/**
 * 基于 MPSC 队列的高性能消息批量发送器
 */
public class MpscMessageBatcher<T> implements MessageBatcher<T> {
    private static final Cleaner clr = Cleaner.create();
    
    private final MpscMessageBatcher0<T> impl;
    private final MpscMessageBatcherReleaser<T> releaser;
    
    public MpscMessageBatcher(int batchSize, long maxDelayNanos) {
        this.impl = new MpscMessageBatcher0<>(batchSize, maxDelayNanos, 65536);
        this.releaser = new MpscMessageBatcherReleaser<>(impl);
        clr.register(this, releaser);
    }
    
    @Override
    public int getBatchSize() { return impl.getBatchSize(); }
    
    @Override
    public long getMaxDelay() { return impl.getMaxDelay(); }
    
    @Override
    public void setConsumer(Consumer<List<T>> consumer) { impl.setConsumer(consumer); }
    
    @Override
    public void putMessage(T message) { impl.putMessage(message); }
    
    @Override
    public void putMessages(List<T> messages) { impl.putMessages(messages); }
    
    @Override
    public void flush() { impl.flush(); }
    
    @Override
    public int getQueueSize() { return impl.getQueueSize(); }
    
    @Override
    public boolean isClosed() { return impl.isClosed(); }
    
    @Override
    public void close() { impl.close(); }
    
    public static void main(String[] args) throws InterruptedException {
        MpscMessageBatcher<Integer> batcher = 
            new MpscMessageBatcher<>(10, 1_000_000L);
        batcher.setConsumer(System.out::println);
        
        int n = 0;
        for (;;) {
            batcher.putMessage(n++);
            batcher.putMessage(n++);
           batcher.putMessage(n++);
            batcher.putMessage(n++);
            Thread.sleep(1);
        }
    }
}

class MpscMessageBatcherReleaser<T> extends Releaser<MpscMessageBatcher0<T>> {
    public MpscMessageBatcherReleaser(MpscMessageBatcher0<T> resource) {
        super(resource);
    }
    
    @Override
    protected void release(MpscMessageBatcher0<T> resource) {
        resource.close();
    }
}

class MpscMessageBatcher0<T> implements Runnable, MessageBatcher<T> {
    
    private final MessagePassingQueue<T> queue;
    private final int batchSize;
    private final long maxDelay;
    private volatile Consumer<List<T>> consumer;
    private final Thread batchThread;
    private final AtomicBoolean closed = new AtomicBoolean(false);
	private long firstTime;
    
    
    public MpscMessageBatcher0(int batchSize, long maxDelay, int queueCapacity) {
        this.batchSize = batchSize;
        this.maxDelay = maxDelay;
        this.queue = new MpscVarHandleArrayQueue<>(queueCapacity);
        
        this.batchThread = ThreadTool.makeVDaemonThread("MpscMessageBatcher", this);
        this.batchThread.start();
    }
    
    @Override
    public void setConsumer(Consumer<List<T>> consumer) {
        this.consumer = consumer;
    }
    
    @Override
    public void putMessage(T message) {
        if (message == null || closed.get()) return;
        // 无锁入队
        while (!queue.offer(message)) {
            // 队列满时，主动尝试 check 一次（自旋等待）
        	Thread.yield();
        }
    	check(false);
    }
    
    @Override
    public void putMessages(List<T> messages) {
        for (T msg : messages) {
            putMessage(msg);
        }
    }
    
    @Override
    public void run() {
    	while(!closed.get()) {
        		check(false);
        	// 更精确的等待策略
            long sleepTimeNanos = calculateSleepTime();
            if (sleepTimeNanos > 0) {
                LockSupport.parkNanos(sleepTimeNanos);
            } else {
                // 避免忙等待
                LockSupport.parkNanos(1_000_000L); // 1ms
            }
    	}
    }
    
    private void check(boolean force) {
    	long currTime=System.nanoTime();
		if(queue.size()>=batchSize||(currTime-firstTime)>maxDelay||force) {
			if(consumer!=null) {
				ArrayList<T>batch=new ArrayList<>(batchSize);
				queue.drain(batch::add, batchSize);
				if(!batch.isEmpty())
				try {
				consumer.accept(batch);
				}catch(Exception e) {
					e.printStackTrace();
				}
				firstTime=currTime;
			}
		}
	}
    
    /**
     * 计算需要等待的时间
     * @return 等待时间(纳秒)
     */
    private long calculateSleepTime() {
      
            
            long elapsed = System.nanoTime() - firstTime;
            long remaining = maxDelay - elapsed;
            
            if (remaining <= 0) {
                return 0; // 立即处理
            }
            
            // 返回剩余时间或10ms中的较小值
            return Math.min(remaining, 10_000_000L);
    }
    
    @Override
    public void flush() {
    	check(true);
    }
    
    @Override
    public int getQueueSize() {
        return queue.size();
    }
    
    @Override
    public boolean isClosed() {
        return closed.get();
    }
    
    @Override
    public void close() {
        closed.set(true);
        LockSupport.unpark(batchThread);
        flush(); // 确保剩余消息被处理
    }
    
    @Override
    public int getBatchSize() {
        return batchSize;
    }
    
    @Override
    public long getMaxDelay() {
        return maxDelay;
    }
}