package com.aicodeassistant.llm;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Service;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 模型注册表 — 内置模型映射表 + Provider 动态查询 + 用户自定义覆盖。
 */
@Service
public class ModelRegistry {

    private static final Logger log = LoggerFactory.getLogger(ModelRegistry.class);

    private final Map<String, ModelCapabilities> customModels = new ConcurrentHashMap<>();
    private final Set<String> warnedUnknownModels = ConcurrentHashMap.newKeySet();
    private final LlmProviderRegistry providerRegistry;
    private final ModelCapabilitiesProperties properties;

    // 内置模型映射表
    private static final Map<String, ModelCapabilities> BUILTIN_MODELS = Map.ofEntries(
        // OpenAI
        // https://developers.openai.com/api/docs/models/gpt-6-astra (2026-11-02).
        // 4 张图片为应用保守上限；费用显示为基础档 per-1k，不含长上下文加价。
        entry("gpt-6-astra",       cacheCaps("gpt-6-astra", "GPT-6 Astra（OpenAI）", 128000, 1050000, true, true, true, 4, true, 0.010, 0.050, ModelCapabilities.ImageInputMode.URL)),
        entry("gpt-5.6-sol",       caps("gpt-5.6-sol",       "GPT-5.6-Sol",      128000, 1050000, true, true, true, 10, true, 0.005, 0.03)),
        entry("gpt-5.4-mini",      caps("gpt-5.4-mini",      "GPT-5.4 Mini",     128000, 400000,  true, true, true, 10, true, 0.00075, 0.0045)),
        // Anthropic
        entry("claude-sonnet-4-6", caps("claude-sonnet-4-6", "Claude Sonnet 4.6",  16384, 200000, true, true, true, 10, true, 0.003, 0.015)),
        entry("claude-haiku-4-5", caps("claude-haiku-4-5", "Claude Haiku 4.5", 8192, 200000, true, false, true, 10, true, 0.0008, 0.004)),
         // Anthropic via ZenMux (anthropic/ 前缀 = zenmux 中转，1M ctx · 128K 最大输出，我们保守设 64K)
        // 价格为网关 per-M 定价换算 per-1k（2026-10-12 核对）：Fable $10/$50、Opus 5.5 $4/$20、Sonnet 5.5 $2/$10、Haiku 5.5 $0.1/$0.5
        // ZenMux Claude 上游可能拒绝 URL 图片源；由现有 UserImageTranscoder 在请求副本中转为 Base64。
        entry("anthropic/claude-fable-5.1", cacheCaps("anthropic/claude-fable-5.1", "Claude Fable 5.1（ZenMux）", 64000, 1000000, true, false, true, 5, true, 0.010, 0.050, ModelCapabilities.ImageInputMode.BASE64_ONLY)),
        entry("anthropic/claude-opus-5.5",   cacheCaps("anthropic/claude-opus-5.5",   "Claude Opus 5.5（ZenMux）",   64000, 1000000, true, false, true, 5, true, 0.004, 0.020, ModelCapabilities.ImageInputMode.BASE64_ONLY)),
        entry("anthropic/claude-sonnet-5.5", cacheCaps("anthropic/claude-sonnet-5.5", "Claude Sonnet 5.5（ZenMux）", 64000, 1000000, true, false, true, 5, true, 0.002, 0.010, ModelCapabilities.ImageInputMode.BASE64_ONLY)),
        entry("anthropic/claude-haiku-5.5",  cacheCaps("anthropic/claude-haiku-5.5",  "Claude Haiku 5.5（ZenMux）",  64000, 1000000, true, false, true, 5, true, 0.0001, 0.0005, ModelCapabilities.ImageInputMode.BASE64_ONLY)),
        // OpenAI via ZenMux (openai/ 前缀 = zenmux 中转)
        entry("openai/gpt-5.6-sol",   caps("openai/gpt-5.6-sol",   "OpenAI GPT-5.6 Sol",   128000, 1050000, true,  true, true, 4, true, 0.030, 0.180)),
        entry("openai/gpt-6-astra",   caps("openai/gpt-6-astra",   "GPT-6 Astra（ZenMux）",   128000, 1050000, true,  true, true, 4, true, 0.010, 0.050)),
        // Google via ZenMux (google/ 前缀 = zenmux 中转)
        // Gemini 4 Argon：2026-10-17 核对 ZenMux 目录与 Responses 探测（1M 上下文 / 65536 输出 / 图片 + 工具 + Responses 推理）。
        entry("google/gemini-4-argon",     caps("google/gemini-4-argon",     "Gemini 4 Argon（ZenMux）",   65536, 1048576, true, true, true, 4, true, 0.002, 0.010)),
        // xAI via ZenMux
        // Grok 4.7：2026-09-21 核对 ZenMux 目录与 Responses 探测；价格取基础档 $2/$6 per M 换算 per-1k，≥200K 提示按 $4/$12 计且界面估算不按阶梯（同 GPT-6 Astra）。
        entry("x-ai/grok-4.7", caps("x-ai/grok-4.7", "xAI Grok 4.7（ZenMux）", 65536, 500000, true, true, true, 4, true, 0.002, 0.006)),
        // Meta via ZenMux (meta/ 前缀 = zenmux 中转)
        // Muse Wonder 1.0：2026-10-17 核对 ZenMux 目录与 Responses 探测（1M 上下文 / 128000 输出 / 图片 + 工具 + Responses 推理）。
        entry("meta/muse-wonder-1.0",     caps("meta/muse-wonder-1.0",     "Muse Wonder 1.0（ZenMux）",   128000, 1048576, true, true, true, 4, true, 0.004, 0.020)),
        entry(OpenRouterModels.FABLE, OpenRouterModels.capabilities(OpenRouterModels.FABLE)),
        // 国产大模型
        // DeepSeek V4.1 Flash 官方规格：1M 上下文 / 384K 最大输出 / 原生视觉 / 思考与工具调用。
        // 模型规格：https://api-docs.deepseek.com/quick_start/pricing/
        // 视觉限制：https://api-docs.deepseek.com/guides/vision/（单个请求最多 600 张图片）
        entry("deepseek-flash", caps("deepseek-flash", "DeepSeek V4.1 Flash", 384000, 1000000, true, true, true, 600, true, 0.0003, 0.0012)),
        // 百炼官方规格：1M context / 393216 output / vision + tools.
        // https://help.aliyun.com/zh/model-studio/deepseek-v4-1-flash
        // 图片保守限制为 4 张（已实测）；价格为北京忙时按量参考，非 Token Plan Credits 账单。
        entry("deepseek-v4.1-flash", caps("deepseek-v4.1-flash", "DeepSeek V4.1 Flash（百炼）", 393216, 1000000, true, true, true, 4, true, 0.002, 0.008)),
        // Moonshot K3.1 Pro Highspeed：1,048,576 上下文；登记输出预算 131072，主循环统一裁剪到 65536。
        // 官方协议：https://platform.kimi.com/docs/guide/kimi-k3-1-pro-quickstart
        // 美元估算按海外官方 $4/$20 per MTok（2026-10-24）：https://platform.kimi.ai/docs/pricing/chat
        // 国内账单为 ¥26/¥136 per MTok；此处不含缓存写入/命中价格，也不做实时汇率换算。
        entry("kimi-k3.1-pro-highspeed", caps("kimi-k3.1-pro-highspeed", "Kimi K3.1 Pro Highspeed", 131072, 1048576, true, true, true, 8, true, 0.004, 0.020, ModelCapabilities.ImageInputMode.BASE64_ONLY)),
        // Kimi Code Highspeed 订阅（2026-10-24）：需 Pro / Allegretto 及以上套餐，最高 1M 上下文。
        // https://www.kimi.com/code/docs/kimi-code/models.html
        // 登记输出预算 131072，主循环统一裁剪到 65536；8 张图片为保守上限，订阅不按量估价。
        entry("k3.1-pro-highspeed", new ModelCapabilities("k3.1-pro-highspeed", "Kimi K3.1 Pro Highspeed（订阅）", 131072, 1048576,
                true, true, true, 8, true, 0.0, 0.0, 3.5, true, ModelCapabilities.ImageInputMode.BASE64_ONLY)),
        entry("kimi-for-coding", new ModelCapabilities("kimi-for-coding", "Kimi K2.8 Preview（订阅）", 131072, 1048576,
                true, true, true, 8, true, 0.0, 0.0, 3.5, true, ModelCapabilities.ImageInputMode.BASE64_ONLY)),
        entry("qwen-turbo",        caps("qwen-turbo",        "Qwen Turbo",         8192, 1000000,  true, false, false, 0, true, 0.0003, 0.0006)),
        entry("qwen3.8-max-0902", caps("qwen3.8-max-0902", "Qwen 3.8 Max 0902", 65536, 1000000, true, true, true, 4, true, 0.009, 0.054)),
        entry("qwen3.8-max", caps("qwen3.8-max", "Qwen 3.8 Max（百炼）", 65536, 1000000, true, true, true, 4, true, 0.009, 0.054)),
        // Qwen3.8-Flash 官方规格（help.aliyun.com/zh/model-studio/qwen3-8-flash）：
        // 1M 上下文 / 131072 最大输出 / 多模态（Image/Text/Video 输入）/ 支持思考模式（最大思维链 262144）
        // 北京区定价：输入 0.8 元/M、输出 2.7 元/M tokens
        entry("qwen3.8-flash", caps("qwen3.8-flash", "Qwen 3.8 Flash（百炼）", 131072, 1000000, true, true, true, 4, true, 0.0008, 0.0027)),
        entry("glm-5.3",           caps("glm-5.3",           "GLM-5.3",            131072, 1048576,  true, true, false, 0, true, 0.001, 0.001)),
        entry(BailianTokenPlanModels.GLM_53, BailianTokenPlanModels.capabilities(BailianTokenPlanModels.GLM_53)),
        // GLM-5.3-Flash 官方规格（docs.bigmodel.cn/cn/guide/models/vlm/glm-5.3-flash）：
        // 1M 上下文 / 128K 最大输出 / 原生多模态（单图≤5MB，系列上限 50 张）/ 强制思考不可关闭
        // 定价为 GLM-5.3 的 1/10（输入 $0.15/M、输出 $0.50/M）
        entry("glm-5.3-flash",     caps("glm-5.3-flash",     "GLM-5.3-Flash",      131072, 1048576,  true, true, true, 50, true, 0.00015, 0.0005)),
        // MiniMax
        entry("MiniMax-M3",       caps("MiniMax-M3",       "MiniMax M3",        16384, 1000000, true, true, true, 4, true, 0.001, 0.004)),
        // Ollama 本地
        entry("ollama/*",          caps("ollama/*",          "Ollama Local",       4096,   8192,  true, false, false, 0, false, 0.0, 0.0))
    );

