AIyaodianzhushou/backend-java/src/main/java/com/pharmacopoeia/service/RetrieverService.java

package com.pharmacopoeia.service;

import org.springframework.dao.DataAccessException;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Service;

import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.time.Duration;
import java.util.*;
import java.util.stream.Collectors;

@Service
public class RetrieverService {

    private final JdbcTemplate jdbc;
    private final LLMService llmService;
    private final StringRedisTemplate redis;
    private final boolean redisAvailable;

    // Embedding 缓存 key 前缀 + TTL（7 天，续期 3 次后不再续）
    private static final String EMBED_CACHE_PREFIX = "aiyaodian:embed:";
    private static final String EMBED_RENEW_PREFIX = "aiyaodian:embed:renew:";
    private static final Duration EMBED_CACHE_TTL = Duration.ofDays(7);
    private static final Duration EMBED_MAX_TTL = Duration.ofDays(30);
    private static final int MAX_RENEWALS = 3;

    private static final List<String> NEGATION_PATTERNS = List.of(
            "不是", "没有", "并非", "算不上", "怎么会是", "不可能", "不会",
            "没得", "没", "无", "不属", "不属于", "不是什么", "这不是"
    );

    public RetrieverService(JdbcTemplate jdbc, LLMService llmService,
                            StringRedisTemplate redis) {
        this.jdbc = jdbc;
        this.llmService = llmService;
        this.redis = redis;
        // 探测 Redis 是否可用，不可用时降级为每次调用远程 API
        this.redisAvailable = pingRedis();
    }

    /** 启动时探测 Redis，失败则降级为无缓存模式，避免阻塞主流程 */
    private boolean pingRedis() {
        try {
            return "PONG".equals(redis.getConnectionFactory().getConnection().ping());
        } catch (Exception e) {
            return false;
        }
    }

    /** 计算 query 的 SHA-256 作为缓存 key */
    private String cacheKey(String text) {
        try {
            MessageDigest md = MessageDigest.getInstance("SHA-256");
            byte[] hash = md.digest(text.getBytes(StandardCharsets.UTF_8));
            StringBuilder sb = new StringBuilder(2 * hash.length);
            for (byte b : hash) sb.append(String.format("%02x", b));
            return EMBED_CACHE_PREFIX + sb;
        } catch (Exception e) {
            return EMBED_CACHE_PREFIX + text.hashCode();
        }
    }

    /** 将 List<Float> 序列化为逗号分隔字符串，便于存入 Redis */
    private String encodeVec(List<Float> vec) {
        StringBuilder sb = new StringBuilder(vec.size() * 8);
        for (int i = 0; i < vec.size(); i++) {
            if (i > 0) sb.append(',');
            sb.append(vec.get(i));
        }
        return sb.toString();
    }

    /** 反序列化 */
    private List<Float> decodeVec(String s) {
        if (s == null || s.isEmpty()) return List.of();
        String[] parts = s.split(",");
        List<Float> vec = new ArrayList<>(parts.length);
        for (String p : parts) vec.add(Float.parseFloat(p));
        return vec;
    }

    /** 带缓存的 Embedding 查询：相同 query 命中缓存则跳过远程 API 调用 */
    private List<Float> embedWithCache(String query) {
        if (!redisAvailable) return llmService.embed(query);

        String key = cacheKey(query);
        try {
            String cached = redis.opsForValue().get(key);
            if (cached != null) {
                // 滑动续期：最多续 3 次，总时长不超 30 天
                var renewKey = EMBED_RENEW_PREFIX + key.substring(EMBED_CACHE_PREFIX.length());
                try {
                    var ops = redis.opsForValue();
                    var cntVal = ops.get(renewKey);
                    int renewals = cntVal != null ? Integer.parseInt(cntVal) : 0;
                    if (renewals < MAX_RENEWALS) {
                        redis.expire(key, EMBED_CACHE_TTL);
                        ops.set(renewKey, String.valueOf(renewals + 1), EMBED_MAX_TTL);
                    }
                } catch (Exception ignored) {}
                return decodeVec(cached);
            }
        } catch (DataAccessException ignored) {
            // Redis 异常时降级到远程调用
        }

        List<Float> vec = llmService.embed(query);
        try {
            redis.opsForValue().set(key, encodeVec(vec), EMBED_CACHE_TTL);
        } catch (DataAccessException ignored) {
            // 缓存写入失败不影响主流程
        }
        return vec;
    }

    /** 检测 query 中关键词之前是否存在否定词，避免误匹配 */
    private boolean hasNegation(String text, String keyword) {
        int idx = text.indexOf(keyword);
        if (idx < 0) return false;
        String prefix = text.substring(0, idx);
        for (String neg : NEGATION_PATTERNS) {
            if (prefix.endsWith(neg) || prefix.contains(neg)) return true;
        }
        return false;
    }

