package tech.easyflow.ai.easyagentsflow.config;

import com.easyagents.flow.core.util.IoBulkhead;
import io.micrometer.core.instrument.Gauge;
import io.micrometer.core.instrument.MeterRegistry;
import org.springframework.stereotype.Component;

import java.util.LinkedHashMap;
import java.util.Map;
import java.util.function.ToDoubleFunction;

/**
 * 将工作流各类 I/O 隔离器运行状态接入 Micrometer。
 */
@Component
public class WorkflowIoBulkheadMetrics {

    /**
     * 注册工作流 I/O 隔离器指标。
     *
     * @param meterRegistry Micrometer 注册表
     * @param properties 工作流 I/O 配置
     */
    public WorkflowIoBulkheadMetrics(
            MeterRegistry meterRegistry,
            WorkflowIoProperties properties) {
        IoBulkhead.configure(
                properties.getHttp().toSettings(),
                properties.getDataset().toSettings(),
                properties.getStorage().toSettings(),
                properties.getDocumentParse().toSettings(),
                properties.getResponseAggregation().toSettings());
        Map<String, IoBulkhead> lanes = new LinkedHashMap<>();
        lanes.put("http", IoBulkhead.shared());
        lanes.put("dataset", IoBulkhead.dataset());
        lanes.put("storage", IoBulkhead.storage());
        lanes.put("document_parse", IoBulkhead.documentParse());
        lanes.put(
                "response_aggregation",
                IoBulkhead.responseAggregation());
        lanes.forEach((lane, bulkhead) -> registerLane(
                meterRegistry, lane, bulkhead));
    }

    /**
     * 注册一个资源 lane 的核心容量和等待指标。
     *
     * @param registry 指标注册表
     * @param lane lane 名
     * @param bulkhead 隔离器
     */
    private void registerLane(
            MeterRegistry registry,
            String lane,
            IoBulkhead bulkhead) {
        gauge(registry, lane, bulkhead, "in_flight",
                snapshot -> snapshot.inFlightCount());
        gauge(registry, lane, bulkhead, "acquired_total",
                snapshot -> snapshot.acquiredCount());
        gauge(registry, lane, bulkhead, "rejected_total",
                snapshot -> snapshot.rejectedCount());
        gauge(registry, lane, bulkhead, "wait_nanos_total",
                snapshot -> snapshot.totalWaitNanos());
        gauge(registry, lane, bulkhead, "available_permits",
                snapshot -> snapshot.availableGlobalPermits());
        gauge(registry, lane, bulkhead, "tracked_targets",
                snapshot -> snapshot.trackedTargetCount());
    }

    /**
     * 注册从快照读取的 Gauge。
     *
     * @param registry 指标注册表
     * @param lane lane 名
     * @param bulkhead 隔离器
     * @param metric 指标后缀
     * @param valueFunction 快照取值函数
     */
    private void gauge(
            MeterRegistry registry,
            String lane,
            IoBulkhead bulkhead,
            String metric,
            ToDoubleFunction<IoBulkhead.Snapshot> valueFunction) {
        Gauge.builder(
                        "easyflow.workflow.io." + metric,
                        bulkhead,
                        value -> valueFunction.applyAsDouble(
                                value.snapshot()))
                .tag("lane", lane)
                .register(registry);
    }
}