    public ModelRegistry(LlmProviderRegistry providerRegistry) {
        this(providerRegistry, new ModelCapabilitiesProperties());
    }

    @org.springframework.beans.factory.annotation.Autowired
    public ModelRegistry(LlmProviderRegistry providerRegistry, ModelCapabilitiesProperties properties) {
        this.providerRegistry = providerRegistry;
        this.properties = properties;
        applyConfiguredOverrides();
        log.info("ModelRegistry initialized with {} built-in models", BUILTIN_MODELS.size());
    }

    /** 获取模型的 contextWindow — 核心 API */
    public int getContextWindowForModel(String modelId) {
        return getCapabilities(modelId).contextWindow();
    }

    /** 获取模型的 maxOutputTokens */
    public int getMaxOutputTokensForModel(String modelId) {
        return getCapabilities(modelId).maxOutputTokens();
    }

    public double getTokenCharRatio(String modelId) { return getCapabilities(modelId).tokenCharRatio(); }

    /**
     * 获取完整模型能力 — 三级查询优先级:
     * 1. 用户自定义覆盖（YAML 配置）
     * 2. Provider 动态查询（已注册的 LlmProvider）
     * 3. 内置映射表（BUILTIN_MODELS）
     * 4. 保守默认值（ModelCapabilities.DEFAULT）
     */
    public ModelCapabilities getCapabilities(String modelId) {
        if (modelId == null || modelId.isBlank()) {
            if (warnedUnknownModels.add("<blank>")) {
                log.warn("Blank model id; using conservative DEFAULT capabilities");
            }
            return ModelCapabilities.DEFAULT;
        }
        // Level 1: 用户自定义覆盖
        ModelCapabilities custom = customModels.get(modelId);
        if (custom != null) return custom;

        // Level 2: 从 Provider 动态查询
        // 注意: getProvider() 在无匹配时抛 IllegalArgumentException，不返回 null
        try {
            ModelCapabilities providerCapabilities =
                    providerRegistry.getProvider(modelId).getModelCapabilities(modelId);
            if (providerCapabilities != null) {
                return providerCapabilities;
            }
            // Provider 属于扩展边界，错误实现可能违反非空契约；注册表必须继续
            // 回退，确保所有调用方始终获得完整的能力对象。
            log.warn("Provider returned null capabilities for model '{}'; falling back to built-in", modelId);
        } catch (Exception e) {
            log.trace("No provider for model {}, falling back to built-in", modelId);
        }

        // Level 3: 内置映射表（精确匹配 + 前缀匹配）
        ModelCapabilities builtin = BUILTIN_MODELS.get(modelId);
        if (builtin != null) return builtin;

        // 前缀匹配 (ollama/* 等通配符)
        for (var e : BUILTIN_MODELS.entrySet()) {
            String key = e.getKey();
            if (key.endsWith("/*") && modelId.startsWith(key.substring(0, key.length() - 1))) {
                return e.getValue();
            }
        }

        // Level 4: 保守默认值
        if (warnedUnknownModels.add(String.valueOf(modelId))) {
            log.warn("Unknown model '{}'; using conservative DEFAULT capabilities", modelId);
        }
        return ModelCapabilities.DEFAULT;
    }

