import os
from appPublic.log import debug, error, info
from base_connection import connection_register
from typing import Dict, List, Any
import numpy as np
import aiohttp
from aiohttp import ClientSession, ClientTimeout
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
import uuid
from datetime import datetime
from filetxt.loader import fileloader
import time
from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
import traceback
import asyncio
import re
# 嵌入缓存
EMBED_CACHE = {}

class MilvusConnection:
    def __init__(self):
        pass

    @retry(stop = stop_after_attempt(3))
    async def _make_neo4japi_request(self, action: str, params: Dict[str, Any]) -> Dict[str, Any]:
        debug(f"开始API请求：action={action}, params={params}")
        try:
            async with ClientSession(timeout=ClientTimeout(total=300)) as session:
                url = f"http://localhost:8885/v1/{action}"
                debug(f"发起POST请求：{url}")
                async with session.post(
                    url,
                    headers={'Content-Type': 'application/json'},
                    json=params
                ) as response:
                    debug(f"收到相应: status={response.status}, headers={response.headers}")
                    respose_text = await response.text()
                    debug(f"响应内容: {respose_text}")
                    result = await response.json()
                    debug(f"API响应内容: {result}")
                    if response.status == 400:
                        debug(f"客户端错误，状态码: {response.status},返回响应: {result}")
                        return result
                    if response.status != 200:
                        error(f"API 调用失败，动作: {action}, 状态码: {response.status}, 响应: {response_text}")
                        raise RuntimeError(f"API 调用失败: {response.status}")
                    debug(f"API 调用成功: {action}, 响应: {result}")
                    return result
        except Exception as e:
            error(f"API 调用失败: {action}, 错误: {str(e)}, 堆栈: {traceback.format_exc()}")
            raise RuntimeError(f"API 调用失败: {str(e)}")

    @retry(stop=stop_after_attempt(3))
    async def _make_api_request(self, action: str, params: Dict[str, Any]) -> Dict[str, Any]:
        debug(f"开始 API 请求: action={action}, params={params}")
        try:
            async with ClientSession(timeout=ClientTimeout(total=300)) as session:
                url = f"http://localhost:8886/v1/{action}"
                debug(f"发起 POST 请求: {url}")
                async with session.post(
                        url,
                        headers={"Content-Type": "application/json"},
                        json=params
                ) as response:
                    debug(f"收到响应: status={response.status}, headers={response.headers}")
                    response_text = await response.text()
                    debug(f"响应内容: {response_text}")
                    result = await response.json()
                    debug(f"API 响应内容: {result}")
                    if response.status == 400:  # 客户端错误，直接返回
                        debug(f"客户端错误，状态码: {response.status}, 返回响应: {result}")
                        return result
                    if response.status != 200:
                        error(f"API 调用失败，动作: {action}, 状态码: {response.status}, 响应: {response_text}")
                        raise RuntimeError(f"API 调用失败: {response.status}")
                    debug(f"API 调用成功: {action}, 响应: {result}")
                    return result
        except Exception as e:
            error(f"API 调用失败: {action}, 错误: {str(e)}, 堆栈: {traceback.format_exc()}")
            raise RuntimeError(f"API 调用失败: {str(e)}")

    async def handle_connection(self, action: str, params: Dict = None) -> Dict:
        """处理数据库操作"""
        try:
            debug(f"处理操作: action={action}, params={params}")
            if not params:
                params = {}
            # 通用 db_type 验证
            db_type = params.get("db_type", "")
            collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
            if db_type and "_" in db_type:
                return {"status": "error", "message": "db_type 不能包含下划线", "collection_name": collection_name,
                        "document_id": "", "status_code": 400}
            if db_type and len(db_type) > 100:
                return {"status": "error", "message": "db_type 的长度应小于 100", "collection_name": collection_name,
                        "document_id": "", "status_code": 400}

            if action == "initialize":
                return {"status": "success", "message": "Milvus 服务已初始化"}
            elif action == "get_params":
                return {"status": "success", "params": {}}
            elif action == "create_collection":
                return await self._create_collection(db_type)
            elif action == "delete_collection":
                return await self._delete_collection(db_type)
            elif action == "insert_document":
                file_path = params.get("file_path", "")
                userid = params.get("userid", "")
                knowledge_base_id = params.get("knowledge_base_id", "")
                document_id = params.get("document_id", "")
                if not file_path or not userid or not knowledge_base_id or not document_id:
                    return {"status": "error", "message": "file_path、userid document_id和 knowledge_base_id 不能为空",
                            "collection_name": collection_name, "document_id": "", "status_code": 400}
                if "_" in userid or "_" in knowledge_base_id:
                    return {"status": "error", "message": "userid 和 knowledge_base_id 不能包含下划线",
                            "collection_name": collection_name, "document_id": document_id, "status_code": 400}
                if len(knowledge_base_id) > 100:
                    return {"status": "error", "message": "knowledge_base_id 的长度应小于 100",
                            "collection_name": collection_name, "document_id": "", "status_code": 400}
                return await self._insert_document(file_path, userid, knowledge_base_id, document_id, db_type)
            elif action == "delete_document":
                userid = params.get("userid", "")
                file_path = params.get("file_path", "")
                knowledge_base_id = params.get("knowledge_base_id", "")
                document_id = params.get("document_id", "")
                if not userid or not file_path or not knowledge_base_id or not document_id:
                    return {"status": "error", "message": "userid、file_path document_id和 knowledge_base_id 不能为空",
                            "collection_name": collection_name, "document_id": "", "status_code": 400}
                if "_" in userid or "_" in knowledge_base_id:
                    return {"status": "error", "message": "userid 和 knowledge_base_id 不能包含下划线",
                            "collection_name": collection_name, "document_id": "", "status_code": 400}
                if len(userid) > 100 or len(file_path) > 255 or len(knowledge_base_id) > 100:
                    return {"status": "error", "message": "userid、file_path 或 knowledge_base_id 的长度超出限制",
                            "collection_name": collection_name, "document_id": "", "status_code": 400}
                return await self._delete_document(userid, file_path, knowledge_base_id, document_id, db_type)
            elif action == "delete_knowledge_base":
                userid = params.get("userid", "")
                knowledge_base_id = params.get("knowledge_base_id", "")
                if not userid or not knowledge_base_id:
                    return {"status": "error", "message": "userid 和 knowledge_base_id 不能为空",
                            "collection_name": collection_name, "document_id": "", "status_code": 400}
                if "_" in userid or "_" in knowledge_base_id:
                    return {"status": "error", "message": "userid 和 knowledge_base_id 不能包含下划线",
                            "collection_name": collection_name, "document_id": "", "status_code": 400}
                if len(userid) > 100 or len(knowledge_base_id) > 100:
                    return {"status": "error", "message": "userid 或 knowledge_base_id 的长度超出限制",
                            "collection_name": collection_name, "document_id": "", "status_code": 400}
                return await self._delete_knowledge_base(db_type, userid, knowledge_base_id)
            elif action == "search_query":
                query = params.get("query", "")
                userid = params.get("userid", "")
                knowledge_base_ids = params.get("knowledge_base_ids", [])
                limit = params.get("limit", 5)
                offset = params.get("offset", 0)
                db_type = params.get("db_type", "")
                use_rerank = params.get("use_rerank", True)
                if not query or not userid or not knowledge_base_ids:
                    return {"status": "error", "message": "query、userid 或 knowledge_base_ids 不能为空",
                            "collection_name": collection_name, "document_id": "", "status_code": 400}
                return await self._search_query(query, userid, knowledge_base_ids, limit, offset, use_rerank, db_type)
            elif action == "fused_search":
                query = params.get("query", "")
                userid = params.get("userid", "")
                knowledge_base_ids = params.get("knowledge_base_ids", [])
                limit = params.get("limit", 5)
                offset = params.get("offset", 0)
                db_type = params.get("db_type", "")
                use_rerank = params.get("use_rerank", True)
                if not query or not userid or not knowledge_base_ids:
                    return {"status": "error", "message": "query、userid 或 knowledge_base_ids 不能为空",
                            "collection_name": collection_name, "document_id": "", "status_code": 400}
                return await self._fused_search(query, userid, knowledge_base_ids, limit, offset, use_rerank, db_type)
            elif action == "list_user_files":
                userid = params.get("userid", "")
                if not userid:
                    return {"status": "error", "message": "userid 不能为空", "collection_name": collection_name,
                            "document_id": "", "status_code": 400}
                return await self._list_user_files(userid, db_type)
            elif action == "list_all_knowledge_bases":
                return await self._list_all_knowledge_bases(db_type)
            else:
                return {"status": "error", "message": f"未知的 action: {action}", "collection_name": collection_name,
                        "document_id": "", "status_code": 400}
        except Exception as e:
            error(f"处理操作失败: action={action}, 错误: {str(e)}")
            return {
                "status": "error",
                "message": f"服务器错误: {str(e)}",
                "collection_name": collection_name,
                "document_id": "",
                "status_code": 400
            }

    async def _create_collection(self, db_type: str = "") -> Dict[str, Any]:
        """创建 Milvus 集合"""
        collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
        try:
            if len(collection_name) > 255:
                raise ValueError(f"集合名称 {collection_name} 超过 255 个字符")
            if len(db_type) > 100:
                raise ValueError("db_type 的长度应小于 100")
            debug(f"调用创建集合端点: {collection_name}, 参数: {{'db_type': '{db_type}'}}")
            result = await self._make_api_request("createcollection", {"db_type": db_type})
            return result
        except Exception as e:
            error(f"创建集合失败: {str(e)}, 堆栈: {traceback.format_exc()}")
            return {
                "status": "error",
                "collection_name": collection_name,
                "message": f"创建集合失败: {str(e)}",
                "status_code": 400
            }

    async def _delete_collection(self, db_type: str = "") -> Dict:
        """删除 Milvus 集合通过服务化端点"""
        try:
            collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
            if len(collection_name) > 255:
                raise ValueError(f"集合名称 {collection_name} 超过 255 个字符")
            if db_type and "_" in db_type:
                raise ValueError("db_type 不能包含下划线")
            if db_type and len(db_type) > 100:
                raise ValueError("db_type 的长度应小于 100")
            debug(f"调用删除集合端点: {collection_name}")

            result = await self._make_api_request("deletecollection", {"db_type": db_type})
            return result
        except Exception as e:
            error(f"删除集合失败: {str(e)}")
            return {
                "status": "error",
                "collection_name": collection_name,
                "message": str(e),
                "status_code": 400
            }

    async def _insert_document(self, file_path: str, userid: str, knowledge_base_id: str, document_id:str, db_type: str = "") -> Dict[
        str, Any]:
        """将文档插入 Milvus 并抽取三元组到 Neo4j"""
        collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
        debug(
            f'Inserting document: file_path={file_path}, userid={userid}, db_type={db_type}, knowledge_base_id={knowledge_base_id}, document_id={document_id}')

        timings = {}
        start_total = time.time()

        try:
            # 验证参数
            if not userid or not knowledge_base_id:
                raise ValueError("userid 和 knowledge_base_id 不能为空")
            if "_" in userid or "_" in knowledge_base_id:
                raise ValueError("userid 和 knowledge_base_id 不能包含下划线")
            if len(userid) > 100 or len(knowledge_base_id) > 100:
                raise ValueError("userid 或 knowledge_base_id 的长度超出限制")
            if not os.path.exists(file_path):
                raise ValueError(f"文件 {file_path} 不存在")

            supported_formats = {'pdf', 'docx', 'xlsx', 'pptx', 'csv', 'txt'}
            ext = file_path.rsplit('.', 1)[1].lower() if '.' in file_path else ''
            if ext not in supported_formats:
                raise ValueError(f"不支持的文件格式: {ext}, 支持的格式: {', '.join(supported_formats)}")

            info(f"生成 document_id: {document_id} for file: {file_path}")

            # 文件加载
            debug(f"加载文件: {file_path}")
            start_load = time.time()
            text = fileloader(file_path)
            text = re.sub(r'[^\u4e00-\u9fa5a-zA-Z0-9\s.;,\n]', '', text)
            timings["load_file"] = time.time() - start_load
            debug(f"加载文件耗时: {timings['load_file']:.2f} 秒, 文本长度: {len(text)}")
            if not text or not text.strip():
                raise ValueError(f"文件 {file_path} 加载为空")

            # 文本分片
            document = Document(page_content=text)
            text_splitter = RecursiveCharacterTextSplitter(
                chunk_size=500,
                chunk_overlap=100,
                length_function=len,
            )
            debug("开始分片文件内容")
            start_split = time.time()
            chunks = text_splitter.split_documents([document])
            timings["split_text"] = time.time() - start_split
            debug(
                f"文本分片耗时: {timings['split_text']:.2f} 秒, 分片数量: {len(chunks)}, 分片内容: {[chunk.page_content[:50] for chunk in chunks[:5]]}")
            if not chunks:
                raise ValueError(f"文件 {file_path} 未生成任何文档块")

            filename = os.path.basename(file_path).rsplit('.', 1)[0]
            upload_time = datetime.now().isoformat()

            # 生成嵌入向量
            debug("调用嵌入服务生成向量")
            start_embedding = time.time()
            texts = [chunk.page_content for chunk in chunks]
            embeddings = await self._get_embeddings(texts)
            if not embeddings or not all(len(vec) == 1024 for vec in embeddings):
                raise ValueError("所有嵌入向量必须是长度为 1024 的浮点数列表")
            timings["generate_embeddings"] = time.time() - start_embedding
            debug(f"生成嵌入向量耗时: {timings['generate_embeddings']:.2f} 秒, 嵌入数量: {len(embeddings)}")

            # 构造 chunks 参数（展平结构）
            chunks_data = []
            for i, chunk in enumerate(chunks):
                chunks_data.append({
                    "userid": userid,
                    "knowledge_base_id": knowledge_base_id,
                    "text": chunk.page_content,
                    "vector": embeddings[i].tolist(),
                    "document_id": document_id,
                    "filename": filename + '.' + ext,
                    "file_path": file_path,
                    "upload_time": upload_time,
                    "file_type": ext,
                })

            # 调用 Milvus 插入端点
            debug(f"调用插入文件端点: {file_path}")
            start_milvus = time.time()
            result = await self._make_api_request("insertdocument", {
                "chunks": chunks_data,
                "db_type": db_type,
            })
            timings["insert_milvus"] = time.time() - start_milvus
            debug(f"Milvus 插入耗时: {timings['insert_milvus']:.2f} 秒")

            if result.get("status") != "success":
                timings["total"] = time.time() - start_total
                return {
                    "status": "error",
                    "document_id": document_id,
                    "collection_name": collection_name,
                    "timings": timings,
                    "message": result.get("message", "未知错误"),
                    "status_code": 400
                }

            # 三元组抽取
            debug("调用三元组抽取服务")
            start_triples = time.time()
            try:
                chunk_texts = [doc.page_content for doc in chunks]
                debug(f"处理 {len(chunk_texts)} 个分片进行三元组抽取")

                tasks = [self._extract_triples(chunk) for chunk in chunk_texts]
                results = await asyncio.gather(*tasks, return_exceptions=True)

                triples = []
                for i, result in enumerate(results):
                    if isinstance(result, list):
                        triples.extend(result)
                        debug(f"分片 {i + 1} 抽取到 {len(result)} 个三元组: {result[:5]}")
                    else:
                        error(f"分片 {i + 1} 处理失败: {str(result)}")

                # 去重
                unique_triples = []
                seen = set()
                for t in triples:
                    identifier = (t['head'].lower(), t['tail'].lower(), t['type'].lower())
                    if identifier not in seen:
                        seen.add(identifier)
                        unique_triples.append(t)
                    else:
                        for existing in unique_triples:
                            if (existing['head'].lower() == t['head'].lower() and
                                    existing['tail'].lower() == t['tail'].lower() and
                                    len(t['type']) > len(existing['type'])):
                                unique_triples.remove(existing)
                                unique_triples.append(t)
                                debug(f"替换三元组为更具体类型: {t}")
                                break

                timings["extract_triples"] = time.time() - start_triples
                debug(
                    f"三元组抽取耗时: {timings['extract_triples']:.2f} 秒, 抽取到 {len(unique_triples)} 个三元组: {unique_triples[:5]}")

                # Neo4j 插入
                debug(f"抽取到 {len(unique_triples)} 个三元组，调用Neo4j服务插入")
                start_neo4j = time.time()
                if unique_triples:
                    neo4j_result = await self._make_neo4japi_request("inserttriples", {
                        "triples":unique_triples,
                        "document_id": document_id,
                        "knowledge_base_id": knowledge_base_id,
                        "userid": userid
                    })
                    debug(f"Neo4j服务响应: {neo4j_result}")
                    if neo4j_result.get("status") != "success":
                        timings["insert_neo4j"] = time.time() - start_neo4j
                        timings["total"] = time.time() - start_total
                        return{
                            "status": "error",
                            "document_id": document_id,
                            "collection_name": collection_name,
                            "timings": timings,
                            "message": f"Neo4j 三元组插入失败: {neo4j_result.get('message', '未知错误')}",
                            "status_code": 400
                        }
                    info(f"文件 {file_path} 三元组成功插入 Neo4j: {neo4j_result.get('message')}")
                else:
                    debug(f"文件 {file_path} 未抽取到三元组")
                timings["insert_neo4j"] = time.time() - start_neo4j
                debug(f"Neo4j 插入耗时: {timings['insert_neo4j']:.2f} 秒")

            except Exception as e:
                timings["extract_triples"] = time.time() - start_triples if "extract_triples" not in timings else \
                timings["extract_triples"]
                timings["insert_neo4j"] = time.time() - start_neo4j
                debug(f"处理三元组或 Neo4j 插入失败: {str(e)}, 堆栈: {traceback.format_exc()}")
                timings["total"] = time.time() - start_total
                return {
                    "status": "success",
                    "document_id": document_id,
                    "collection_name": collection_name,
                    "timings": timings,
                    "unique_triples": unique_triples,
                    "message": f"文件 {file_path} 成功嵌入，但三元组处理或 Neo4j 插入失败: {str(e)}",
                    "status_code": 200
                }

            timings["total"] = time.time() - start_total
            debug(f"总耗时: {timings['total']:.2f} 秒")
            return {
                "status": "success",
                "document_id": document_id,
                "collection_name": collection_name,
                "timings": timings,
                "unique_triples": unique_triples,
                "message": f"文件 {file_path} 成功嵌入并处理三元组",
                "status_code": 200
            }

        except Exception as e:
            error(f"插入文档失败: {str(e)}, 堆栈: {traceback.format_exc()}")
            timings["total"] = time.time() - start_total
            return {
                "status": "error",
                "document_id": document_id,
                "collection_name": collection_name,
                "timings": timings,
                "message": f"插入文档失败: {str(e)}",
                "status_code": 400
            }

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=1, max=10),
        retry=retry_if_exception_type((aiohttp.ClientError, RuntimeError)),
        before_sleep=lambda retry_state: debug(f"重试三元组抽取服务，第 {retry_state.attempt_number} 次")
    )
    async def _get_embeddings(self, texts: List[str]) -> List[List[float]]:
        """调用嵌入服务获取文本的向量，带缓存"""
        try:
            # 检查缓存
            uncached_texts = [text for text in texts if text not in EMBED_CACHE]
            if uncached_texts:
                async with aiohttp.ClientSession() as session:
                    async with session.post(
                        "http://localhost:9998/v1/embeddings",
                        headers={"Content-Type": "application/json"},
                        json={"input": uncached_texts}
                    ) as response:
                        if response.status != 200:
                            error(f"嵌入服务调用失败，状态码: {response.status}")
                            raise RuntimeError(f"嵌入服务调用失败: {response.status}")
                        result = await response.json()
                        if result.get("object") != "list" or not result.get("data"):
                            error(f"嵌入服务响应格式错误: {result}")
                            raise RuntimeError("嵌入服务响应格式错误")
                        embeddings = [item["embedding"] for item in result["data"]]
                        for text, embedding in zip(uncached_texts, embeddings):
                            EMBED_CACHE[text] = np.array(embedding) / np.linalg.norm(embedding)
                        debug(f"成功获取 {len(embeddings)} 个新嵌入向量，缓存大小: {len(EMBED_CACHE)}")
            # 返回缓存中的嵌入
            return [EMBED_CACHE[text] for text in texts]
        except Exception as e:
            error(f"嵌入服务调用失败: {str(e)}")
            raise RuntimeError(f"嵌入服务调用失败: {str(e)}")

    async def _extract_triples(self, text: str) -> List[Dict]:
        """调用三元组抽取服务，无超时限制"""
        request_id = str(uuid.uuid4())  # 为每个请求生成唯一 ID
        start_time = time.time()
        debug(f"Request #{request_id} started for triples extraction")
        try:
            async with aiohttp.ClientSession(
                connector=aiohttp.TCPConnector(limit=30),
                timeout=aiohttp.ClientTimeout(total=None)  # 无限等待
            ) as session:
                async with session.post(
                    "http://localhost:9991/v1/triples",
                    headers={"Content-Type": "application/json; charset=utf-8"},
                    json={"text": text}
                ) as response:
                    elapsed_time = time.time() - start_time
                    debug(f"Request #{request_id} received response, status: {response.status}, took {elapsed_time:.2f} seconds")
                    if response.status != 200:
                        error_text = await response.text()
                        error(f"Request #{request_id} failed, status: {response.status}, response: {error_text}")
                        raise RuntimeError(f"三元组抽取服务调用失败: {response.status}, {error_text}")
                    result = await response.json()
                    if result.get("object") != "list" or not result.get("data"):
                        error(f"Request #{request_id} invalid response format: {result}")
                        raise RuntimeError("三元组抽取服务响应格式错误")
                    triples = result["data"]
                    debug(f"Request #{request_id} extracted {len(triples)} triples, total time: {elapsed_time:.2f} seconds")
                    return triples
        except Exception as e:
            elapsed_time = time.time() - start_time
            error(f"Request #{request_id} failed to extract triples: {str(e)}, took {elapsed_time:.2f} seconds")
            debug(f"Request #{request_id} traceback: {traceback.format_exc()}")
            raise RuntimeError(f"三元组抽取服务调用失败: {str(e)}")

    async def _delete_document(self, userid: str, file_path: str, knowledge_base_id: str, document_id:str, db_type: str = "") -> Dict[str, Any]:
        """删除用户指定文件数据，包括 Milvus 和 Neo4j 中的记录"""
        collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
        try:
            # 调用 Milvus 删除文件端点
            debug(f"调用删除文件端点: userid={userid}, file_path={file_path}, knowledge_base_id={knowledge_base_id}, document_id={document_id}")
            milvus_result = await self._make_api_request("deletedocument", {
                "userid": userid,
                "file_path": file_path,
                "knowledge_base_id": knowledge_base_id,
                "document_id": document_id,
                "db_type": db_type
            })

            if milvus_result.get("status") != "success":
                error(f"Milvus 删除文件失败: {milvus_result.get('message', '未知错误')}")
                return milvus_result

            # 调用 Neo4j 删除端点
            neo4j_deleted_nodes = 0
            neo4j_deleted_rels = 0
            try:
                debug(f"调用 Neo4j 删除文档端点: document_id={document_id}")
                neo4j_result = await self._make_neo4japi_request("deletedocument", {
                    "document_id": document_id
                })
                if neo4j_result.get("status") != "success":
                    error(
                        f"Neo4j 删除文档失败: document_id={document_id}, 错误: {neo4j_result.get('message', '未知错误')}")
                nodes_deleted = neo4j_result.get("nodes_deleted", 0)
                rels_deleted = neo4j_result.get("rels_deleted", 0)
                neo4j_deleted_nodes += nodes_deleted
                neo4j_deleted_rels += rels_deleted
                info(f"成功删除 document_id={document_id} 的 {nodes_deleted} 个 Neo4j 节点和 {rels_deleted} 个关系")
            except Exception as e:
                error(f"删除 document_id={document_id} 的 Neo4j 数据失败: {str(e)}")

            return {
                "status": "success",
                "collection_name": collection_name,
                "document_id": document_id,
                "message": f"成功删除 Milvus 记录，{neo4j_deleted_nodes} 个 Neo4j 节点，{neo4j_deleted_rels} 个 Neo4j 关系",
                "status_code": 200
            }

        except Exception as e:
            error(f"删除文档失败: {str(e)}")
            return {
                "status": "error",
                "collection_name": collection_name,
                "document_id": document_id,
                "message": f"删除文档失败: {str(e)}",
                "status_code": 400
            }

    async def _delete_knowledge_base(self, db_type: str, userid: str, knowledge_base_id: str) -> Dict[str, Any]:
        """删除用户的整个知识库，包括 Milvus 和 Neo4j 中的记录"""
        collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
        try:
            # 调用 Milvus 删除知识库端点
            debug(f"调用删除知识库端点: userid={userid}, knowledge_base_id={knowledge_base_id}")
            milvus_result = await self._make_api_request("deleteknowledgebase", {
                "userid": userid,
                "knowledge_base_id": knowledge_base_id,
                "db_type": db_type
            })

            if milvus_result.get("status") != "success":
                error(f"Milvus 删除知识库失败: {milvus_result.get('message', '未知错误')}")
                return milvus_result

            deleted_files = milvus_result.get("deleted_files", [])

            # 新增：调用 Neo4j 删除知识库端点
            neo4j_deleted_nodes = 0
            neo4j_deleted_rels = 0
            try:
                debug(f"调用 Neo4j 删除知识库端点: userid={userid}, knowledge_base_id={knowledge_base_id}")
                neo4j_result = await self._make_neo4japi_request("deleteknowledgebase", {
                    "userid": userid,
                    "knowledge_base_id": knowledge_base_id
                })
                if neo4j_result.get("status") == "success":
                    neo4j_deleted_nodes = neo4j_result.get("nodes_deleted", 0)
                    neo4j_deleted_rels = neo4j_result.get("rels_deleted", 0)
                    info(f"成功删除 {neo4j_deleted_nodes} 个 Neo4j 节点和 {neo4j_deleted_rels} 个关系")
                else:
                    error(f"Neo4j 删除知识库失败: {neo4j_result.get('message', '未知错误')}")
                    return {
                        "status": "success",
                        "collection_name": collection_name,
                        "deleted_files": deleted_files,
                        "message": f"成功删除 Milvus 知识库，{neo4j_deleted_nodes} 个 Neo4j 节点，{neo4j_deleted_rels} 个 Neo4j 关系，但 Neo4j 删除失败: {neo4j_result.get('message')}",
                        "status_code": 200
                    }
            except Exception as e:
                error(f"Neo4j 删除知识库失败: {str(e)}")
                return {
                    "status": "success",
                    "collection_name": collection_name,
                    "deleted_files": deleted_files,
                    "message": f"成功删除 Milvus 知识库，{neo4j_deleted_nodes} 个 Neo4j 节点，{neo4j_deleted_rels} 个 Neo4j 关系，但 Neo4j 删除失败: {str(e)}",
                    "status_code": 200
                }

            if not deleted_files and neo4j_deleted_nodes == 0 and neo4j_deleted_rels == 0:
                debug(f"没有删除任何记录，userid={userid}, knowledge_base_id={knowledge_base_id}")
                return {
                    "status": "success",
                    "collection_name": collection_name,
                    "deleted_files": [],
                    "message": f"没有找到 userid={userid}, knowledge_base_id={knowledge_base_id} 的记录，无需删除",
                    "status_code": 200
                }

            info(
                f"总计删除 Milvus 记录，{neo4j_deleted_nodes} 个 Neo4j 节点，{neo4j_deleted_rels} 个 Neo4j 关系，删除文件: {deleted_files}")
            return {
                "status": "success",
                "collection_name": collection_name,
                "deleted_files": deleted_files,
                "message": f"成功删除 Milvus 记录，{neo4j_deleted_nodes} 个 Neo4j 节点，{neo4j_deleted_rels} 个 Neo4j 关系，删除文件: {deleted_files}",
                "status_code": 200
            }

        except Exception as e:
            error(f"删除知识库失败: {str(e)}")
            return {
                "status": "error",
                "collection_name": collection_name,
                "deleted_files": [],
                "message": f"删除知识库失败: {str(e)}",
                "status_code": 400
            }

    async def _extract_entities(self, query: str) -> List[str]:
        """调用实体识别服务"""
        try:
            if not query:
                raise ValueError("查询文本不能为空")
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    "http://localhost:9990/v1/entities",
                    headers={"Content-Type": "application/json"},
                    json={"query": query}
                ) as response:
                    if response.status != 200:
                        error(f"实体识别服务调用失败，状态码: {response.status}")
                        raise RuntimeError(f"实体识别服务调用失败: {response.status}")
                    result = await response.json()
                    if result.get("object") != "list" or not result.get("data"):
                        error(f"实体识别服务响应格式错误: {result}")
                        raise RuntimeError("实体识别服务响应格式错误")
                    entities = result["data"]
                    unique_entities = list(dict.fromkeys(entities))  # 去重
                    debug(f"成功提取 {len(unique_entities)} 个唯一实体: {unique_entities}")
                    return unique_entities
        except Exception as e:
            error(f"实体识别服务调用失败: {str(e)}")
            return []

    async def _rerank_results(self, query: str, results: List[Dict], top_n: int) -> List[Dict]:
        """调用重排序服务"""
        try:
            if not results:
                debug("无结果需要重排序")
                return results

            if not isinstance(top_n, int) or top_n < 1:
                debug(f"无效的 top_n 参数: {top_n}，使用 len(results)={len(results)}")
                top_n = len(results)
            else:
                top_n = min(top_n, len(results))
            debug(f"重排序 top_n={top_n}, 原始结果数={len(results)}")

            documents = [result["text"] for result in results]
            async with aiohttp.ClientSession() as session:
                async with session.post(
                        "http://localhost:9997/v1/rerank",
                        headers={"Content-Type": "application/json"},
                        json={
                            "model": "rerank-001",
                            "query": query,
                            "documents": documents,
                            "top_n": top_n
                        }
                ) as response:
                    if response.status != 200:
                        error(f"重排序服务调用失败，状态码: {response.status}")
                        raise RuntimeError(f"重排序服务调用失败: {response.status}")
                    result = await response.json()
                    if result.get("object") != "rerank.result" or not result.get("data"):
                        error(f"重排序服务响应格式错误: {result}")
                        raise RuntimeError("重排序服务响应格式错误")
                    rerank_data = result["data"]
                    reranked_results = []
                    for item in rerank_data:
                        index = item["index"]
                        if index < len(results):
                            results[index]["rerank_score"] = item["relevance_score"]
                            reranked_results.append(results[index])
                    debug(f"成功重排序 {len(reranked_results)} 条结果")
                    return reranked_results[:top_n]
        except Exception as e:
            error(f"重排序服务调用失败: {str(e)}")
            return results

    async def _search_query(self, query: str, userid: str, knowledge_base_ids: List[str], limit: int = 5,
                            offset: int = 0, use_rerank: bool = True, db_type: str = "") -> Dict[str, Any]:
        """纯向量搜索，调用服务化端点"""
        start_time = time.time()
        collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
        timing_stats = {}
        try:
            info(
                f"开始纯向量搜索: query={query}, userid={userid}, db_type={db_type}, knowledge_base_ids={knowledge_base_ids}, limit={limit}, offset={offset}, use_rerank={use_rerank}")

            if not query:
                raise ValueError("查询文本不能为空")
            if not userid:
                raise ValueError("userid 不能为空")
            if "_" in userid or (db_type and "_" in db_type):
                raise ValueError("userid 和 db_type 不能包含下划线")
            if (db_type and len(db_type) > 100) or len(userid) > 100:
                raise ValueError("userid 或 db_type 的长度超出限制")
            if limit <= 0 or limit > 16384:
                raise ValueError("limit 必须在 1 到 16384 之间")
            if offset < 0:
                raise ValueError("offset 不能为负数")
            if limit + offset > 16384:
                raise ValueError("limit + offset 不能超过 16384")
            if not knowledge_base_ids:
                raise ValueError("knowledge_base_ids 不能为空")
            for kb_id in knowledge_base_ids:
                if not isinstance(kb_id, str):
                    raise ValueError(f"knowledge_base_id 必须是字符串: {kb_id}")
                if len(kb_id) > 100:
                    raise ValueError(f"knowledge_base_id 长度超出 100 个字符: {kb_id}")
                if "_" in kb_id:
                    raise ValueError(f"knowledge_base_id 不能包含下划线: {kb_id}")

            # 将查询文本转换为向量
            vector_start = time.time()
            query_vector = await self._get_embeddings([query])
            if not query_vector or not all(len(vec) == 1024 for vec in query_vector):
                raise ValueError("查询向量必须是长度为 1024 的浮点数列表")
            query_vector = query_vector[0]  # 取第一个向量
            timing_stats["vector_generation"] = time.time() - vector_start
            debug(f"生成查询向量耗时: {timing_stats['vector_generation']:.3f} 秒")

            # 调用纯向量搜索端点
            search_start = time.time()
            result = await self._make_api_request("searchquery", {
                "query_vector": query_vector.tolist(),
                "userid": userid,
                "knowledge_base_ids": knowledge_base_ids,
                "limit": limit,
                "offset": offset,
                "db_type": db_type
            })
            timing_stats["vector_search"] = time.time() - search_start
            debug(f"向量搜索耗时: {timing_stats['vector_search']:.3f} 秒")

            if result.get("status") != "success":
                error(f"纯向量搜索失败: {result.get('message', '未知错误')}")
                return {"results": [], "timing": timing_stats}

            unique_results = result.get("results", [])
            if use_rerank and unique_results:
                rerank_start = time.time()
                debug("开始重排序")
                unique_results = await self._rerank_results(query, unique_results, limit)
                unique_results = sorted(unique_results, key=lambda x: x.get('rerank_score', 0), reverse=True)
                timing_stats["reranking"] = time.time() - rerank_start
                debug(f"重排序耗时: {timing_stats['reranking']:.3f} 秒")
                debug(f"重排序分数分布: {[round(r.get('rerank_score', 0), 3) for r in unique_results]}")
            else:
                unique_results = [{k: v for k, v in r.items() if k != 'rerank_score'} for r in unique_results]

            timing_stats["total_time"] = time.time() - start_time
            info(f"纯向量搜索完成，返回 {len(unique_results)} 条结果，总耗时: {timing_stats['total_time']:.3f} 秒")
            return {"results": unique_results[:limit], "timing": timing_stats}

        except Exception as e:
            error(f"纯向量搜索失败: {str(e)}, 堆栈: {traceback.format_exc()}")
            return {"results": [], "timing": timing_stats}

    async def _fused_search(self, query: str, userid: str, knowledge_base_ids: List[str], limit: int = 5,
                            offset: int = 0, use_rerank: bool = True, db_type: str = "") -> Dict[str, Any]:
        """融合搜索，调用服务化端点"""
        start_time = time.time()
        collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
        timing_stats = {}
        try:
            info(
                f"开始融合搜索: query={query}, userid={userid}, db_type={db_type}, knowledge_base_ids={knowledge_base_ids}, limit={limit}, offset={offset}, use_rerank={use_rerank}")

            if not query or not userid or not knowledge_base_ids:
                raise ValueError("query、userid 和 knowledge_base_ids 不能为空")
            if "_" in userid or (db_type and "_" in db_type):
                raise ValueError("userid 和 db_type 不能包含下划线")
            if (db_type and len(db_type) > 100) or len(userid) > 100:
                raise ValueError("db_type 或 userid 的长度超出限制")
            if limit < 1 or limit > 16384 or offset < 0:
                raise ValueError("limit 必须在 1 到 16384 之间，offset 必须大于或等于 0")

            # 提取实体
            entity_extract_start = time.time()
            query_entities = await self._extract_entities(query)
            timing_stats["entity_extraction"] = time.time() - entity_extract_start
            debug(f"提取实体: {query_entities}, 耗时: {timing_stats['entity_extraction']:.3f} 秒")

            # 调用 Neo4j 服务进行三元组匹配
            all_triplets = []
            triplet_match_start = time.time()
            for kb_id in knowledge_base_ids:
                debug(f"调用 Neo4j 三元组匹配: knowledge_base_id={kb_id}")
                try:
                    neo4j_result = await self._make_neo4japi_request("matchtriplets", {
                        "query": query,
                        "query_entities": query_entities,
                        "userid": userid,
                        "knowledge_base_id": kb_id
                    })
                    if neo4j_result.get("status") == "success":
                        triplets = neo4j_result.get("triplets", [])
                        all_triplets.extend(triplets)
                        debug(f"知识库 {kb_id} 匹配到 {len(triplets)} 个三元组: {triplets[:5]}")
                    else:
                        error(
                            f"Neo4j 三元组匹配失败: knowledge_base_id={kb_id}, 错误: {neo4j_result.get('message', '未知错误')}")
                except Exception as e:
                    error(f"Neo4j 三元组匹配失败: knowledge_base_id={kb_id}, 错误: {str(e)}")
                    continue
            timing_stats["triplet_matching"] = time.time() - triplet_match_start
            debug(f"三元组匹配总耗时: {timing_stats['triplet_matching']:.3f} 秒")

            # 拼接三元组文本
            triplet_text_start = time.time()
            triplet_texts = []
            for triplet in all_triplets:
                head = triplet.get('head', '')
                type_ = triplet.get('type', '')
                tail = triplet.get('tail', '')
                if head and type_ and tail:
                    triplet_texts.append(f"{head} {type_} {tail}")
                else:
                    debug(f"无效三元组: {triplet}")
            combined_text = query
            if triplet_texts:
                combined_text += " [三元组] " + "; ".join(triplet_texts)
            debug(
                f"拼接文本: {combined_text[:200]}... (总长度: {len(combined_text)}, 三元组数量: {len(triplet_texts)})")
            timing_stats["triplet_text_combine"] = time.time() - triplet_text_start
            debug(f"拼接三元组文本耗时: {timing_stats['triplet_text_combine']:.3f} 秒")

            # 将拼接文本转换为向量
            vector_start = time.time()
            query_vector = await self._get_embeddings([combined_text])
            if not query_vector or not all(len(vec) == 1024 for vec in query_vector):
                raise ValueError("查询向量必须是长度为 1024 的浮点数列表")
            query_vector = query_vector[0]  # 取第一个向量
            timing_stats["vector_generation"] = time.time() - vector_start
            debug(f"生成查询向量耗时: {timing_stats['vector_generation']:.3f} 秒")

            # 调用融合搜索端点
            search_start = time.time()
            result = await self._make_api_request("searchquery", {
                "query_vector": query_vector.tolist(),
                "userid": userid,
                "knowledge_base_ids": knowledge_base_ids,
                "limit": limit,
                "offset": offset,
                "db_type": db_type
            })
            timing_stats["vector_search"] = time.time() - search_start
            debug(f"向量搜索耗时: {timing_stats['vector_search']:.3f} 秒")

            if result.get("status") != "success":
                error(f"融合搜索失败: {result.get('message', '未知错误')}")
                return {"results": [], "timing": timing_stats}

            unique_results = result.get("results", [])
            if use_rerank and unique_results:
                rerank_start = time.time()
                debug("开始重排序")
                unique_results = await self._rerank_results(combined_text, unique_results, limit)
                unique_results = sorted(unique_results, key=lambda x: x.get('rerank_score', 0), reverse=True)
                timing_stats["reranking"] = time.time() - rerank_start
                debug(f"重排序耗时: {timing_stats['reranking']:.3f} 秒")
                debug(f"重排序分数分布: {[round(r.get('rerank_score', 0), 3) for r in unique_results]}")
            else:
                unique_results = [{k: v for k, v in r.items() if k != 'rerank_score'} for r in unique_results]

            timing_stats["total_time"] = time.time() - start_time
            info(f"融合搜索完成，返回 {len(unique_results)} 条结果，总耗时: {timing_stats['total_time']:.3f} 秒")
            return {"results": unique_results[:limit], "timing": timing_stats}

        except Exception as e:
            error(f"融合搜索失败: {str(e)}, 堆栈: {traceback.format_exc()}")
            return {"results": [], "timing": timing_stats}

    async def _list_user_files(self, userid: str, db_type: str = "") -> Dict[str, List[Dict]]:
        """列出用户的所有知识库及其文件，按 knowledge_base_id 分组"""
        collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
        try:
            info(f"列出用户文件: userid={userid}, db_type={db_type}")

            if not userid:
                raise ValueError("userid 不能为空")
            if "_" in userid or (db_type and "_" in db_type):
                raise ValueError("userid 和 db_type 不能包含下划线")
            if (db_type and len(db_type) > 100) or len(userid) > 100:
                raise ValueError("userid 或 db_type 的长度超出限制")

            # 调用列出用户文件端点
            result = await self._make_api_request("listuserfiles", {
                "userid": userid,
                "db_type": db_type
            })

            if result.get("status") != "success":
                error(f"列出用户文件失败: {result.get('message', '未知错误')}")
                return {}

            return result.get("files_by_knowledge_base", {})

        except Exception as e:
            error(f"列出用户文件失败: {str(e)}")
            return {}

    async def _list_all_knowledge_bases(self, db_type: str = "") -> Dict[str, Any]:
        """列出数据库中所有用户的知识库及其文件，按用户分组"""
        collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
        try:
            info(f"列出所有用户的知识库: db_type={db_type}")

            if db_type and "_" in db_type:
                raise ValueError("db_type 不能包含下划线")
            if db_type and len(db_type) > 100:
                raise ValueError("db_type 的长度应小于 100")

            # 调用列出所有知识库端点
            result = await self._make_api_request("listallknowledgebases", {
                "db_type": db_type
            })

            return result

        except Exception as e:
            error(f"列出所有用户知识库失败: {str(e)}")
            return {
                "status": "error",
                "users_knowledge_bases": {},
                "collection_name": collection_name,
                "message": f"列出所有用户知识库失败: {str(e)}",
                "status_code": 400
            }

connection_register('Rag', MilvusConnection)
info("MilvusConnection registered")