package tech.easyflow.ai.service.capability;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.core.io.ClassPathResource;
import org.springframework.stereotype.Component;
import tech.easyflow.ai.entity.Model;

import java.io.IOException;
import java.io.InputStream;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Locale;
import java.util.Map;
import java.util.Optional;
import java.util.Set;

/**
 * 从静态 {@code llm.json} 加载模型能力，并建立常量时间查询索引。
 */
@Component
public class ModelCapabilityCatalog {

    /** classpath 模型能力库资源。 */
    private static final String CATALOG_RESOURCE = "llm.json";

    /** 按规范化完整模型 ID 建立的目录索引。 */
    private final Map<String, ModelCatalogMetadata> metadataById;
    /** 仅在模型短 ID 唯一时建立的目录别名索引。 */
    private final Map<String, ModelCatalogMetadata> metadataByAlias;

    /**
     * 加载并索引静态模型能力库。
     *
     * @param objectMapper JSON 解析器
     * @throws IllegalStateException 静态资源缺失或格式非法时抛出
     */
    public ModelCapabilityCatalog(ObjectMapper objectMapper) {
        Map<String, ModelCatalogMetadata> fullIdIndex = new HashMap<>();
        Map<String, ModelCatalogMetadata> aliasCandidates = new HashMap<>();
        Set<String> ambiguousAliases = new HashSet<>();
        loadCatalog(objectMapper, fullIdIndex, aliasCandidates, ambiguousAliases);
        ambiguousAliases.forEach(aliasCandidates::remove);
        this.metadataById = Map.copyOf(fullIdIndex);
        this.metadataByAlias = Map.copyOf(aliasCandidates);
    }

    /**
     * 查询模型能力。
     *
     * @param providerType EasyFlow 供应商类型
     * @param modelId 用户配置的模型 ID
     * @return 命中的模型能力
     */
    public Optional<ModelCapabilityResolution> find(String providerType, String modelId) {
        return findMetadata(providerType, modelId).map(ModelCatalogMetadata::getCapability);
    }

    /**
     * 查询模型目录元数据。
     *
     * @param providerType EasyFlow 供应商类型
     * @param modelId 用户配置的模型 ID
     * @return 命中的模型目录元数据
     */
    public Optional<ModelCatalogMetadata> findMetadata(String providerType, String modelId) {
        String normalizedId = normalize(modelId);
        if (normalizedId.isEmpty()) {
            return Optional.empty();
        }

        ModelCatalogMetadata direct = metadataById.get(normalizedId);
        if (direct != null) {
            return Optional.of(direct);
        }

        String providerPrefix = providerPrefix(providerType);
        if (!providerPrefix.isEmpty() && !normalizedId.contains("/")) {
            direct = metadataById.get(providerPrefix + "/" + normalizedId);
            if (direct != null) {
                return Optional.of(direct);
            }
        }

        String alias = shortId(normalizedId);
        return Optional.ofNullable(metadataByAlias.get(alias));
    }

    /**
     * 读取静态 JSON 并建立完整 ID 与唯一短 ID 索引。
     *
     * @param objectMapper JSON 解析器
     * @param fullIdIndex 完整 ID 索引
     * @param aliasCandidates 短 ID 候选索引
     * @param ambiguousAliases 存在冲突的短 ID
     */
    private void loadCatalog(ObjectMapper objectMapper,
                             Map<String, ModelCatalogMetadata> fullIdIndex,
                             Map<String, ModelCatalogMetadata> aliasCandidates,
                             Set<String> ambiguousAliases) {
        ClassPathResource resource = new ClassPathResource(CATALOG_RESOURCE);
        try (InputStream inputStream = resource.getInputStream()) {
            JsonNode root = objectMapper.readTree(inputStream);
            if (root == null || !root.isObject()) {
                throw new IllegalStateException("模型能力库根节点必须是 JSON 对象");
            }
            Iterator<Map.Entry<String, JsonNode>> fields = root.fields();
            while (fields.hasNext()) {
                Map.Entry<String, JsonNode> field = fields.next();
                String normalizedId = normalize(field.getKey());
                if (normalizedId.isEmpty() || !field.getValue().isObject()) {
                    continue;
                }
                ModelCatalogMetadata metadata = toMetadata(normalizedId, field.getValue());
                fullIdIndex.put(normalizedId, metadata);
                registerAlias(shortId(normalizedId), metadata, aliasCandidates, ambiguousAliases);
            }
        } catch (IOException exception) {
            throw new IllegalStateException("无法加载模型能力库 " + CATALOG_RESOURCE, exception);
        }
    }

    /**
     * 将目录条目转换为展示元数据与能力信息。
     *
     * @param normalizedId 规范化模型 ID
     * @param node 模型目录条目
     * @return 模型目录元数据
     */
    private ModelCatalogMetadata toMetadata(String normalizedId, JsonNode node) {
        return new ModelCatalogMetadata(
                normalizedId,
                textValue(node, "name"),
                textValue(node, "family"),
                modalities(node, "input"),
                modalities(node, "output"),
                toCapability(normalizedId, node));
    }