    /** Summary-specific resolution: never consult another provider or DEFAULT capabilities. */
    public java.util.Optional<ModelCapabilities> findExplicitCapabilities(String modelId, LlmProvider selected) {
        ModelCapabilities base = null;
        try {
            if (selected != null && selected.getSupportedModels().contains(modelId)) {
                var candidate = selected.getModelCapabilities(modelId);
                if (candidate != null && modelId.equals(candidate.modelId())) base = candidate;
            }
        } catch (RuntimeException ignored) { }
        if (base == null) base = BUILTIN_MODELS.get(modelId);
        var override = properties.getCapabilities().get(modelId);
        if (override == null) return java.util.Optional.ofNullable(base);
        if (base == null && (override.getContextWindow() == null || override.getOutputMaxTokens() == null
                || override.getTokenCharRatio() == null)) return java.util.Optional.empty();
        try {
            return java.util.Optional.of(new ModelCapabilities(modelId, modelId,
                    override.getOutputMaxTokens() != null ? override.getOutputMaxTokens() : base.maxOutputTokens(),
                    override.getContextWindow() != null ? override.getContextWindow() : base.contextWindow(),
                    true, true, false, 0, false, 0, 0,
                    override.getTokenCharRatio() != null ? override.getTokenCharRatio() : base.tokenCharRatio()));
        } catch (IllegalArgumentException e) { return java.util.Optional.empty(); }
    }

