package com.vincent.rsf.server.ai.store;

import com.vincent.rsf.server.ai.dto.AiObserveStatsDto;
import com.vincent.rsf.server.ai.store.support.AiRedisExecutor;
import com.vincent.rsf.server.ai.store.support.AiRedisKeys;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

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

@Component
@RequiredArgsConstructor
public class AiObserveStatsStore {

    private static final String FIELD_CALL_COUNT = "callCount";
    private static final String FIELD_SUCCESS_COUNT = "successCount";
    private static final String FIELD_FAILURE_COUNT = "failureCount";
    private static final String FIELD_ELAPSED_SUM = "elapsedSum";
    private static final String FIELD_ELAPSED_COUNT = "elapsedCount";
    private static final String FIELD_FIRST_TOKEN_SUM = "firstTokenSum";
    private static final String FIELD_FIRST_TOKEN_COUNT = "firstTokenCount";
    private static final String FIELD_TOTAL_TOKENS_SUM = "totalTokensSum";
    private static final String FIELD_TOTAL_TOKENS_COUNT = "totalTokensCount";
    private static final String FIELD_TOOL_CALL_COUNT = "toolCallCount";
    private static final String FIELD_TOOL_SUCCESS_COUNT = "toolSuccessCount";
    private static final String FIELD_TOOL_FAILURE_COUNT = "toolFailureCount";

    private final AiRedisExecutor aiRedisExecutor;
    private final AiRedisKeys aiRedisKeys;

    public void recordObserveCallStarted(Long tenantId) {
        aiRedisExecutor.executeVoid(jedis -> jedis.hincrBy(aiRedisKeys.buildObserveStatsKey(tenantId), FIELD_CALL_COUNT, 1));
    }

    public void recordObserveCallFinished(Long tenantId, String status, Long elapsedMs, Long firstTokenLatencyMs, Integer totalTokens) {
        aiRedisExecutor.executeVoid(jedis -> {
            String key = aiRedisKeys.buildObserveStatsKey(tenantId);
            if ("COMPLETED".equals(status)) {
                jedis.hincrBy(key, FIELD_SUCCESS_COUNT, 1);
            } else if ("FAILED".equals(status)) {
                jedis.hincrBy(key, FIELD_FAILURE_COUNT, 1);
            }
            if (elapsedMs != null) {
                jedis.hincrBy(key, FIELD_ELAPSED_SUM, elapsedMs);
                jedis.hincrBy(key, FIELD_ELAPSED_COUNT, 1);
            }
            if (firstTokenLatencyMs != null) {
                jedis.hincrBy(key, FIELD_FIRST_TOKEN_SUM, firstTokenLatencyMs);
                jedis.hincrBy(key, FIELD_FIRST_TOKEN_COUNT, 1);
            }
            if (totalTokens != null) {
                jedis.hincrBy(key, FIELD_TOTAL_TOKENS_SUM, totalTokens.longValue());
                jedis.hincrBy(key, FIELD_TOTAL_TOKENS_COUNT, 1);
            }
        });
    }

    public void recordObserveToolCall(Long tenantId, String toolName, String status) {
        aiRedisExecutor.executeVoid(jedis -> {
            String key = aiRedisKeys.buildObserveStatsKey(tenantId);
            jedis.hincrBy(key, FIELD_TOOL_CALL_COUNT, 1);
            if ("COMPLETED".equals(status)) {
                jedis.hincrBy(key, FIELD_TOOL_SUCCESS_COUNT, 1);
            } else if ("FAILED".equals(status)) {
                jedis.hincrBy(key, FIELD_TOOL_FAILURE_COUNT, 1);
            }
            if (StringUtils.hasText(toolName)) {
                jedis.zincrby(aiRedisKeys.buildToolRankKey(tenantId), 1D, toolName);
                if ("FAILED".equals(status)) {
                    jedis.zincrby(aiRedisKeys.buildToolFailRankKey(tenantId), 1D, toolName);
                }
            }
        });
    }

    public AiObserveStatsDto getObserveStats(Long tenantId, Supplier<AiObserveStatsDto> fallbackLoader) {
        AiObserveStatsDto cached = readObserveStats(tenantId);
        if (cached != null) {
            return cached;
        }
        AiObserveStatsDto snapshot = fallbackLoader.get();
        if (snapshot != null) {
            seedObserveStats(tenantId, snapshot);
        }
        return snapshot;
    }

