# -*- coding: utf-8 -*-
"""
search_engine.py — 问题知识库搜索引擎

功能：
  基于关键词匹配 + TF-IDF 相似度，从问题知识库中检索最相关的历史案例。
  P2-1 新增：Embedding 向量搜索（优先），TF-IDF 降级兜底。

用法：
  from search_engine import SearchEngine

  engine = SearchEngine()
  results = engine.search("门口屏 MQTT 连接失败", top_k=5)
"""

import os
import re
import json
import math
from pathlib import Path
from collections import defaultdict, Counter, OrderedDict

try:
    import numpy as np
    HAS_NUMPY = True
except ImportError:
    HAS_NUMPY = False

from utils.logger import get_logger
from utils.paths import DATA_DIR, SEARCH_INDEX_DIR, VECTOR_INDEX_FILENAME

logger = get_logger(__name__)

# ============================================================
# 配置
# ============================================================

SCRIPT_DIR = Path(__file__).resolve().parent  # .../web

# 搜索索引路径（按优先级查找：统一路径 → 部署环境 → 开发环境）
SEARCH_INDEX_PATHS = [
    SEARCH_INDEX_DIR / "搜索索引.json",                              # 统一路径（容器/裸机）
    DATA_DIR / "搜索索引.json",                                      # 裸机部署（/opt/troubleshoot/）
    Path("E:/github/ubains-module-test/develop/Docs/PRD/问题知识库/搜索索引.json"),  # 开发环境绝对路径
    SCRIPT_DIR.parent.parent.parent.parent / "Docs" / "PRD" / "问题知识库" / "搜索索引.json",  # 相对路径
]

# 向量索引路径（与搜索索引同目录查找）
VECTOR_INDEX_PATHS = [
    SEARCH_INDEX_DIR / VECTOR_INDEX_FILENAME,                        # 统一路径（容器/裸机）
    DATA_DIR / VECTOR_INDEX_FILENAME,                                # 裸机部署
    SCRIPT_DIR.parent.parent.parent.parent / "Docs" / "PRD" / "问题知识库" / VECTOR_INDEX_FILENAME,  # 相对路径
]

# query 向量缓存上限
_QUERY_VECTOR_CACHE_MAX = 100


def find_search_index():
    """按优先级查找搜索索引文件"""
    for path in SEARCH_INDEX_PATHS:
        if path.exists():
            print(f"[搜索引擎] 找到索引文件：{path}")
            return path
    raise FileNotFoundError(f"搜索索引不存在，已查找路径：{SEARCH_INDEX_PATHS}")


def find_vector_index():
    """按优先级查找向量索引文件（可选，缺失不报错）"""
    for path in VECTOR_INDEX_PATHS:
        if path.exists():
            print(f"[搜索引擎] 找到向量文件：{path}")
            return path
    return None

# 停用词（过滤无意义的词）
STOPWORDS = {
    '的', '了', '是', '在', '有', '我', '不', '和', '与', '或',
    '这个', '那个', '怎么', '为什么', '哪里', '什么', '可以',
    '应该', '需要', '请', '帮', '一下', '吗', '呢', '啊', '吧',
    'the', 'a', 'an', 'is', 'are', 'was', 'were', 'to', 'of', 'in',
    'and', 'or', 'not', 'for', 'on', 'at', 'by', 'be', 'it', 'its',
    'as', 'if', 'then', 'than', 'that', 'this', 'these', 'those',
}


def load_search_index():
    """加载搜索索引"""
    index_file = find_search_index()
    with open(index_file, 'r', encoding='utf-8') as f:
        return json.load(f)


def tokenize(text):
    """中文分词（简单实现：按字切分 + 关键词提取）"""
    # 提取中文词组（2-4字）
    cn_words = re.findall(r'[一-龥]{2,4}', text)
    # 提取英文单词
    en_words = re.findall(r'[a-zA-Z]{2,}', text.lower())
    # 提取数字
    numbers = re.findall(r'\d+', text)

    tokens = cn_words + en_words + numbers
    # 过滤停用词
    tokens = [t for t in tokens if t not in STOPWORDS]

    return tokens


def compute_tf(tokens):
    """计算词频 (TF)"""
    tf = Counter(tokens)
    total = len(tokens) if tokens else 1
    return {word: count / total for word, count in tf.items()}


def compute_idf(documents):
    """计算逆文档频率 (IDF)"""
    N = len(documents)
    df = defaultdict(int)

    for doc in documents:
        unique_tokens = set(doc)
        for token in unique_tokens:
            df[token] += 1

    idf = {}
    for token, count in df.items():
        # 平滑处理：避免除零
        idf[token] = math.log((N + 1) / (count + 1)) + 1

    return idf