    public boolean isKnownModel(String modelId) {
        if (modelId == null) return false;
        if (customModels.containsKey(modelId) || BUILTIN_MODELS.containsKey(modelId)) return true;
        if (BUILTIN_MODELS.keySet().stream().anyMatch(k -> k.endsWith("/*")
                && modelId.startsWith(k.substring(0, k.length() - 1)))) return true;
        try { return providerRegistry.getProvider(modelId).getModelCapabilities(modelId) != null; }
        catch (Exception ignored) { return false; }
    }

    /** 注册自定义模型能力（YAML 配置加载时调用） */
    public void registerCustomModel(ModelCapabilities capabilities) {
        customModels.put(capabilities.modelId(), capabilities);
        log.info("Custom model registered: {}", capabilities.modelId());
    }

    /** 列出所有已知模型（内置 + 自定义 + Provider） */
    public List<ModelCapabilities> listAllModels() {
        Map<String, ModelCapabilities> all = new LinkedHashMap<>(BUILTIN_MODELS);
        all.putAll(customModels);
        // 追加 Provider 注册的模型
        try {
            providerRegistry.listModelCapabilities().forEach(mc -> all.putIfAbsent(mc.modelId(), mc));
        } catch (Exception ignored) {}
        return List.copyOf(all.values());
    }

