package com.ycwl.basic.integration.kafka.scheduler;

import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.*;

/**
 * QPS 调度器
 * 定期从队列取任务,严格控制 QPS
 * 每秒调度固定数量的任务,确保不超过云端 API 的 QPS 限制
 */
@Slf4j
public class QpsScheduler {
    private final BlockingQueue<Runnable> taskQueue;
    private final ThreadPoolExecutor workerPool;
    private final ScheduledExecutorService scheduler;
    private final int qps;
    private final Semaphore concurrentLimiter; // 并发数限制,防止瞬时抖动

    /**
     * 创建 QPS 调度器
     *
     * @param qps           每秒允许的最大请求数
     * @param maxConcurrent 最大并发数
     * @param workerPool    工作线程池
     */
    public QpsScheduler(int qps, int maxConcurrent, ThreadPoolExecutor workerPool) {
        this.qps = qps;
        this.taskQueue = new LinkedBlockingQueue<>();
        this.workerPool = workerPool;
        this.scheduler = new ScheduledThreadPoolExecutor(1, r -> {
            Thread thread = new Thread(r);
            thread.setName("qps-scheduler-" + workerPool.getThreadFactory().newThread(() -> {}).getName());
            thread.setDaemon(true);
            return thread;
        });
        this.concurrentLimiter = new Semaphore(maxConcurrent);

        // 每秒调度一次,取 qps 个任务
        scheduler.scheduleAtFixedRate(this::dispatch, 0, 1, TimeUnit.SECONDS);

        log.info("QPS调度器已启动: qps={}, maxConcurrent={}", qps, maxConcurrent);
    }

    /**
     * 调度任务
     * 每秒执行一次,从队列中取出 qps 个任务提交到工作线程池
     */
    private void dispatch() {
        int dispatched = 0;
        for (int i = 0; i < qps; i++) {
            Runnable task = taskQueue.poll();
            if (task == null) {
                break; // 队列为空,结束本次调度
            }

            // 检查并发数限制
            if (concurrentLimiter.tryAcquire()) {
                try {
                    workerPool.execute(() -> {
                        try {
                            task.run();
                        } catch (Exception e) {
                            log.error("任务执行失败", e);
                        } finally {
                            concurrentLimiter.release();
                        }
                    });
                    dispatched++;
                } catch (RejectedExecutionException e) {
                    // 线程池拒绝,释放并发许可,任务丢弃
                    concurrentLimiter.release();
                    log.warn("任务被线程池拒绝", e);
                }
            } else {
                // 并发数已满,任务放回队列,等待下次调度
                taskQueue.offer(task);
                break;
            }
        }

        if (dispatched > 0 || taskQueue.size() > 0) {
            log.debug("QPS调度完成: dispatched={}, remainQueue={}, availableConcurrent={}",
                    dispatched, taskQueue.size(), concurrentLimiter.availablePermits());
        }
    }

    /**
     * 提交任务到调度队列
     *
     * @param task 待执行的任务
     * @return 是否成功提交
     */
    public boolean submit(Runnable task) {
        return taskQueue.offer(task);
    }

    /**
     * 获取队列中等待调度的任务数量
     *
     * @return 队列大小
     */
    public int getQueueSize() {
        return taskQueue.size();
    }

    /**
     * 关闭调度器
     */
    public void shutdown() {
        scheduler.shutdown();
        log.info("QPS调度器已关闭, qps={}", qps);
    }
}