def compute_tfidf(tf, idf):
    """计算 TF-IDF 向量"""
    return {word: tf_val * idf.get(word, 1) for word, tf_val in tf.items()}


def cosine_similarity(vec1, vec2):
    """计算余弦相似度"""
    # 找到共同词
    common_words = set(vec1.keys()) & set(vec2.keys())
    if not common_words:
        return 0.0

    # 计算点积
    dot_product = sum(vec1[w] * vec2[w] for w in common_words)

    # 计算模长
    norm1 = math.sqrt(sum(v ** 2 for v in vec1.values()))
    norm2 = math.sqrt(sum(v ** 2 for v in vec2.values()))

    if norm1 == 0 or norm2 == 0:
        return 0.0

    return dot_product / (norm1 * norm2)


def _cosine_sim_python(vec1, vec2):
    """纯 Python 余弦相似度（用于向量列表，非 dict）"""
    if len(vec1) != len(vec2) or not vec1:
        return 0.0
    dot_product = sum(a * b for a, b in zip(vec1, vec2))
    norm1 = math.sqrt(sum(a * a for a in vec1))
    norm2 = math.sqrt(sum(b * b for b in vec2))
    if norm1 == 0 or norm2 == 0:
        return 0.0
    return dot_product / (norm1 * norm2)


