# -*- coding: utf-8 -*-
# @file: doc_processor.py
# @author: lyg
# @date: 2025-5-13
# @version: 
# @description: 处理文档，拆分文档，将拆分后的章节保存到数据库中。
from langchain_core.messages import HumanMessage

from knowledgebase.doc.docx_split import DocSplit
import asyncio
from knowledgebase.db.doc_db_helper import doc_dbh
from knowledgebase.doc.entity_helper import entity_helper
from knowledgebase.doc.entity_recognition import EntityRecognition
import os.path

from knowledgebase.doc.models import DocInfo, ParagraphInfo
from knowledgebase.llm import llm
from knowledgebase.log import Log
from knowledgebase import utils


class DocProcessor:
    def __init__(self, docx_file: str):
        """
        文档处理
        :param docx_file: 要处理的文档
        """
        Log.info(f'开始处理文档：{docx_file}')
        self.docx_file = docx_file
        self.doc_split = DocSplit(docx_file)
        self.doc_type = self.get_doc_type()
        self.entity_recognition = EntityRecognition(self.doc_type)
        self.doc_id = 0

    def get_doc_type(self):
        Log.info(f'识别文档类型：{self.docx_file}')
        rules = '；\n'.join([f'- {it}：{entity_helper.doc_prompt_map[it]}' for it in entity_helper.doc_prompt_map.keys()])
        msg = HumanMessage(f'''
# 指令
请从下面的文件名中识别文档类型，如果识别失败不要输出任何字符。
文件名：{os.path.basename(self.docx_file)}
# 识别规则
{rules}
# 示例
遥测大纲
''')
        resp = llm.invoke([msg])
        Log.info(f'识别结果：{resp.content}')
        return resp.content

    async def gen_sect_entities(self, paragraph: ParagraphInfo):
        # Log.info(f'生成章节实体词：{paragraph.full_text}')
        # 获取章节实体词
        entities = await asyncio.to_thread(lambda: self.entity_recognition.run(paragraph.full_text))
        Log.info(f'章节实体词：{entities}')
        if entities:
            paragraph.entities = [next(filter(lambda x: x.name == e, entity_helper.entities), None) for e in entities]
            paragraph.entities = [e for e in paragraph.entities if e]

    def process(self):
        self.doc_split.split()
        # 分批并发处理，每批10个
        batch_size = 10
        for i in range(0, len(self.doc_split.paragraphs), batch_size):
            batch_paragraphs = self.doc_split.paragraphs[i:i + batch_size]
            tasks = []
            for paragraph in batch_paragraphs:
                tasks.append(self.gen_sect_entities(paragraph))

            async def run():
                await asyncio.gather(*tasks)

            asyncio.run(run())
        # 保存到数据库
        self.save_to_db()

    def save_to_db(self):
        """
        保存段落和段落实体词关系到数据库。
        """
        Log.info('保存段落和段落实体词关系到数据库...')
        with open(self.docx_file, 'rb') as f:
            file_bytes = f.read()
            md5 = utils.generate_bytes_md5(file_bytes)
        doc = DocInfo(os.path.basename(self.docx_file), md5, self.doc_type, self.doc_split.paragraph_tree)
        self.doc_id = doc_dbh.add_doc(doc)
        for paragraph in doc.paragraphs:
            doc_dbh.add_paragraph(self.doc_id, None, paragraph)
        Log.info('保存段落和段落实体词关系到数据库完成')


if __name__ == '__main__':
    files = [
        r"D:\workspace\PythonProjects\KnowledgeBase\doc\XA-5D无人机1553B总线传输通信帧分配（公开）.docx",
        r"D:\workspace\PythonProjects\KnowledgeBase\doc\XA-5D无人机分系统遥测源包设计报告（公开）.docx",
        r"D:\workspace\PythonProjects\KnowledgeBase\doc\XA-5D无人机软件用户需求（公开）.docx",
        r"D:\workspace\PythonProjects\KnowledgeBase\doc\XA-5D无人机遥测大纲（公开）.docx",
        r"D:\workspace\PythonProjects\KnowledgeBase\doc\XA-5D无人机遥测信号分配表（公开）.docx",
        # r"D:\workspace\PythonProjects\KnowledgeBase\doc\XA-5D无人机指令格式与编码定义（公开）.docx",
        r"D:\workspace\PythonProjects\KnowledgeBase\doc\指令格式(公开).docx"
    ]
    for file in files:
        doc_processor = DocProcessor(file)
        doc_processor.process()

    # doc_dbh.get_docs()