    // ── 内部工具方法 ──

    private static Map.Entry<String, ModelCapabilities> entry(String k, ModelCapabilities v) {
        return Map.entry(k, v);
    }

    private static ModelCapabilities caps(String id, String name, int maxOut, int ctx,
            boolean stream, boolean think, boolean img, int maxImages, boolean tool, double in$, double out$) {
        return new ModelCapabilities(id, name, maxOut, ctx, stream, think, img, maxImages, tool, in$, out$, 3.5);
    }

    private static ModelCapabilities caps(String id, String name, int maxOut, int ctx,
            boolean stream, boolean think, boolean img, int maxImages, boolean tool,
            double in$, double out$, ModelCapabilities.ImageInputMode imageInputMode) {
        return new ModelCapabilities(id, name, maxOut, ctx, stream, think, img,
                maxImages, tool, in$, out$, 3.5, false, imageInputMode);
    }

    private static ModelCapabilities cacheCaps(String id, String name, int maxOut, int ctx,
            boolean stream, boolean think, boolean img, int maxImages, boolean tool,
            double in$, double out$, ModelCapabilities.ImageInputMode imageInputMode) {
        return new ModelCapabilities(id, name, maxOut, ctx, stream, think, img,
                maxImages, tool, in$, out$, 3.5, true, imageInputMode);
    }

    private void applyConfiguredOverrides() {
        for (var entry : properties.getCapabilities().entrySet()) {
            String id = entry.getKey();
            ModelCapabilitiesProperties.Override override = entry.getValue();
            if (override.getTokenCharRatio() == null)
                throw new IllegalStateException("model.capabilities." + id + ".tokenCharRatio is required");
            ModelCapabilities base = resolveBase(id);
            int context = override.getContextWindow() == null ? base.contextWindow() : override.getContextWindow();
            int output = override.getOutputMaxTokens() == null ? base.maxOutputTokens() : override.getOutputMaxTokens();
            customModels.put(id, new ModelCapabilities(id,
                    "unknown".equals(base.modelId()) ? id : base.displayName(), output, context,
                    value(override.getSupportsStreaming(), base.supportsStreaming()), base.supportsThinking(),
                    value(override.getSupportsVision(), base.supportsImages()), base.maxImages(),
                    value(override.getSupportsToolUse(), base.supportsToolUse()),
                    base.costPer1kInput(), base.costPer1kOutput(), override.getTokenCharRatio(),
                    value(override.getSupportsCache(), base.supportsCache()),
                    override.getImageInputMode() != null ? override.getImageInputMode() : base.imageInputMode()));
        }
    }

    private ModelCapabilities resolveBase(String modelId) {
        try {
            ModelCapabilities providerCapabilities =
                    providerRegistry.getProvider(modelId).getModelCapabilities(modelId);
            if (providerCapabilities != null) return providerCapabilities;
        }
        catch (Exception ignored) { }
        ModelCapabilities builtin = BUILTIN_MODELS.get(modelId);
        return builtin == null ? ModelCapabilities.DEFAULT : builtin;
    }

    private static boolean value(Boolean configured, boolean fallback) {
        return configured == null ? fallback : configured;
    }
}