package com.nanri.aiimage.config;

import io.micrometer.core.instrument.MeterRegistry;
import io.micrometer.core.instrument.Timer;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.core.task.TaskExecutor;
import org.springframework.core.task.TaskRejectedException;

import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * Task 75：虚拟线程任务排队闸门。Coze 执行池的信号量只限制"正在执行"的
 * 并发度，提交侧仍会在虚拟线程里无限排队。此闸门在提交时统计"已受理未启动"
 * 的等待数，达到上限立即拒绝并记录指标，防止等待队列无界堆积：
 * <ul>
 *   <li>等待数：提交时 +1，任务执行（或执行器拒绝）时 -1；上限钳制到 [1, +∞)；</li>
 *   <li>指标：等待耗时、执行耗时、拒绝次数（queue-full / delegate-rejected / invalid-input）；</li>
 *   <li>执行器不可用（provider 为空）时拒绝新提交，不产生死等任务。</li>
 * </ul>
 */
@Slf4j
public class CozeTaskQueueGate implements TaskExecutor {

    private final TaskExecutor delegate;
    private final int maxWaiting;
    private final ObjectProvider<MeterRegistry> meterRegistryProvider;
    private final AtomicInteger waiting = new AtomicInteger();

    public CozeTaskQueueGate(TaskExecutor delegate, int maxWaiting,
                             ObjectProvider<MeterRegistry> meterRegistryProvider) {
        this.delegate = delegate;
        this.maxWaiting = Math.max(1, maxWaiting);
        this.meterRegistryProvider = meterRegistryProvider;
    }

    public int waiting() {
        return waiting.get();
    }

    @Override
    public void execute(Runnable command) {
        if (command == null) {
            recordRejected("invalid-input");
            throw new IllegalArgumentException("coze 任务不能为 null");
        }
        if (waiting.get() >= maxWaiting) {
            recordRejected("queue-full");
            log.warn("[coze-task][gate] waiting queue full, reject submit waiting={} limit={}",
                    waiting.get(), maxWaiting);
            throw new TaskRejectedException("coze 等待队列已满，limit=" + maxWaiting
                    + ", waiting=" + waiting.get());
        }
        waiting.incrementAndGet();
        long queuedAt = System.nanoTime();
        try {
            delegate.execute(() -> {
                long waitNanos = System.nanoTime() - queuedAt;
                try {
                    run(command);
                } finally {
                    waiting.decrementAndGet();
                    recordQueueWait(waitNanos);
                }
            });
        } catch (RuntimeException ex) {
            waiting.decrementAndGet();
            recordQueueWait(System.nanoTime() - queuedAt);
            recordRejected("delegate-rejected");
            log.warn("[coze-task][gate] delegate rejected submit waiting={} limit={} msg={}",
                    waiting.get(), maxWaiting, ex.getMessage(), ex);
            throw ex;
        }
    }

    private void run(Runnable command) {
        long startedAt = System.nanoTime();
        try {
            command.run();
        } finally {
            recordExecution(System.nanoTime() - startedAt);
        }
    }

    private void recordRejected(String reason) {
        MeterRegistry registry = meterRegistry();
        if (registry != null) {
            registry.counter("aiimage.coze-task.submit.rejected.total", "reason", reason).increment();
        }
    }

    private void recordQueueWait(long waitNanos) {
        MeterRegistry registry = meterRegistry();
        if (registry != null && waitNanos >= 0L) {
            Timer.builder("aiimage.coze-task.queue.wait.duration")
                    .register(registry)
                    .record(waitNanos, TimeUnit.NANOSECONDS);
        }
    }

    private void recordExecution(long durationNanos) {
        MeterRegistry registry = meterRegistry();
        if (registry != null && durationNanos >= 0L) {
            Timer.builder("aiimage.coze-task.execution.duration")
                    .register(registry)
                    .record(durationNanos, TimeUnit.NANOSECONDS);
        }
    }

    private MeterRegistry meterRegistry() {
        return meterRegistryProvider == null ? null : meterRegistryProvider.getIfAvailable();
    }
}