# -*- coding: utf-8 -*-
# @file: doc_processor.py
# @author: lyg
# @date: 20250427
# @version: 
# @description: 处理文档，提取章节信息，提取页码信息，提取实体词，写入图数据库（neo4j）。
from knowledgebase.db.neo4j import Neo4jHelper
from knowledgebase.doc.doc_split import DocSplit
from knowledgebase.doc.entity_recognition import EntityRecognition
import asyncio


class DocProcessor:
    def __init__(self, pdf_file):
        self.doc_split = DocSplit(pdf_file)
        self.entity_recognition = EntityRecognition()
        self.neo4j = Neo4jHelper()

    async def gen_page_entities(self, page_info):
        # 获取页面实体词
        page_entities = await asyncio.to_thread(lambda: self.entity_recognition.run(page_info.text))
        page_info.entities = page_entities

    def process(self):
        # 分批并发处理，每批10页
        batch_size = 10
        for i in range(0, len(self.doc_split.page_infos), batch_size):
            batch_page_infos = self.doc_split.page_infos[i:i + batch_size]
            tasks = []
            for page_info in batch_page_infos:
                tasks.append(self.gen_page_entities(page_info))
            asyncio.run(asyncio.gather(*tasks))
        self.save_to_neo4j()

    def save_to_neo4j(self):
        """
        保存页和页实体词到neo4j数据库。

        1.每一页为一个Node；
        2.每一个实体词为一个Node；
        3.页和实体词直接建立关系 页->实体词
        :return:
        """
        for page_info in self.doc_split.page_infos:
            # 创建页节点
            page_node = self.neo4j.create_page_node(page_info)
            entity_nodes = []
            for entity in page_info.entities:
                # 创建实体词节点
                entity_node = self.neo4j.create_entity_node(entity)
                # 建立关系 页->实体词
                self.neo4j.create_page_entity_relationship(page_node, entity_node)
                entity_nodes.append(entity_node)
            if len(entity_nodes) > 0:
                for i in range(len(entity_nodes)):
                    prev_entity_node = entity_nodes[i]
                    for entity_node in entity_nodes[i + 1:]:
                        # 建立关系 一页中的 实体词1->实体词2
                        self.neo4j.create_entity_relationship(prev_entity_node, entity_node)


if __name__ == '__main__':
    pdf_file = "D:/workspace/PythonProjects/KnowledgeBase/doc/XA-5D无人机探测大纲（公开）111.pdf"
    doc_processor = DocProcessor(pdf_file)
    doc_processor.process()