"""Word 文档解析引擎"""
from io import BytesIO
from typing import Any, Dict, Iterable, List

from docx import Document
from docx.document import Document as DocxDocument
from docx.oxml.table import CT_Tbl
from docx.oxml.text.paragraph import CT_P
from docx.table import Table as DocxTable
from docx.table import _Cell
from docx.text.paragraph import Paragraph


def parse_document(source: str | bytes) -> List[Dict[str, Any]]:
    """
    解析 Word 文档，返回带 block_id、parent_block_id、sort_order 的 block 列表。
    """
    doc = Document(BytesIO(source)) if isinstance(source, bytes) else Document(source)
    blocks: List[Dict[str, Any]] = []

    for item in iter_block_items(doc):
        if isinstance(item, Paragraph):
            block = classify_paragraph(item)
        else:
            block = classify_table(item)

        if block:
            blocks.append(block)

    assign_block_ids(blocks)
    build_tree(blocks)
    return blocks


def iter_block_items(parent: DocxDocument | _Cell) -> Iterable[Paragraph | DocxTable]:
    """按 Word 文档中的实际顺序遍历段落和表格。"""
    parent_element = parent.element.body if isinstance(parent, DocxDocument) else parent._tc

    for child in parent_element.iterchildren():
        if isinstance(child, CT_P):
            yield Paragraph(child, parent)
        elif isinstance(child, CT_Tbl):
            yield DocxTable(child, parent)


def classify_paragraph(para: Paragraph) -> Dict[str, Any] | None:
    """识别段落类型"""
    text = para.text.strip()
    style_name = para.style.name if para.style else ""

    if not text:
        return None  # 跳过空段落

    if is_heading_style(style_name):
        level = parse_heading_level(style_name)
        return {"type": "heading", "text": text, "level": level, "style": style_name}

    return {"type": "paragraph", "text": text, "level": 0, "style": style_name}


def is_heading_style(style_name: str) -> bool:
    return style_name.startswith("Heading") or style_name.startswith("标题")


def parse_heading_level(style_name: str) -> int:
    parts = style_name.replace("Heading", "").replace("标题", "").strip()
    return int(parts) if parts.isdigit() and 1 <= int(parts) <= 6 else 1


def classify_table(table: DocxTable) -> Dict[str, Any] | None:
    """识别表格类型并提取行列数。"""
    rows = len(table.rows)
    cols = len(table.columns)
    if rows == 0 or cols == 0:
        return None

    first_row = [cell.text.strip() for cell in table.rows[0].cells]
    text = " | ".join(cell for cell in first_row if cell)
    if not text:
        text = f"表格（{rows} 行 x {cols} 列）"

    return {
        "type": "table",
        "text": text,
        "level": 0,
        "rows": rows,
        "cols": cols,
        "headers": first_row,
    }


def generate_block_id(index: int) -> str:
    """生成 block_id，格式：block_001"""
    return f"block_{index:03d}"


def assign_block_ids(blocks: List[Dict[str, Any]]) -> None:
    for index, block in enumerate(blocks, start=1):
        block["block_id"] = generate_block_id(index)
        block["sort_order"] = index
        block["parent_block_id"] = None


def build_tree(blocks: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    """构建层级树结构"""
    tree: List[Dict[str, Any]] = []
    heading_stack: List[Dict[str, Any]] = []

    for block in blocks:
        node = {
            "block_id": block["block_id"],
            "block_type": block["type"],
            "text": block["text"][:50],
            "level": block.get("level", 0),
            "children": [],
        }

        if block["type"] == "heading":
            level = block.get("level", 1)

            while heading_stack and heading_stack[-1]["level"] >= level:
                heading_stack.pop()

            if heading_stack:
                parent = heading_stack[-1]["node"]
                parent["children"].append(node)
                block["parent_block_id"] = parent["block_id"]
            else:
                tree.append(node)

            heading_stack.append({"level": level, "node": node})
        elif heading_stack:
            parent = heading_stack[-1]["node"]
            parent["children"].append(node)
            block["parent_block_id"] = parent["block_id"]
        else:
            tree.append(node)

    return tree