"""
knowledge/rag.py - 轻量级 RAG 知识检索模块
==========================================
基于 TF-IDF 的文本检索系统，支持:
  - 读取知识库目录中的文本文件
  - 文本分块 (chunk_size, overlap)
  - TF-IDF 向量化与余弦相似度搜索
  - 支持组织级别和 Agent 级别的知识库

支持文件类型: .md, .txt, .json, .csv, .py, .js, .html
"""
from __future__ import annotations

import json
import math
import os
import re
from collections import Counter
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

from core.logger import get_logger
from core.utils import timestamp

logger = get_logger("myagent.knowledge.rag")

# 支持的文件扩展名
SUPPORTED_EXTENSIONS = {".md", ".txt", ".json", ".csv", ".py", ".js", ".html"}


@dataclass
class KnowledgeChunk:
    """知识文本块"""
    chunk_id: str = ""
    file_path: str = ""         # 相对路径
    file_name: str = ""
    content: str = ""
    chunk_index: int = 0        # 在原文件中的块索引
    score: float = 0.0          # 搜索得分

    def to_dict(self) -> dict:
        return {
            "chunk_id": self.chunk_id,
            "file_path": self.file_path,
            "file_name": self.file_name,
            "content": self.content[:2000],  # 限制返回长度
            "chunk_index": self.chunk_index,
            "score": round(self.score, 4),
        }


