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

import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;

@Component
public class AiRedisKeys {

    private static final String CONFIG_KEY_PREFIX = "AI:CONFIG:";
    private static final String RUNTIME_KEY_PREFIX = "AI:RUNTIME:";
    private static final String MEMORY_KEY_PREFIX = "AI:MEMORY:";
    private static final String SESSIONS_KEY_PREFIX = "AI:SESSIONS:";
    private static final String MCP_PREVIEW_KEY_PREFIX = "AI:MCP:PREVIEW:";
    private static final String MCP_HEALTH_KEY_PREFIX = "AI:MCP:HEALTH:";
    private static final String STREAM_STATE_KEY_PREFIX = "AI:STREAM:";
    private static final String TOOL_RESULT_KEY_PREFIX = "AI:TOOL:RESULT:";
    private static final String RATE_LIMIT_KEY_PREFIX = "AI:RATE:";
    private static final String OBSERVE_STATS_KEY_PREFIX = "AI:OBSERVE:STATS:";
    private static final String OBSERVE_TOOL_RANK_KEY_PREFIX = "AI:OBSERVE:TOOL:RANK:";
    private static final String OBSERVE_TOOL_FAIL_RANK_KEY_PREFIX = "AI:OBSERVE:TOOL:FAIL:RANK:";
    private static final String CONFIG_INDEX_PREFIX = "AI:IDX:CONFIG:";
    private static final String CONVERSATION_INDEX_PREFIX = "AI:IDX:CONVERSATION:";
    private static final String TENANT_RUNTIME_INDEX_PREFIX = "AI:IDX:RUNTIME:TENANT:";

    public String buildConfigKey(Long tenantId, String promptCode, Long aiParamId) {
        return CONFIG_KEY_PREFIX + tenantId + ":" + safeToken(promptCode) + ":" + aiParamToken(aiParamId);
    }

    public String buildRuntimeKey(Long tenantId, Long userId, String promptCode, Long sessionId, Long aiParamId) {
        return RUNTIME_KEY_PREFIX + tenantId + ":" + userId + ":" + safeToken(promptCode) + ":" + sessionToken(sessionId) + ":" + aiParamToken(aiParamId);
    }

    public String buildMemoryKey(Long tenantId, Long userId, String promptCode, Long sessionId) {
        return MEMORY_KEY_PREFIX + tenantId + ":" + userId + ":" + safeToken(promptCode) + ":" + sessionToken(sessionId);
    }

    public String buildSessionsKey(Long tenantId, Long userId, String promptCode, String keyword) {
        return SESSIONS_KEY_PREFIX + tenantId + ":" + userId + ":" + safeToken(promptCode) + ":" + safeToken(keyword);
    }

    public String buildMcpPreviewKey(Long tenantId, Long mountId) {
        return MCP_PREVIEW_KEY_PREFIX + tenantId + ":" + mountId;
    }

    public String buildMcpHealthKey(Long tenantId, Long mountId) {
        return MCP_HEALTH_KEY_PREFIX + tenantId + ":" + mountId;
    }

    public String buildStreamStateKey(Long tenantId, String requestId) {
        return STREAM_STATE_KEY_PREFIX + tenantId + ":" + safeToken(requestId);
    }

    public String buildToolResultKey(Long tenantId, String requestId, String toolName, String toolInput) {
        return TOOL_RESULT_KEY_PREFIX + tenantId + ":" + safeToken(requestId) + ":" + safeToken(toolName) + ":" + digest(toolInput);
    }

    public String buildRateLimitKey(Long tenantId, Long userId, String promptCode) {
        return RATE_LIMIT_KEY_PREFIX + tenantId + ":" + userId + ":" + safeToken(promptCode);
    }

    public String buildObserveStatsKey(Long tenantId) {
        return OBSERVE_STATS_KEY_PREFIX + tenantId;
    }

    public String buildToolRankKey(Long tenantId) {
        return OBSERVE_TOOL_RANK_KEY_PREFIX + tenantId;
    }

    public String buildToolFailRankKey(Long tenantId) {
        return OBSERVE_TOOL_FAIL_RANK_KEY_PREFIX + tenantId;
    }

    public String buildConfigIndexKey(Long tenantId) {
        return CONFIG_INDEX_PREFIX + tenantId;
    }

    public String buildConversationIndexKey(Long tenantId, Long userId) {
        return CONVERSATION_INDEX_PREFIX + tenantId + ":" + userId;
    }

    public String buildTenantRuntimeIndexKey(Long tenantId) {
        return TENANT_RUNTIME_INDEX_PREFIX + tenantId;
    }

    public String sessionToken(Long sessionId) {
        return sessionId == null ? "LATEST" : String.valueOf(sessionId);
    }

    public String aiParamToken(Long aiParamId) {
        return aiParamId == null ? "DEFAULT" : String.valueOf(aiParamId);
    }

    public String safeToken(String source) {
        if (!StringUtils.hasText(source)) {
            return "_";
        }
        return URLEncoder.encode(source.trim(), StandardCharsets.UTF_8);
    }

    public String digest(String source) {
        try {
            MessageDigest messageDigest = MessageDigest.getInstance("SHA-256");
            byte[] bytes = messageDigest.digest((source == null ? "" : source).getBytes(StandardCharsets.UTF_8));
            StringBuilder builder = new StringBuilder();
            for (byte value : bytes) {
                builder.append(String.format("%02x", value));
            }
            return builder.toString();
        } catch (Exception e) {
            return safeToken(source);
        }
    }
}