    /** 已知药物过敏检测：用户明确询问对某药物的过敏处理方法 */
    private static final List<String> KNOWN_ALLERGY_DRUGS = List.of(
        "青霉素", "头孢", "磺胺", "阿莫西林", "布洛芬",
        "阿司匹林", "链霉素", "庆大霉素", "四环素",
        "红霉素", "氯霉素", "万古霉素", "喹诺酮",
        "普鲁卡因", "利多卡因", "碘", "破伤风"
    );

    private static final java.util.regex.Pattern[] KNOWN_ALLERGY_REGEX = {
        java.util.regex.Pattern.compile("对.{1,6}过敏"),      // 对青霉素过敏
        java.util.regex.Pattern.compile(".{1,6}药过敏"),      // 头孢药过敏
        java.util.regex.Pattern.compile("药物过敏"),           // 药物过敏
        java.util.regex.Pattern.compile(".+过敏(处理|怎么办|急救|替代|注意|救治|抢救|应急)"),
    };

    private boolean isKnownDrugAllergy(String query) {
        // 必须包含"过敏"
        if (!query.contains("过敏")) return false;

        // 精确匹配已知过敏药物清单
        for (String drug : KNOWN_ALLERGY_DRUGS) {
            if (query.contains(drug)) return true;
        }

        // 正则模式匹配（对XX过敏 / XX药过敏 / 药物过敏 / 过敏处理/急救等）
        for (var pattern : KNOWN_ALLERGY_REGEX) {
            if (pattern.matcher(query).find()) return true;
        }

        return false;
    }

    public String classifyIntent(String query) {
        String q = query.trim();

        // 0. 已知药物过敏检测（最高优先级，必须在 "过敏" 通用匹配之前）
        //    如 "青霉素过敏处理" "头孢过敏怎么办" → 用药安全咨询，非症状求药
        if (isKnownDrugAllergy(q)) {
            return "usage_guide";
        }

        // 1. 用药安全/用法用量（优先级最高，避免"过敏"与症状类冲突）
        if (anyMatch(q, "怎么吃", "吃多少", "怎么用", "怎么服用", "孕妇", "儿童用量",
                "副作用多大", "伤肝", "伤肾", "安全吗", "副作用", "不良反应",
                "禁忌", "过敏", "能不能", "可以吗", "用法", "用量", "剂量",
                "用药指导", "一天几次", "一次多少", "饭前", "饭后", "空腹",
                "能不能一起吃", "相互作用", "过量", "停用", "停药", "忌口",
                "饮酒", "肝功能", "肾功能")) {
            return "usage_guide";
        }

        // 2. 考试辅导（高优先级，关键词明确）
        if (anyMatch(q, "执业药师", "考点", "历年真题", "考试大纲", "高频考点",
                "药物化学", "药剂学", "药理学", "药分", "药物分析")) {
            return "exam_tutor";
        }

        // 3. 法规条款（关键词明确）
        if (anyMatch(q, "凡例", "通则规定", "制剂通则", "一般规定", "通则")) {
            return "regulation";
        }

        // 4. 症状用药建议（安全类关键词已在上方处理，"过敏"不会落到这里）
        if (anyMatch(q, "发烧", "咳嗽", "感冒", "腹泻", "头疼", "头痛", "嗓子疼",
                "吃了什么药", "吃什么药", "该吃", "推荐用药", "推荐下用药",
                "体温", "多少度", "退烧", "止痛", "止泻", "鼻塞", "流鼻涕",
                "头晕", "乏力", "呕吐", "腹痛", "咽痛", "打喷嚏")) {
            return "symptom_advice";
        }

        // 5. 兜底：药品查询
        return "drug_query";
    }

    /** 检索：仅 2025 年版药典，AI 兜底补充 */
    public List<Map<String, Object>> search(String query, String intent, int topK) {
        List<Float> vec = embedWithCache(query);
        if (vec.isEmpty()) return List.of();
        String vecStr = vec.stream()
                .map(String::valueOf)
                .collect(Collectors.joining(",", "[", "]"));

        String drugName = extractDrugName(query);
        return searchWithVersion(vecStr, drugName, "2025年版", topK);
    }

    /** 响应式检索：供 SSE 端点使用，避免 reactor 线程中 block() */
    public reactor.core.publisher.Mono<List<Map<String, Object>>> searchReactive(String query, String intent, int topK) {
        return llmService.embedReactive(query)
                .map(vec -> {
                    if (vec.isEmpty()) return List.<Map<String, Object>>of();
                    String vecStr = vec.stream()
                            .map(String::valueOf)
                            .collect(Collectors.joining(",", "[", "]"));
                    String drugName = extractDrugName(query);
                    return searchWithVersion(vecStr, drugName, "2025年版", topK);
                });
    }