    private AiObserveStatsDto readObserveStats(Long tenantId) {
        Map<String, String> fields = aiRedisExecutor.execute(jedis -> {
            String key = aiRedisKeys.buildObserveStatsKey(tenantId);
            if (!jedis.exists(key)) {
                return null;
            }
            return jedis.hgetAll(key);
        });
        if (fields == null || fields.isEmpty()) {
            return null;
        }
        long callCount = parseLong(fields.get(FIELD_CALL_COUNT));
        long successCount = parseLong(fields.get(FIELD_SUCCESS_COUNT));
        long failureCount = parseLong(fields.get(FIELD_FAILURE_COUNT));
        long elapsedSum = parseLong(fields.get(FIELD_ELAPSED_SUM));
        long elapsedCount = parseLong(fields.get(FIELD_ELAPSED_COUNT));
        long firstTokenSum = parseLong(fields.get(FIELD_FIRST_TOKEN_SUM));
        long firstTokenCount = parseLong(fields.get(FIELD_FIRST_TOKEN_COUNT));
        long totalTokensSum = parseLong(fields.get(FIELD_TOTAL_TOKENS_SUM));
        long totalTokensCount = parseLong(fields.get(FIELD_TOTAL_TOKENS_COUNT));
        long toolCallCount = parseLong(fields.get(FIELD_TOOL_CALL_COUNT));
        long toolSuccessCount = parseLong(fields.get(FIELD_TOOL_SUCCESS_COUNT));
        long toolFailureCount = parseLong(fields.get(FIELD_TOOL_FAILURE_COUNT));
        return AiObserveStatsDto.builder()
                .callCount(callCount)
                .successCount(successCount)
                .failureCount(failureCount)
                .avgElapsedMs(elapsedCount == 0 ? 0L : elapsedSum / elapsedCount)
                .avgFirstTokenLatencyMs(firstTokenCount == 0 ? 0L : firstTokenSum / firstTokenCount)
                .totalTokens(totalTokensSum)
                .avgTotalTokens(totalTokensCount == 0 ? 0L : totalTokensSum / totalTokensCount)
                .toolCallCount(toolCallCount)
                .toolSuccessCount(toolSuccessCount)
                .toolFailureCount(toolFailureCount)
                .toolSuccessRate(toolCallCount == 0 ? 0D : (toolSuccessCount * 100D) / toolCallCount)
                .build();
    }

    private void seedObserveStats(Long tenantId, AiObserveStatsDto snapshot) {
        aiRedisExecutor.executeVoid(jedis -> {
            String key = aiRedisKeys.buildObserveStatsKey(tenantId);
            Map<String, String> values = new LinkedHashMap<>();
            values.put(FIELD_CALL_COUNT, String.valueOf(defaultLong(snapshot.getCallCount())));
            values.put(FIELD_SUCCESS_COUNT, String.valueOf(defaultLong(snapshot.getSuccessCount())));
            values.put(FIELD_FAILURE_COUNT, String.valueOf(defaultLong(snapshot.getFailureCount())));
            values.put(FIELD_ELAPSED_SUM, String.valueOf(defaultLong(snapshot.getAvgElapsedMs()) * defaultLong(snapshot.getCallCount())));
            values.put(FIELD_ELAPSED_COUNT, String.valueOf(defaultLong(snapshot.getCallCount())));
            values.put(FIELD_FIRST_TOKEN_SUM, String.valueOf(defaultLong(snapshot.getAvgFirstTokenLatencyMs()) * defaultLong(snapshot.getCallCount())));
            values.put(FIELD_FIRST_TOKEN_COUNT, String.valueOf(defaultLong(snapshot.getCallCount())));
            values.put(FIELD_TOTAL_TOKENS_SUM, String.valueOf(defaultLong(snapshot.getTotalTokens())));
            values.put(FIELD_TOTAL_TOKENS_COUNT, String.valueOf(defaultLong(snapshot.getCallCount())));
            values.put(FIELD_TOOL_CALL_COUNT, String.valueOf(defaultLong(snapshot.getToolCallCount())));
            values.put(FIELD_TOOL_SUCCESS_COUNT, String.valueOf(defaultLong(snapshot.getToolSuccessCount())));
            values.put(FIELD_TOOL_FAILURE_COUNT, String.valueOf(defaultLong(snapshot.getToolFailureCount())));
            jedis.hset(key, values);
        });
    }

    private long parseLong(String source) {
        if (!StringUtils.hasText(source)) {
            return 0L;
        }
        try {
            return Long.parseLong(source);
        } catch (Exception e) {
            return 0L;
        }
    }

    private long defaultLong(Long value) {
        return value == null ? 0L : value;
    }
}