# -*- coding: utf-8 -*-
# doc_split.py
# @author: lyg
# @date: 2025-04-24
# @version: 0.0.1
# @description: 文档处理，按页切割文档并提取每一页的文本，提文档的章节信息，将章节信息和页码信息关联起来，提取每一页的实体词。

from dataclasses import dataclass, field
import docx
import pymupdf

import pymupdf4llm

from knowledgebase.doc.entity_recognition import EntityRecognition


@dataclass
class ChapterInfo:
    """
    表示章节信息的类，包含章节标题、起始页码、结束页码和层级。
    """
    title: str  # 章节标题
    start_page: int  # 起始页码
    end_page: int  # 结束页码
    level: int  # 章节层级

    def __post_init__(self):
        """
        在初始化后执行的钩子函数，用于验证输入参数的有效性。
        """
        # 验证标题是否为空
        if not self.title or not isinstance(self.title, str):
            raise ValueError("标题必须是非空字符串")

        # 验证页码是否为正整数
        if not isinstance(self.start_page, int) or self.start_page < 1:
            raise ValueError("起始页码必须是大于0的整数")
        if not isinstance(self.end_page, int) or self.end_page < 1:
            raise ValueError("结束页码必须是大于0的整数")

        # 验证起始页码是否小于等于结束页码
        if self.start_page > self.end_page:
            raise ValueError("起始页码不能大于结束页码")

        # 验证层级是否为正整数
        if not isinstance(self.level, int) or self.level < 1:
            raise ValueError("层级必须是大于0的整数")


@dataclass
class PageInfo:
    """
    页面信息类，用于存储页面的基本信息。

    参数：
    - page_num (int): 页面编号，必须为正整数。
    - chapter_info (str): 章节信息，描述当前页面所属章节。
    - text (str, optional): 页面文本内容，默认为空字符串。
    - entities (list[str]): 页面实体词列表。

    异常：
    - ValueError: 如果 `page_num` 不是正整数，或者 `chapter_info` 和 `text` 不是字符串。
    """
    page_num: int
    chapter_info: str
    text: str = ""
    entities: list[str] = field(default_factory=list)

    def __post_init__(self):
        """
        初始化后验证参数类型和值。
        """
        if not isinstance(self.page_num, int) or self.page_num < 0:
            raise ValueError("page_num 必须是正整数")
        if not isinstance(self.chapter_info, str):
            raise ValueError("chapter_info 必须是字符串")
        if not isinstance(self.text, str):
            raise ValueError("text 必须是字符串")
        if self.entities is None or not isinstance(self.entities, list):
            raise ValueError("entities 必须是列表")


class DocSplit:
    """
    文档处理
    按页切割文档并提取每一页都文本，使用
    """
    pdf_file: str
    page_infos: list[PageInfo]
    chapter_infos: list[ChapterInfo]

    def __init__(self, pdf_file):
        self.pdf_file = pdf_file
        self.page_infos = []
        self.chapter_infos = []
        self.doc = pymupdf.open(self.pdf_file)
        self.extract_chapter_info()
        self.extract_page_info()

    def extract_chapter_info(self):
        """
        提取章节信息
        """
        toc = self.doc.get_toc()
        for item in toc:
            idx = toc.index(item)
            end = len(self.doc)
            for i in range(idx + 1, len(toc)):
                if toc[i][0] >= item[0]:
                    end = toc[i][2]
                    break
            self.chapter_infos.append(ChapterInfo(title=item[1], start_page=item[2], end_page=end, level=item[0]))

    def extract_page_info(self):
        """
        提取页面信息
        """
        for page in self.doc:
            # 通过pymupdf4llm获取页面文本，markdown格式
            page_text = pymupdf4llm.to_markdown(self.doc, pages=[page.number])
            # 创建页面信息对象
            page_info = PageInfo(page_num=page.number, chapter_info="", text=page_text)
            self.page_infos.append(page_info)

    def get_page_info(self, page_num):
        """
        获取指定页码的页面信息
        """
        for page in self.page_infos:
            if page.page_num == page_num:
                return page
        return None


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