    /** 带版本过滤的向量检索。version 为 null 时不限制版本。 */
    private List<Map<String, Object>> searchWithVersion(
            String vecStr, String drugName, String version, int topK) {

        String versionFilter = (version != null)
                ? " AND d.source_version = '" + version + "' "
                : " ";

        String sql;
        Object[] params;
        if (!drugName.isEmpty()) {
            // 精确匹配
            sql = """
                SELECT c.content, c.source, c.drug_id, c.section,
                       d.name, d.category, d.source_version, d.source_volume,
                       1 - (c.vec <=> ?::vector) AS similarity
                FROM drug_chunks c
                JOIN drugs d ON d.drug_id = c.drug_id
                WHERE c.vec IS NOT NULL AND d.name = ?""" + versionFilter + """
                ORDER BY c.vec <=> ?::vector
                LIMIT ?
                """;
            params = new Object[]{vecStr, drugName, vecStr, topK};
            List<Map<String, Object>> results = jdbc.queryForList(sql, params);
            if (!results.isEmpty()) return results;

            // 前缀匹配兜底
            sql = """
                SELECT c.content, c.source, c.drug_id, c.section,
                       d.name, d.category, d.source_version, d.source_volume,
                       1 - (c.vec <=> ?::vector) AS similarity
                FROM drug_chunks c
                JOIN drugs d ON d.drug_id = c.drug_id
                WHERE c.vec IS NOT NULL AND d.name LIKE ?""" + versionFilter + """
                ORDER BY c.vec <=> ?::vector
                LIMIT ?
                """;
            params = new Object[]{vecStr, drugName + "%", vecStr, topK};
        } else {
            sql = """
                SELECT c.content, c.source, c.drug_id, c.section,
                       d.name, d.category, d.source_version, d.source_volume,
                       1 - (c.vec <=> ?::vector) AS similarity
                FROM drug_chunks c
                JOIN drugs d ON d.drug_id = c.drug_id
                WHERE c.vec IS NOT NULL""" + versionFilter + """
                ORDER BY c.vec <=> ?::vector
                LIMIT ?
                """;
            params = new Object[]{vecStr, vecStr, topK};
        }

        return jdbc.queryForList(sql, params);
    }

    // 剂型后缀（长后缀优先，避免"缓释胶囊"被错误截断为"缓释"）
    private static final List<String> FORMULATION_SUFFIXES = List.of(
        "缓释胶囊", "缓释片", "肠溶胶囊", "肠溶片", "分散片", "咀嚼片",
        "口服混悬液", "口服液", "混悬液", "滴眼液", "注射液",
        "缓释", "肠溶", "胶囊", "颗粒", "糖浆", "软膏", "栓剂",
        "片", "剂", "栓"
    );

    /** 从 query 中提取已知药品名：查 drugs 表，支持剂型后缀剥离和模糊匹配 */
    private String extractDrugName(String query) {
        String cleaned = query.trim();

        // 第一步：去掉尾部常见修饰词
        String[] querySuffixes = {
            "的用法与用量", "的用法用量", "用法与用量", "用法用量", "的用量", "的用法",
            "的副作用", "不良反应", "的禁忌", "禁忌", "的注意事项", "注意事项",
            "是什么", "说明书", "怎么用", "怎么吃", "的用量", "用量", "的剂量", "剂量"
        };
        for (String s : querySuffixes) {
            if (cleaned.endsWith(s)) {
                cleaned = cleaned.substring(0, cleaned.length() - s.length()).trim();
                break;
            }
        }
        // 去掉问句前缀
        cleaned = cleaned.replaceAll("^(什么是|怎么|如何|告诉我|请问|查询|搜索|查一下)", "").trim();

        if (cleaned.length() < 2) return "";

        // 第二步：精确匹配原始 query（含剂型名如"布洛芬缓释胶囊"）
        String exact = tryExactMatch(cleaned);
        if (!exact.isEmpty()) return exact;

        // 第三步：逐步剥剂型后缀再试（"布洛芬缓释胶囊"→"布洛芬"）
        for (String suffix : FORMULATION_SUFFIXES) {
            if (cleaned.endsWith(suffix)) {
                String base = cleaned.substring(0, cleaned.length() - suffix.length()).trim();
                if (base.length() >= 2) {
                    String match = tryExactMatch(base);
                    if (!match.isEmpty()) return match;
                }
            }
        }

        // 第四步：ILIKE 模糊匹配兜底
        return tryFuzzyMatch(cleaned);
    }

    private String tryExactMatch(String name) {
        try {
            List<String> matches = jdbc.queryForList(
                "SELECT name FROM drugs WHERE name = ? AND is_active = TRUE LIMIT 1",
                String.class, name);
            if (!matches.isEmpty()) return matches.get(0);
        } catch (Exception ignored) {}
        return "";
    }

    private String tryFuzzyMatch(String name) {
        try {
            List<String> matches = jdbc.queryForList(
                "SELECT name FROM drugs WHERE name ILIKE ? AND is_active = TRUE ORDER BY name LIMIT 1",
                String.class, "%" + name + "%");
            if (!matches.isEmpty()) return matches.get(0);
        } catch (Exception ignored) {}
        return "";
    }

    private boolean anyMatch(String text, String... keywords) {
        for (String kw : keywords) {
            int idx = text.indexOf(kw);
            if (idx >= 0 && !hasNegation(text, kw)) {
                return true;
            }
        }
        return false;
    }
}