    /**
     * 将模型目录条目转换为 EasyFlow 能力结果。
     *
     * @param normalizedId 规范化模型 ID
     * @param node 模型目录条目
     * @return EasyFlow 能力结果
     */
    private ModelCapabilityResolution toCapability(String normalizedId, JsonNode node) {
        String modelType = resolveModelType(normalizedId);
        if (!Model.MODEL_TYPES[0].equals(modelType)) {
            return new ModelCapabilityResolution(
                    modelType,
                    Boolean.FALSE,
                    Boolean.FALSE,
                    Boolean.FALSE,
                    ModelCapabilitySource.CATALOG);
        }
        return new ModelCapabilityResolution(
                modelType,
                hasInputModality(node, "image"),
                booleanValue(node, "reasoning"),
                booleanValue(node, "tool_call"),
                ModelCapabilitySource.CATALOG);
    }

    /**
     * 根据模型目录 ID 识别互斥模型类型。
     *
     * @param normalizedId 规范化模型 ID
     * @return EasyFlow 模型类型
     */
    private String resolveModelType(String normalizedId) {
        if (ModelCapabilityNameRules.isRerankModel(normalizedId)) {
            return Model.MODEL_TYPES[2];
        }
        if (ModelCapabilityNameRules.isEmbeddingModel(normalizedId)) {
            return Model.MODEL_TYPES[1];
        }
        return Model.MODEL_TYPES[0];
    }

    /**
     * 读取布尔字段，缺失时按 false 处理。
     *
     * @param node 模型条目
     * @param fieldName 字段名
     * @return 布尔字段值
     */
    private boolean booleanValue(JsonNode node, String fieldName) {
        JsonNode value = node.get(fieldName);
        return value != null && value.asBoolean(false);
    }

    /**
     * 判断模型输入模态是否包含指定类型。
     *
     * @param node 模型条目
     * @param modality 输入模态
     * @return 包含指定模态返回 true
     */
    private boolean hasInputModality(JsonNode node, String modality) {
        JsonNode inputs = node.path("modalities").path("input");
        if (!inputs.isArray()) {
            return false;
        }
        for (JsonNode input : inputs) {
            if (modality.equalsIgnoreCase(input.asText())) {
                return true;
            }
        }
        return false;
    }

    /**
     * 读取非空文本字段。
     *
     * @param node 模型目录条目
     * @param fieldName 字段名
     * @return 去除首尾空白的文本，缺失时返回 null
     */
    private String textValue(JsonNode node, String fieldName) {
        JsonNode value = node.get(fieldName);
        if (value == null || !value.isTextual() || value.asText().isBlank()) {
            return null;
        }
        return value.asText().trim();
    }

    /**
     * 读取并规范化模型模态集合。
     *
     * @param node 模型目录条目
     * @param direction input 或 output
     * @return 小写模态集合
     */
    private Set<String> modalities(JsonNode node, String direction) {
        JsonNode values = node.path("modalities").path(direction);
        if (!values.isArray()) {
            return Set.of();
        }
        Set<String> modalities = new HashSet<>();
        values.forEach(value -> {
            if (value.isTextual() && !value.asText().isBlank()) {
                modalities.add(value.asText().trim().toLowerCase(Locale.ROOT));
            }
        });
        return Set.copyOf(modalities);
    }

    /**
     * 注册无冲突的模型短 ID。
     *
     * @param alias 模型短 ID
     * @param capability 模型能力
     * @param aliasCandidates 短 ID 候选索引
     * @param ambiguousAliases 冲突短 ID 集合
     */
    private void registerAlias(String alias,
                               ModelCatalogMetadata metadata,
                               Map<String, ModelCatalogMetadata> aliasCandidates,
                               Set<String> ambiguousAliases) {
        if (alias.isEmpty() || ambiguousAliases.contains(alias)) {
            return;
        }
        ModelCatalogMetadata previous = aliasCandidates.putIfAbsent(alias, metadata);
        if (previous != null) {
            aliasCandidates.remove(alias);
            ambiguousAliases.add(alias);
        }
    }

    /**
     * 将 EasyFlow 供应商类型映射为 models.dev 前缀。
     *
     * @param providerType EasyFlow 供应商类型
     * @return models.dev 供应商前缀，未知时返回空字符串
     */
    private String providerPrefix(String providerType) {
        return switch (normalize(providerType)) {
            case "dashscope", "bailian", "aliyun" -> "alibaba";
            case "gemini" -> "google";
            case "kimi" -> "moonshotai";
            case "zhipu" -> "zhipuai";
            case "minimax" -> "minimax";
            case "azure-openai", "azure_openai" -> "openai";
            case "openai", "anthropic", "deepseek", "google", "xai", "mistral", "cohere" ->
                    normalize(providerType);
            default -> "";
        };
    }

    /**
     * 规范化模型或供应商标识。
     *
     * @param value 原始值
     * @return 小写且去除首尾空白的标识
     */
    private String normalize(String value) {
        return value == null ? "" : value.trim().toLowerCase(Locale.ROOT);
    }

    /**
     * 提取最后一个路径段作为模型短 ID。
     *
     * @param normalizedId 规范化模型 ID
     * @return 模型短 ID
     */
    private String shortId(String normalizedId) {
        int separator = normalizedId.lastIndexOf('/');
        return separator < 0 ? normalizedId : normalizedId.substring(separator + 1);
    }
}