class SearchEngine:
    """问题知识库搜索引擎（TF-IDF + Embedding 向量搜索）"""

    def __init__(self):
        """初始化：加载索引并构建 TF-IDF 模型，可选加载向量索引"""
        self.index = load_search_index()
        self.records = self.index.get('records', [])

        # 预处理：为每条记录构建 TF-IDF 向量
        self.doc_tokens = []
        self.doc_vectors = []
        self.idf = {}

        if self.records:
            # 对所有记录进行分词
            self.doc_tokens = [
                tokenize(r.get('full_text', '') + ' ' + r.get('title', ''))
                for r in self.records
            ]

            # 计算 IDF
            self.idf = compute_idf(self.doc_tokens)

            # 预计算每条记录的 TF-IDF 向量
            self.doc_vectors = [
                compute_tfidf(compute_tf(tokens), self.idf)
                for tokens in self.doc_tokens
            ]

        # P2-1：向量搜索相关字段
        self._vectors = None           # dict: {rec_key: [float]}
        self._vector_ids = []          # 有序 ID 列表（与矩阵行对齐）
        self._vector_matrix = None     # numpy array (N, D)，加速批量余弦
        self._embedding_enabled = False
        self._query_vector_cache = OrderedDict()  # LRU: query_text → vector
        self._load_vectors()

        print(f"[搜索引擎] 初始化完成：{len(self.records)} 条记录，搜索模式：{'vector' if self._embedding_enabled else 'tfidf'}")

    # ============================================================
    # P2-1：向量搜索内部方法
    # ============================================================

    def _load_vectors(self):
        """尝试加载向量文件，失败则静默降级 TF-IDF"""
        try:
            # 1. 检查配置是否启用
            import container
            config = container.get_config()
            if not config.get('embedding_enabled', True):
                logger.info("[搜索引擎] embedding_enabled=False，跳过向量加载")
                return

            # 2. 查找向量文件
            vector_path = find_vector_index()
            if vector_path is None:
                logger.warning("[搜索引擎] 向量文件未找到，降级 TF-IDF")
                return

            # 3. 加载并校验
            with open(vector_path, 'r', encoding='utf-8') as f:
                vector_data = json.load(f)

            expected_dims = vector_data.get('dimensions', 0)
            vectors_dict = vector_data.get('vectors', {})

            if not vectors_dict:
                logger.warning("[搜索引擎] 向量文件为空，降级 TF-IDF")
                return

            # 4. 构建有序矩阵（与 records 列表对齐）
            if HAS_NUMPY:
                vector_ids = []
                matrix_rows = []
                for i in range(len(self.records)):
                    key = f"rec_{i:03d}"
                    if key in vectors_dict:
                        vec = vectors_dict[key]
                        if len(vec) != expected_dims:
                            logger.warning(f"[搜索引擎] 向量维度不匹配（期望 {expected_dims}，实际 {len(vec)}），降级 TF-IDF")
                            return
                        vector_ids.append(key)
                        matrix_rows.append(vec)

                if not matrix_rows:
                    logger.warning("[搜索引擎] 无有效向量与记录匹配，降级 TF-IDF")
                    return

                self._vector_ids = vector_ids
                self._vector_matrix = np.array(matrix_rows, dtype=np.float32)
                # 归一化（用于余弦相似度）
                norms = np.linalg.norm(self._vector_matrix, axis=1, keepdims=True)
                norms[norms == 0] = 1.0
                self._vector_matrix = self._vector_matrix / norms
                self._vectors = vectors_dict
                self._embedding_enabled = True
            else:
                # 无 numpy 时用纯 Python
                self._vector_ids = []
                for i in range(len(self.records)):
                    key = f"rec_{i:03d}"
                    if key in vectors_dict:
                        self._vector_ids.append(key)
                if self._vector_ids:
                    self._vectors = vectors_dict
                    self._embedding_enabled = True
                else:
                    logger.warning("[搜索引擎] 无有效向量与记录匹配，降级 TF-IDF")

        except Exception as e:
            logger.warning(f"[搜索引擎] 向量加载失败，降级 TF-IDF：{e}")
            self._vectors = None
            self._embedding_enabled = False

    def _get_query_vector(self, query):
        """获取 query 向量，带 LRU 缓存。失败返回 None（触发降级 TF-IDF）"""
        # 1. 查缓存
        if query in self._query_vector_cache:
            # LRU：移到末尾
            self._query_vector_cache.move_to_end(query)
            return self._query_vector_cache[query]

        # 2. 调 API
        vector = self._call_embedding_api(query)
        if vector is None:
            return None

        # 3. 存缓存（LRU 淘汰）
        self._query_vector_cache[query] = vector
        if len(self._query_vector_cache) > _QUERY_VECTOR_CACHE_MAX:
            self._query_vector_cache.popitem(last=False)  # 淘汰最旧

        return vector

    def _call_embedding_api(self, text):
        """调用 OpenAI 兼容 /v1/embeddings API"""
        try:
            import container
            import requests

            config = container.get_config()
            api_base = config.get('claude_api_base', os.environ.get('CLAUDE_API_BASE', ''))
            api_key = config.get('claude_api_key', os.environ.get('CLAUDE_API_KEY', ''))
            model = config.get('embedding_model', 'text-embedding-3-small')
            timeout = config.get('embedding_api_timeout', 10)

            if not api_base or not api_key:
                logger.warning("[搜索引擎] API 配置缺失，降级 TF-IDF")
                return None

            url = f"{api_base.rstrip('/')}/v1/embeddings"
            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {api_key}",
            }
            payload = {"model": model, "input": [text]}

            resp = requests.post(url, json=payload, headers=headers, timeout=timeout)
            resp.raise_for_status()
            data = resp.json()

            embedding = data["data"][0]["embedding"]
            return embedding

        except Exception as e:
            logger.warning(f"[搜索引擎] Embedding API 调用失败，降级 TF-IDF：{e}")
            return None

    def _search_by_vector(self, query_vector, top_k, project_filter, category_filter):
        """向量余弦相似度搜索，返回格式与 TF-IDF 路径一致"""
        if HAS_NUMPY and self._vector_matrix is not None:
            # numpy 加速路径
            qv = np.array(query_vector, dtype=np.float32)
            qv_norm = np.linalg.norm(qv)
            if qv_norm == 0:
                return []
            qv = qv / qv_norm

            # 批量余弦相似度
            similarities = self._vector_matrix @ qv  # (N,)

            # 构建结果（含过滤）
            scores = []
            for idx, key in enumerate(self._vector_ids):
                record_idx = int(key.split('_')[1])
                record = self.records[record_idx]

                if project_filter:
                    if record.get('project', '').lower() != project_filter.lower():
                        continue
                if category_filter:
                    if category_filter not in record.get('category', []):
                        continue

                score = float(similarities[idx])

                # 关键词匹配加分（与 TF-IDF 路径一致）
                matched_keywords = 0
                for kw in record.get('keywords', []):
                    if kw.lower() in self._last_query_lower:
                        matched_keywords += 1
                keyword_bonus = matched_keywords * 0.05

                # 项目匹配加分
                project_bonus = 0
                if project_filter and record.get('project', '').lower() == project_filter.lower():
                    project_bonus = 0.1

                final_score = min(1.0, score + keyword_bonus + project_bonus)
                if final_score > 0:
                    scores.append((record_idx, final_score))
        else:
            # 纯 Python 降级路径
            scores = []
            for key in self._vector_ids:
                record_idx = int(key.split('_')[1])
                record = self.records[record_idx]

                if project_filter:
                    if record.get('project', '').lower() != project_filter.lower():
                        continue
                if category_filter:
                    if category_filter not in record.get('category', []):
                        continue

                doc_vec = self._vectors[key]
                score = _cosine_sim_python(query_vector, doc_vec)

                matched_keywords = 0
                for kw in record.get('keywords', []):
                    if kw.lower() in self._last_query_lower:
                        matched_keywords += 1
                keyword_bonus = matched_keywords * 0.05

                project_bonus = 0
                if project_filter and record.get('project', '').lower() == project_filter.lower():
                    project_bonus = 0.1

                final_score = min(1.0, score + keyword_bonus + project_bonus)
                if final_score > 0:
                    scores.append((record_idx, final_score))

        # 排序 + top_k
        scores.sort(key=lambda x: x[1], reverse=True)
        results = []
        for rank, (idx, score) in enumerate(scores[:top_k], start=1):
            results.append({
                'rank': rank,
                'score': round(score, 3),
                'record': self.records[idx],
            })
        return results

    # ============================================================
    # 公开接口（签名零变更）
    # ============================================================

    def search(self, query, top_k=5, project_filter=None, category_filter=None):
        """
        搜索相似问题记录。

        参数：
          query: 搜索查询字符串
          top_k: 返回结果数量
          project_filter: 项目名称过滤（可选）
          category_filter: 分类过滤（可选）

        返回：
          [
            {
              'rank': 1,
              'score': 0.85,
              'record': {...}  # 原始记录
            },
            ...
          ]
        """
        if not self.records:
            return []

        # P2-1：优先向量搜索
        if self._vectors and self._embedding_enabled:
            self._last_query_lower = query.lower()
            query_vector = self._get_query_vector(query)
            if query_vector is not None:
                return self._search_by_vector(query_vector, top_k, project_filter, category_filter)
            # API 失败，降级 TF-IDF
            logger.warning("[搜索引擎] query 向量获取失败，降级 TF-IDF")

        # TF-IDF 路径（原有逻辑，降级兜底）
        return self._search_by_tfidf(query, top_k, project_filter, category_filter)

    def _search_by_tfidf(self, query, top_k, project_filter, category_filter):
        """TF-IDF 搜索（原有逻辑，从 search() 拆出）"""
        # 对查询进行分词和向量化
        query_tokens = tokenize(query)
        query_tf = compute_tf(query_tokens)
        query_vector = compute_tfidf(query_tf, self.idf)

        # 计算与每条记录的相似度
        scores = []
        for i, doc_vector in enumerate(self.doc_vectors):
            # 应用过滤器
            record = self.records[i]

            if project_filter:
                if record.get('project', '').lower() != project_filter.lower():
                    continue

            if category_filter:
                if category_filter not in record.get('category', []):
                    continue

            # 计算相似度
            similarity = cosine_similarity(query_vector, doc_vector)

            # 关键词匹配加分
            matched_keywords = 0
            for kw in record.get('keywords', []):
                if kw.lower() in query.lower():
                    matched_keywords += 1
            keyword_bonus = matched_keywords * 0.05  # 每匹配一个关键词加 5%

            # 项目匹配加分
            project_bonus = 0
            if project_filter and record.get('project', '').lower() == project_filter.lower():
                project_bonus = 0.1  # 同项目加 10%

            # 综合得分
            final_score = min(1.0, similarity + keyword_bonus + project_bonus)

            if final_score > 0:
                scores.append((i, final_score))

        # 按得分排序
        scores.sort(key=lambda x: x[1], reverse=True)

        # 返回 Top K
        results = []
        for rank, (idx, score) in enumerate(scores[:top_k], start=1):
            results.append({
                'rank': rank,
                'score': round(score, 3),
                'record': self.records[idx],
            })

        return results

    def get_projects(self):
        """获取所有项目名称列表"""
        return self.index.get('project_names', [])

    def get_categories(self):
        """获取所有分类列表"""
        return list(self.index.get('categories', {}).keys())

    def get_search_mode(self):
        """获取当前搜索模式（供 /api/health 使用）"""
        return {
            'mode': 'vector' if self._embedding_enabled else 'tfidf',
            'vector_loaded': self._vectors is not None,
            'embedding_model': self._get_embedding_model_name(),
        }

    def _get_embedding_model_name(self):
        """获取配置中的 embedding 模型名"""
        try:
            import container
            return container.get_config().get('embedding_model', 'text-embedding-3-small')
        except Exception:
            return 'text-embedding-3-small'


# 模块级单例
_engine = None


def get_engine():
    """获取搜索引擎单例"""
    global _engine
    if _engine is None:
        _engine = SearchEngine()
    return _engine


if __name__ == "__main__":
    # 测试
    engine = SearchEngine()

    test_queries = [
        "门口屏 MQTT 连接失败",
        "redis 连接失败",
        "预定系统启动失败",
        "配置文件不生效",
    ]

    for q in test_queries:
        print(f"\n{'='*50}")
        print(f"查询：{q}")
        print(f"{'='*50}")

        results = engine.search(q, top_k=3)
        for r in results:
            rec = r['record']
            print(f"[{r['rank']}] 得分：{r['score']}")
            print(f"    项目：{rec.get('project', '无')}")
            print(f"    标题：{rec.get('title', '无')}")
            print(f"    分类：{rec.get('category', [])}")
