package org.dromara.ai.monitor;

import java.time.LocalDateTime;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.LongAdder;

/**
 * 基础监控指标类
 *
 * @author pengles
 * @since 2024-01-01
 */
public abstract class BaseMetrics {

    protected final AtomicLong totalRequests = new AtomicLong(0);
    protected final AtomicLong successRequests = new AtomicLong(0);
    protected final AtomicLong failedRequests = new AtomicLong(0);
    protected final LongAdder totalResponseTime = new LongAdder();
    protected final AtomicLong minResponseTime = new AtomicLong(Long.MAX_VALUE);
    protected final AtomicLong maxResponseTime = new AtomicLong(0);
    protected final AtomicLong totalInputTokens = new AtomicLong(0);
    protected final AtomicLong totalOutputTokens = new AtomicLong(0);
    protected final Map<String, AtomicLong> errorCounts = new ConcurrentHashMap<>();
    protected final Map<String, Long> activeRequests = new ConcurrentHashMap<>();
    protected volatile LocalDateTime lastUpdateTime = LocalDateTime.now();

    /**
     * 增加总请求数
     */
    public void incrementTotalRequests() {
        totalRequests.incrementAndGet();
        updateLastUpdateTime();
    }

    /**
     * 增加成功请求数
     */
    public void incrementSuccessRequests() {
        successRequests.incrementAndGet();
        updateLastUpdateTime();
    }

    /**
     * 增加失败请求数
     */
    public void incrementFailedRequests() {
        failedRequests.incrementAndGet();
        updateLastUpdateTime();
    }

    /**
     * 添加响应时间
     *
     * @param responseTime 响应时间(毫秒)
     */
    public void addResponseTime(long responseTime) {
        totalResponseTime.add(responseTime);

        // 更新最小响应时间
        long currentMin = minResponseTime.get();
        while (responseTime < currentMin && !minResponseTime.compareAndSet(currentMin, responseTime)) {
            currentMin = minResponseTime.get();
        }

        // 更新最大响应时间
        long currentMax = maxResponseTime.get();
        while (responseTime > currentMax && !maxResponseTime.compareAndSet(currentMax, responseTime)) {
            currentMax = maxResponseTime.get();
        }

        updateLastUpdateTime();
    }

    /**
     * 添加Token使用量
     *
     * @param inputTokens  输入Token数
     * @param outputTokens 输出Token数
     */
    public void addTokenUsage(int inputTokens, int outputTokens) {
        totalInputTokens.addAndGet(inputTokens);
        totalOutputTokens.addAndGet(outputTokens);
        updateLastUpdateTime();
    }

    /**
     * 增加错误计数
     *
     * @param errorCode 错误码
     */
    public void incrementErrorCount(String errorCode) {
        if (errorCode != null && !errorCode.isEmpty()) {
            errorCounts.computeIfAbsent(errorCode, k -> new AtomicLong(0)).incrementAndGet();
            updateLastUpdateTime();
        }
    }

    /**
     * 添加活跃请求
     *
     * @param requestId 请求ID
     * @param startTime 开始时间
     */
    public void addActiveRequest(String requestId, long startTime) {
        activeRequests.put(requestId, startTime);
        updateLastUpdateTime();
    }

    /**
     * 移除活跃请求
     *
     * @param requestId 请求ID
     * @return 开始时间
     */
    public Long removeActiveRequest(String requestId) {
        Long startTime = activeRequests.remove(requestId);
        if (startTime != null) {
            updateLastUpdateTime();
        }
        return startTime;
    }

    /**
     * 清理过期的活跃请求
     *
     * @param expireTime 过期时间
     */
    public void cleanupExpiredRequests(long expireTime) {
        activeRequests.entrySet().removeIf(entry -> entry.getValue() < expireTime);
    }

    /**
     * 获取总请求数
     */
    public long getTotalRequests() {
        return totalRequests.get();
    }

    /**
     * 获取成功请求数
     */
    public long getSuccessRequests() {
        return successRequests.get();
    }

    /**
     * 获取失败请求数
     */
    public long getFailedRequests() {
        return failedRequests.get();
    }

    /**
     * 获取活跃请求数
     */
    public int getActiveRequestCount() {
        return activeRequests.size();
    }

    /**
     * 获取成功率
     */
    public double getSuccessRate() {
        long total = totalRequests.get();
        if (total == 0) {
            return 0.0;
        }
        return (double) successRequests.get() / total;
    }

    /**
     * 获取平均响应时间
     */
    public double getAvgResponseTime() {
        long completedRequests = successRequests.get() + failedRequests.get();
        if (completedRequests == 0) {
            return 0.0;
        }
        return (double) totalResponseTime.sum() / completedRequests;
    }

    /**
     * 获取最小响应时间
     */
    public long getMinResponseTime() {
        long min = minResponseTime.get();
        return min == Long.MAX_VALUE ? 0 : min;
    }

    /**
     * 获取最大响应时间
     */
    public long getMaxResponseTime() {
        return maxResponseTime.get();
    }

    /**
     * 获取总输入Token数
     */
    public long getTotalInputTokens() {
        return totalInputTokens.get();
    }

    /**
     * 获取总输出Token数
     */
    public long getTotalOutputTokens() {
        return totalOutputTokens.get();
    }

    /**
     * 获取错误计数
     */
    public Map<String, Long> getErrorCounts() {
        Map<String, Long> result = new ConcurrentHashMap<>();
        errorCounts.forEach((key, value) -> result.put(key, value.get()));
        return result;
    }

    /**
     * 获取最后更新时间
     */
    public LocalDateTime getLastUpdateTime() {
        return lastUpdateTime;
    }

    /**
     * 重置指标
     */
    public void reset() {
        totalRequests.set(0);
        successRequests.set(0);
        failedRequests.set(0);
        totalResponseTime.reset();
        minResponseTime.set(Long.MAX_VALUE);
        maxResponseTime.set(0);
        totalInputTokens.set(0);
        totalOutputTokens.set(0);
        errorCounts.clear();
        activeRequests.clear();
        updateLastUpdateTime();
    }

    /**
     * 更新最后更新时间
     */
    private void updateLastUpdateTime() {
        lastUpdateTime = LocalDateTime.now();
    }
}