class KnowledgeRAG:
    """
    轻量级 RAG 知识检索引擎。

    使用示例:
        rag = KnowledgeRAG(kb_dir="~/.myagent/data/organization/knowledge")
        rag.build_index()
        results = rag.search("公司报销流程", top_k=5)
    """

    def __init__(
        self,
        kb_dir: str | Path = "",
        chunk_size: int = 500,
        overlap: int = 50,
    ):
        """
        Args:
            kb_dir: 知识库目录路径
            chunk_size: 文本分块大小（字符数）
            overlap: 分块重叠大小（字符数）
        """
        self.kb_dir = Path(kb_dir) if kb_dir else Path()
        self.chunk_size = chunk_size
        self.overlap = overlap
        # 索引数据: {chunk_id: KnowledgeChunk}
        self._chunks: Dict[str, KnowledgeChunk] = {}
        # 原始文档缓存: {相对路径: 完整内容}
        self._documents: Dict[str, str] = {}

    # ==========================================================================
    # TF-IDF 分词（复用 memory/manager.py 的中文分词策略）
    # ==========================================================================

    @staticmethod
    def _tokenize(text: str) -> List[str]:
        """
        中文分词（简单实现：单字+双字组合）+ 英文词提取。
        """
        if not text:
            return []

        tokens: List[str] = []
        # 提取中文连续片段
        chinese_segments = re.findall(r'[\u4e00-\u9fff]+', text)
        for seg in chinese_segments:
            tokens.extend(list(seg))
            for i in range(len(seg) - 1):
                tokens.append(seg[i:i + 2])

        # 提取英文/数字单词
        english_words = re.findall(r'[a-zA-Z0-9]+', text.lower())
        tokens.extend(english_words)

        return tokens

    @staticmethod
    def _compute_tfidf(
        query: str,
        documents: List[Tuple[str, str]],
    ) -> Dict[str, float]:
        """
        计算 TF-IDF 相似度得分。

        Args:
            query: 查询文本
            documents: 文档列表 [(doc_id, text), ...]

        Returns:
            {doc_id: tfidf_score} 按得分降序排列
        """
        if not documents or not query:
            return {}

        _tokenize = KnowledgeRAG._tokenize

        query_tokens = _tokenize(query)
        if not query_tokens:
            return {}

        n_docs = len(documents)

        # 计算每个文档的 TF
        doc_tfs: Dict[str, Counter] = {}
        for doc_id, text in documents:
            doc_tfs[doc_id] = Counter(_tokenize(text))

        # 计算 IDF
        doc_freq: Counter = Counter()
        for doc_id, tf in doc_tfs.items():
            for token in set(tf.keys()):
                doc_freq[token] += 1

        idf: Dict[str, float] = {}
        for token, df in doc_freq.items():
            idf[token] = math.log((n_docs + 1) / (1 + df)) + 1

        # 查询的 TF
        query_tf = Counter(query_tokens)

        # 计算查询向量的模
        query_norm = math.sqrt(
            sum((query_tf.get(t, 0) * idf.get(t, 0)) ** 2 for t in query_tf)
        )
        if query_norm == 0:
            return {}

        # 余弦相似度
        scores: Dict[str, float] = {}
        for doc_id in doc_tfs:
            doc_tf = doc_tfs[doc_id]
            dot_product = 0.0
            for token in query_tf:
                if token in doc_tf:
                    dot_product += (query_tf[token] * idf.get(token, 0)) * \
                                   (doc_tf[token] * idf.get(token, 0))

            doc_norm = math.sqrt(
                sum((doc_tf.get(t, 0) * idf.get(t, 0)) ** 2 for t in doc_tf)
            ) if doc_tf else 0

            if doc_norm > 0:
                scores[doc_id] = dot_product / (query_norm * doc_norm)
            else:
                scores[doc_id] = 0.0

        return dict(sorted(scores.items(), key=lambda x: x[1], reverse=True))

    # ==========================================================================
    # 文本分块
    # ==========================================================================

    def _split_text(self, text: str) -> List[str]:
        """
        将文本按段落分块，保持语义完整性。

        优先在段落边界（空行）处分割，不足 chunk_size 时按字符切分。
        """
        if not text or not text.strip():
            return []

        # 按空行分段
        paragraphs = re.split(r'\n\s*\n', text)
        paragraphs = [p.strip() for p in paragraphs if p.strip()]

        if not paragraphs:
            return []

        chunks: List[str] = []
        current_chunk = ""

        for para in paragraphs:
            if not current_chunk:
                current_chunk = para
            elif len(current_chunk) + len(para) + 2 <= self.chunk_size:
                current_chunk += "\n\n" + para
            else:
                # 当前块已满，先保存
                chunks.append(current_chunk)
                # 如果段落本身超过 chunk_size，按行切分
                if len(para) > self.chunk_size:
                    # 先保存已有部分
                    lines = para.split('\n')
                    current_chunk = ""
                    for line in lines:
                        if not current_chunk:
                            current_chunk = line
                        elif len(current_chunk) + len(line) + 1 <= self.chunk_size:
                            current_chunk += "\n" + line
                        else:
                            chunks.append(current_chunk)
                            # 重叠: 保留最后 overlap 个字符
                            if self.overlap > 0 and len(current_chunk) > self.overlap:
                                current_chunk = current_chunk[-self.overlap:] + "\n" + line
                            else:
                                current_chunk = line
                else:
                    # 重叠: 保留前一块的末尾
                    if self.overlap > 0 and len(current_chunk) > self.overlap:
                        current_chunk = current_chunk[-self.overlap:] + "\n\n" + para
                    else:
                        current_chunk = para

        if current_chunk:
            chunks.append(current_chunk)

        return chunks

    # ==========================================================================
    # 索引管理
    # ==========================================================================

    def _read_file_content(self, file_path: Path) -> str:
        """读取文件内容，支持多种文本格式"""
        try:
            content = file_path.read_text(encoding="utf-8", errors="replace")
            # JSON 文件提取纯文本
            if file_path.suffix == ".json":
                try:
                    obj = json.loads(content)
                    content = json.dumps(obj, ensure_ascii=False, indent=2)
                except json.JSONDecodeError:
                    pass
            return content
        except Exception as e:
            logger.warning(f"读取文件失败 {file_path}: {e}")
            return ""

    def build_index(self) -> int:
        """
        扫描知识库目录，构建全文索引。

        Returns:
            索引的块数量
        """
        if not self.kb_dir.exists():
            logger.warning(f"知识库目录不存在: {self.kb_dir}")
            return 0

        self._chunks.clear()
        self._documents.clear()

        file_count = 0
        chunk_count = 0

        # 递归扫描文件
        for file_path in sorted(self.kb_dir.rglob("*")):
            if not file_path.is_file():
                continue
            if file_path.suffix.lower() not in SUPPORTED_EXTENSIONS:
                continue

            content = self._read_file_content(file_path)
            if not content.strip():
                continue

            rel_path = str(file_path.relative_to(self.kb_dir))
            self._documents[rel_path] = content

            # 分块
            text_chunks = self._split_text(content)
            for idx, chunk_text in enumerate(text_chunks):
                chunk_id = f"{rel_path}#chunk_{idx}"
                kc = KnowledgeChunk(
                    chunk_id=chunk_id,
                    file_path=rel_path,
                    file_name=file_path.name,
                    content=chunk_text,
                    chunk_index=idx,
                )
                self._chunks[chunk_id] = kc
                chunk_count += 1

            file_count += 1

        logger.info(
            f"知识库索引构建完成: {file_count} 个文件, {chunk_count} 个文本块 "
            f"(dir={self.kb_dir})"
        )
        return chunk_count

    def add_document(self, file_path: str, content: str = "") -> bool:
        """
        添加单个文档到索引。

        Args:
            file_path: 文件相对路径（相对于知识库根目录）
            content: 文件内容，如果为空则从文件读取

        Returns:
            是否添加成功
        """
        if not content:
            abs_path = self.kb_dir / file_path
            if not abs_path.exists():
                logger.warning(f"文件不存在: {abs_path}")
                return False
            content = self._read_file_content(abs_path)

        if not content.strip():
            return False

        # 如果已有旧块，先移除
        self.remove_document(file_path)

        self._documents[file_path] = content

        # 分块
        text_chunks = self._split_text(content)
        for idx, chunk_text in enumerate(text_chunks):
            chunk_id = f"{file_path}#chunk_{idx}"
            kc = KnowledgeChunk(
                chunk_id=chunk_id,
                file_path=file_path,
                file_name=Path(file_path).name,
                content=chunk_text,
                chunk_index=idx,
            )
            self._chunks[chunk_id] = kc

        logger.info(f"已添加文档到索引: {file_path} ({len(text_chunks)} 块)")
        return True

    def remove_document(self, file_path: str) -> int:
        """
        从索引中移除文档。

        Args:
            file_path: 文件相对路径

        Returns:
            移除的块数量
        """
        removed = 0
        to_remove = [cid for cid in self._chunks if cid.startswith(f"{file_path}#")]
        for cid in to_remove:
            del self._chunks[cid]
            removed += 1

        self._documents.pop(file_path, None)
        if removed > 0:
            logger.info(f"已从索引移除文档: {file_path} ({removed} 块)")
        return removed

    # ==========================================================================
    # 搜索
    # ==========================================================================

    def search(self, query: str, top_k: int = 5) -> List[KnowledgeChunk]:
        """
        搜索知识库，返回最相关的文本块。

        Args:
            query: 查询文本
            top_k: 返回前 k 个结果

        Returns:
            按相关性排序的 KnowledgeChunk 列表
        """
        if not query or not self._chunks:
            return []

        # 构建文档列表
        documents = [
            (cid, chunk.content)
            for cid, chunk in self._chunks.items()
        ]

        # 计算 TF-IDF 得分
        scores = self._compute_tfidf(query, documents)

        # 取 top_k（创建副本避免修改共享对象）
        results = []
        for chunk_id, score in list(scores.items())[:top_k]:
            if chunk_id in self._chunks:
                original = self._chunks[chunk_id]
                # 创建副本并设置 score，避免修改 _chunks 中的共享引用
                chunk = KnowledgeChunk(
                    chunk_id=original.chunk_id,
                    file_path=original.file_path,
                    file_name=original.file_name,
                    content=original.content,
                    chunk_index=original.chunk_index,
                    score=score,
                )
                results.append(chunk)

        return results

    def get_all_files(self) -> List[Dict[str, Any]]:
        """
        列出知识库中所有已索引的文件。

        Returns:
            文件信息列表 [{path, name, chunks, size}]
        """
        file_info: Dict[str, Dict] = {}
        for cid, chunk in self._chunks.items():
            fp = chunk.file_path
            if fp not in file_info:
                content = self._documents.get(fp, "")
                file_info[fp] = {
                    "path": fp,
                    "name": chunk.file_name,
                    "chunks": 0,
                    "size": len(content),
                }
            file_info[fp]["chunks"] += 1

        return list(file_info.values())

    @property
    def total_chunks(self) -> int:
        """当前索引中的总块数"""
        return len(self._chunks)

    @property
    def total_documents(self) -> int:
        """当前索引中的总文档数"""
        return len(self._documents)