from api_service import APIService
from appPublic.registerfunction import RegisterFunction
from appPublic.log import debug, error, info
from sqlor.dbpools import DBPools
import asyncio
import aiohttp
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
import os
import re
import time
import uuid
from datetime import datetime
import traceback
from filetxt.loader import fileloader
from ahserver.serverenv import get_serverenv

async def get_orgid_by_id(kdb_id):
    """
    根据 kdb 的 id 查询对应的 orgid。
    """
    db = DBPools()
    # f = get_serverenv("get_module_dbname")
    # dbname = f("rag")
    dbname = "kyrag"
    sql = "SELECT orgid FROM kdb WHERE id = ${id}$"
    try:
        async with db.sqlorContext(dbname) as sor:
            result = await sor.sqlExe(sql,{"id":kdb_id})
            print(result)
            if result and len(result) > 0:
                return result[0].get('orgid')
            return None
    except Exception as e:
        error(f"查询 orgid 失败: {str(e)}, 堆栈: {traceback.format_exc()}")
        return None

async def file_uploaded(params_kw):
    """将文档插入 Milvus 并抽取三元组到 Neo4j"""
    debug(f'Received params: {params_kw=}')
    api_service = APIService()
    realpath = params_kw.get('realpath', '')
    fiid = params_kw.get('fiid', '')
    id = params_kw.get('id', '')
    orgid = await get_orgid_by_id(fiid)
    db_type = ''
    debug(f'Inserting document: file_path={realpath}, userid={orgid}, db_type={db_type}, knowledge_base_id={fiid}, document_id={id}')

    timings = {}
    start_total = time.time()

    try:
        if not orgid or not fiid or not id:
            raise ValueError("orgid、fiid 和 id 不能为空")
        if len(orgid) > 32 or len(fiid) > 255:
            raise ValueError("orgid 或 fiid 的长度超出限制")
        if not os.path.exists(realpath):
            raise ValueError(f"文件 {realpath} 不存在")

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

        debug(f"加载文件: {realpath}")
        start_load = time.time()
        text = fileloader(realpath)
        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"文件 {realpath} 加载为空")

        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"文件 {realpath} 未生成任何文档块")

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

        debug("调用嵌入服务生成向量")
        start_embedding = time.time()
        texts = [chunk.page_content for chunk in chunks]
        embeddings = await api_service.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_data = []
        for i, chunk in enumerate(chunks):
            chunks_data.append({
                "userid": orgid,
                "knowledge_base_id": fiid,
                "text": chunk.page_content,
                "vector": embeddings[i],
                "document_id": id,
                "filename": filename + '.' + ext,
                "file_path": realpath,
                "upload_time": upload_time,
                "file_type": ext,
            })

        debug(f"调用插入文件端点: {realpath}")
        start_milvus = time.time()
        result = await api_service.milvus_insert_document(chunks_data, 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": id, "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 = [api_service.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]}")

            debug(f"抽取到 {len(unique_triples)} 个三元组，调用 Neo4j 服务插入")
            start_neo4j = time.time()
            if unique_triples:
                neo4j_result = await api_service.neo4j_insert_triples(unique_triples, id, fiid, orgid)
                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": id, "collection_name": "ragdb", "timings": timings,
                            "message": f"Neo4j 三元组插入失败: {neo4j_result.get('message', '未知错误')}", "status_code": 400}
                info(f"文件 {realpath} 三元组成功插入 Neo4j: {neo4j_result.get('message')}")
            else:
                debug(f"文件 {realpath} 未抽取到三元组")
            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": id, "collection_name": "ragdb", "timings": timings,
                    "unique_triples": unique_triples,
                    "message": f"文件 {realpath} 成功嵌入，但三元组处理或 Neo4j 插入失败: {str(e)}", "status_code": 200}

        timings["total"] = time.time() - start_total
        debug(f"总耗时: {timings['total']:.2f} 秒")
        return {"status": "success", "userid": orgid, "document_id": id, "collection_name": "ragdb", "timings": timings,
                "unique_triples": unique_triples, "message": f"文件 {realpath} 成功嵌入并处理三元组", "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": id, "collection_name": "ragdb", "timings": timings,
                "message": f"插入文档失败: {str(e)}", "status_code": 400}

async def file_deleted(params_kw):
    """删除用户指定文件数据，包括 Milvus 和 Neo4j 中的记录"""
    api_service = APIService()
    id = params_kw.get('id', '')
    realpath = params_kw.get('realpath', '')
    fiid = params_kw.get('fiid', '')
    orgid = await get_orgid_by_id(fiid)
    db_type = ''
    collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
    try:
        required_fields = ['id', 'fiid', 'realpath']
        missing_fields = [field for field in required_fields if field not in params_kw or not params_kw[field]]
        if missing_fields:
            raise ValueError(f"缺少必填字段: {', '.join(missing_fields)}")

        debug(f"调用删除文件端点: userid={orgid}, file_path={realpath}, knowledge_base_id={fiid}")
        milvus_result = await api_service.milvus_delete_document(orgid, realpath, fiid, id, db_type)

        if milvus_result.get("status") != "success":
            raise ValueError(milvus_result.get("message", "Milvus 删除失败"))

        neo4j_deleted_nodes = 0
        neo4j_deleted_rels = 0
        try:
            debug(f"调用 Neo4j 删除文档端点: document_id={id}")
            neo4j_result = await api_service.neo4j_delete_document(id)
            if neo4j_result.get("status") != "success":
                raise ValueError(neo4j_result.get("message", "Neo4j 删除失败"))
            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={id} 的 {nodes_deleted} 个 Neo4j 节点和 {rels_deleted} 个关系")
        except Exception as e:
            error(f"删除 document_id={id} 的 Neo4j 数据失败: {str(e)}")

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

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


async def main():
    dbs = {
		"kyrag":{
			"driver":"aiomysql",
			"async_mode":True,
			"coding":"utf8",
			"maxconn":100,
			"dbname":"kyrag",
			"kwargs":{
				"user":"test",
				"db":"kyrag",
				"password":"QUZVcXg5V1p1STMybG5Ia6mX9D0v7+g=",
				"host":"db"
			}
		}
	}
    DBPools(dbs)
    # 测试 file_uploaded
    print("测试 file_uploaded...")
    test_file_path = "/home/wangmeihua/data/kg.txt"
    test_params_upload = {
        "realpath": test_file_path,
        "fiid": "1",
        "id": "doc1"
    }
    upload_result = await file_uploaded(test_params_upload)
    print(f"file_uploaded 结果: {upload_result}")

    # # 测试 file_deleted
    # test_file_path = "/home/wangmeihua/data/kg.txt"
    # print("测试 file_deleted...")
    # test_params_delete = {
    #     "realpath": test_file_path,
    #     "fiid": "1",
    #     "id": "doc1"
    # }
    # delete_result = await file_deleted(test_params_delete)
    # print(f"file_deleted 结果: {delete_result}")


if __name__ == "__main__":
    asyncio.run(main())
