from appPublic.log import debug, error
from typing import Dict, Any, List
import uuid
from ahserver.serverenv import ServerEnv
from uapi.init import load_uapi

load_uapi()

class APIService:
    """处理 API 请求的服务类"""

    # 嵌入服务 (BAAI/bge-m3)
    async def get_embeddings(self, request, texts: list, upappid: str, apiname: str, user: str) -> list:
        """调用嵌入服务获取文本向量"""
        try:
            uapi = UAPI(request, DictObject(**globals()))
            params_kw = {"input": texts}
            d = await uapi.request(upappid, apiname, user, params_kw)
            if d.get("object") != "list" or not d.get("data"):
                error(f"嵌入服务响应格式错误: {d}")
                raise RuntimeError("嵌入服务响应格式错误")
            embeddings = [item["embedding"] for item in d["data"]]
            debug(f"成功获取 {len(embeddings)} 个嵌入向量")
            return embeddings
        except Exception as e:
            error(f"嵌入服务调用失败: {str(e)}")
            raise RuntimeError(f"嵌入服务调用失败: {str(e)}")

    # 实体提取服务 (LTP/small)
    async def extract_entities(self, request, query: str, upappid: str, apiname: str, user: str) -> list:
        """调用实体识别服务"""
        try:
            if not query:
                raise ValueError("查询文本不能为空")
            uapi = UAPI(request, DictObject(**globals()))
            params_kw = {"query": query}
            d = await uapi.request(upappid, apiname, user, params_kw)
            if d.get("object") != "list" or not d.get("data"):
                error(f"实体识别服务响应格式错误: {d}")
                raise RuntimeError("实体识别服务响应格式错误")
            entities = d["data"]
            unique_entities = list(dict.fromkeys(entities))
            debug(f"成功提取 {len(unique_entities)} 个唯一实体")
            return unique_entities
        except Exception as e:
            error(f"实体识别服务调用失败: {str(e)}")
            return []

    # 三元组抽取服务 (Babelscape/mrebel-large)
    async def extract_triples(self, request, text: str, upappid: str, apiname: str, user: str) -> list:
        """调用三元组抽取服务"""
        request_id = str(uuid.uuid4())
        debug(f"Request #{request_id} started for triples extraction")
        try:
            uapi = UAPI(request, DictObject(**globals()))
            params_kw = {"text": text}
            d = await uapi.request(upappid, apiname, user, params_kw)
            if d.get("object") != "list":
                error(f"Request #{request_id} invalid response format: {d}")
                raise RuntimeError("三元组抽取服务响应格式错误")
            triples = d["data"]
            debug(f"Request #{request_id} extracted {len(triples)} triples")
            return triples
        except Exception as e:
            error(f"Request #{request_id} failed to extract triples: {str(e)}")
            raise RuntimeError(f"三元组抽取服务调用失败: {str(e)}")

    # 重排序服务 (BAAI/bge-reranker-v2-m3)
    async def rerank_results(self, request, query: str, results: list, top_n: int, upappid: str, apiname: str, user: str) -> list:
        """调用重排序服务"""
        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))

            documents = [result.get("text", str(result)) for result in results]
            uapi = UAPI(request, DictObject(**globals()))
            params_kw = {
                "model": "rerank-001",
                "query": query,
                "documents": documents,
                "top_n": top_n
            }
            d = await uapi.request(upappid, apiname, user, params_kw)
            if d.get("object") != "rerank.result" or not d.get("data"):
                error(f"重排序服务响应格式错误: {d}")
                raise RuntimeError("重排序服务响应格式错误")
            rerank_data = d["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

    # Neo4j 服务
    async def neo4j_docs(self, request, upappid: str, apiname: str, user: str) -> str:
        """获取 Neo4j 文档（返回文本格式）"""
        try:
            uapi = UAPI(request, DictObject(**globals()))
            params_kw = {}
            d = await uapi.request(upappid, apiname, user, params_kw)
            if d.get("status") != 200:
                error(f"Neo4j 文档请求失败，状态码: {d.get('status')}")
                raise RuntimeError(f"Neo4j 文档请求失败: {d.get('status')}")
            text = d.get("text")
            debug(f"Neo4j 文档内容: {text}")
            return text
        except Exception as e:
            error(f"Neo4j 文档请求失败: {str(e)}")
            raise RuntimeError(f"Neo4j 文档请求失败: {str(e)}")

    async def neo4j_initialize(self, request, upappid: str, apiname: str, user: str) -> Dict[str, Any]:
        """初始化 Neo4j 服务"""
        uapi = UAPI(request, DictObject(**globals()))
        params_kw = {}
        return await uapi.request(upappid, apiname, user, params_kw)

    async def neo4j_insert_triples(self, request, triples: list, document_id: str, knowledge_base_id: str, userid: str, upappid: str, apiname: str, user: str) -> Dict[str, Any]:
        """插入三元组到 Neo4j"""
        uapi = UAPI(request, DictObject(**globals()))
        params_kw = {
            "triples": triples,
            "document_id": document_id,
            "knowledge_base_id": knowledge_base_id,
            "userid": userid
        }
        return await uapi.request(upappid, apiname, user, params_kw)

    async def neo4j_delete_document(self, request, document_id: str, upappid: str, apiname: str, user: str) -> Dict[str, Any]:
        """删除指定文档"""
        uapi = UAPI(request, DictObject(**globals()))
        params_kw = {"document_id": document_id}
        return await uapi.request(upappid, apiname, user, params_kw)

    async def neo4j_delete_knowledgebase(self, request, userid: str, knowledge_base_id: str, upappid: str, apiname: str, user: str) -> Dict[str, Any]:
        """删除用户知识库"""
        uapi = UAPI(request, DictObject(**globals()))
        params_kw = {"userid": userid, "knowledge_base_id": knowledge_base_id}
        return await uapi.request(upappid, apiname, user, params_kw)

    async def neo4j_match_triplets(self, request, query: str, query_entities: list, userid: str, knowledge_base_id: str, upappid: str, apiname: str, user: str) -> Dict[str, Any]:
        """根据实体匹配相关三元组"""
        uapi = UAPI(request, DictObject(**globals()))
        params_kw = {
            "query": query,
            "query_entities": query_entities,
            "userid": userid,
            "knowledge_base_id": knowledge_base_id
        }
        return await uapi.request(upappid, apiname, user, params_kw)

    # Milvus 服务
    async def milvus_create_collection(self, request, upappid: str, apiname: str, user: str, db_type: str = "") -> Dict[str, Any]:
        """创建 Milvus 集合"""
        uapi = UAPI(request, DictObject(**globals()))
        params_kw = {"db_type": db_type}
        return await uapi.request(upappid, apiname, user, params_kw)

    async def milvus_delete_collection(self, request, upappid: str, apiname: str, user: str, db_type: str = "") -> Dict[str, Any]:
        """删除 Milvus 集合"""
        uapi = UAPI(request, DictObject(**globals()))
        params_kw = {"db_type": db_type}
        return await uapi.request(upappid, apiname, user, params_kw)

    async def milvus_insert_document(self, request, chunks: List[Dict], upappid: str, apiname: str, user: str, db_type: str = "") -> Dict[str, Any]:
        """添加 Milvus 记录"""
        uapi = UAPI(request, DictObject(**globals()))
        params_kw = {
            "chunks": chunks,
            "dbtype": db_type
        }
        payload = json.dumps(params_kw)  # 转换为 JSON 字符串
        payload_bytes = payload.encode()  # 编码为字节
        payload_size = len(payload_bytes)  # 获取字节数
        debug(f"Request payload size for insertdocument: {payload_size} bytes")
        return await uapi.request(upappid, apiname, user, params_kw)

    async def milvus_delete_document(self, request, userid: str, file_path: str, knowledge_base_id: str, document_id: str, upappid: str, apiname: str, user: str, db_type: str = "") -> Dict[str, Any]:
        """删除 Milvus 记录"""
        uapi = UAPI(request, DictObject(**globals()))
        params_kw = {
            "userid": userid,
            "file_path": file_path,
            "knowledge_base_id": knowledge_base_id,
            "document_id": document_id,
            "dbtype": db_type
        }
        return await uapi.request(upappid, apiname, user, params_kw)

    async def milvus_delete_knowledgebase(self, request, userid: str, knowledge_base_id: str, upappid: str, apiname: str, user: str) -> Dict[str, Any]:
        """删除 Milvus 知识库"""
        uapi = UAPI(request, DictObject(**globals()))
        params_kw = {"userid": userid, "knowledge_base_id": knowledge_base_id}
        return await uapi.request(upappid, apiname, user, params_kw)

    async def milvus_search_query(self, request, query_vector: List[float], userid: str, knowledge_base_ids: list, limit: int, offset: int, upappid: str, apiname: str, user: str) -> Dict[str, Any]:
        """根据用户知识库检索 Milvus"""
        uapi = UAPI(request, DictObject(**globals()))
        params_kw = {
            "query_vector": query_vector,
            "userid": userid,
            "knowledge_base_ids": knowledge_base_ids,
            "limit": limit,
            "offset": offset
        }
        return await uapi.request(upappid, apiname, user, params_kw)

    async def milvus_list_user_files(self, request, userid: str, upappid: str, apiname: str, user: str) -> Dict[str, Any]:
        """列出 Milvus 用户知识库列表"""
        uapi = UAPI(request, DictObject(**globals()))
        params_kw = {"userid": userid}
        return await uapi.request(upappid, apiname, user, params_kw)

    async def milvus_list_all_knowledgebases(self, request, upappid: str, apiname: str, user: str) -> Dict[str, Any]:
        """列出 Milvus 数据库中所有数据"""
        uapi = UAPI(request, DictObject(**globals()))
        params_kw = {}
        return await uapi.request(upappid, apiname, user, params_kw)

    # RAG 服务
    async def rag_create_collection(self, request, upappid: str, apiname: str, user: str, db_type: str = "") -> Dict[str, Any]:
        """创建 RAG 集合"""
        uapi = UAPI(request, DictObject(**globals()))
        params_kw = {"db_type": db_type}
        return await uapi.request(upappid, apiname, user, params_kw)

    async def rag_delete_collection(self, request, upappid: str, apiname: str, user: str, db_type: str = "") -> Dict[str, Any]:
        """删除 RAG 集合"""
        uapi = UAPI(request, DictObject(**globals()))
        params_kw = {"db_type": db_type}
        return await uapi.request(upappid, apiname, user, params_kw)

    async def rag_insert_file(self, request, file_path: str, userid: str, knowledge_base_id: str, document_id: str, upappid: str, apiname: str, user: str) -> Dict[str, Any]:
        """添加 RAG 记录"""
        uapi = UAPI(request, DictObject(**globals()))
        params_kw = {
            "file_path": file_path,
            "userid": userid,
            "knowledge_base_id": knowledge_base_id,
            "document_id": document_id
        }
        return await uapi.request(upappid, apiname, user, params_kw)

    async def rag_delete_file(self, request, userid: str, file_path: str, knowledge_base_id: str, document_id: str, upappid: str, apiname: str, user: str) -> Dict[str, Any]:
        """删除 RAG 记录"""
        uapi = UAPI(request, DictObject(**globals()))
        params_kw = {
            "userid": userid,
            "file_path": file_path,
            "knowledge_base_id": knowledge_base_id,
            "document_id": document_id
        }
        return await uapi.request(upappid, apiname, user, params_kw)

    async def rag_delete_knowledgebase(self, request, userid: str, knowledge_base_id: str, upappid: str, apiname: str, user: str) -> Dict[str, Any]:
        """删除 RAG 知识库"""
        uapi = UAPI(request, DictObject(**globals()))
        params_kw = {"userid": userid, "knowledge_base_id": knowledge_base_id}
        return await uapi.request(upappid, apiname, user, params_kw)

    async def rag_search_query(self, request, query: str, userid: str, knowledge_base_ids: list, limit: int, offset: int, use_rerank: bool, upappid: str, apiname: str, user: str) -> Dict[str, Any]:
        """根据用户知识库检索 RAG"""
        uapi = UAPI(request, DictObject(**globals()))
        params_kw = {
            "query": query,
            "userid": userid,
            "knowledge_base_ids": knowledge_base_ids,
            "limit": limit,
            "offset": offset,
            "use_rerank": use_rerank
        }
        return await uapi.request(upappid, apiname, user, params_kw)

    async def rag_fused_search_query(self, request, query: str, userid: str, knowledge_base_ids: list, limit: int, offset: int, use_rerank: bool, upappid: str, apiname: str, user: str) -> Dict[str, Any]:
        """根据用户知识库+知识图谱检索 RAG"""
        uapi = UAPI(request, DictObject(**globals()))
        params_kw = {
            "query": query,
            "userid": userid,
            "knowledge_base_ids": knowledge_base_ids,
            "limit": limit,
            "offset": offset,
            "use_rerank": use_rerank
        }
        return await uapi.request(upappid, apiname, user, params_kw)

    async def rag_list_user_files(self, request, userid: str, upappid: str, apiname: str, user: str) -> Dict[str, Any]:
        """列出 RAG 用户知识库列表"""
        uapi = UAPI(request, DictObject(**globals()))
        params_kw = {"userid": userid}
        return await uapi.request(upappid, apiname, user, params_kw)

    async def rag_list_all_knowledgebases(self, request, upappid: str, apiname: str, user: str) -> Dict[str, Any]:
        """列出 RAG 数据库中所有数据"""
        uapi = UAPI(request, DictObject(**globals()))
        params_kw = {}
        return await uapi.request(upappid, apiname, user, params_kw)

    async def rag_docs(self, request, upappid: str, apiname: str, user: str) -> Dict[str, Any]:
        """获取 RAG 帮助文档"""
        uapi = UAPI(request, DictObject(**globals()))
        params_kw = {}
        return await uapi.request(upappid, apiname, user, params_kw)