ax-backend/app/services/document_service.py

"""document_service.py — 文档 CRUD，含 Word 下载解析。"""

import re
import tempfile
from datetime import datetime, timezone
from pathlib import Path

import httpx
from docx import Document as DocxDocument
from docx.oxml.ns import qn
from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.exceptions import DocumentNotFoundError, DocumentParseError
from app.models.document import Document
from app.schemas.document import CreateDocumentRequest, UpdateDocumentRequest


# ------------------------------------------------------------------ #
# Word → Markdown 解析（完整实现，支持往返转换）
# ------------------------------------------------------------------ #

# 标准样式列表（不需要添加样式注释）
STANDARD_STYLES = {
    "Normal",
    "Heading 1", "Heading 2", "Heading 3", "Heading 4", "Heading 5", "Heading 6",
    "List Bullet", "List Bullet 2", "List Bullet 3",
    "List Number", "List Number 2", "List Number 3",
    "Quote", "Quote Char",
    "No Spacing",
    "Table Grid",
}


def _is_standard_style(style_name: str) -> bool:
    """判断是否为标准样式"""
    if not style_name or style_name in STANDARD_STYLES:
        return True
    if style_name.startswith("Heading ") or "List Bullet" in style_name or "List Number" in style_name:
        return True
    return False


def _render_run_with_format(run) -> str:
    """将单个 Run 转换为带格式的 Markdown 文本"""
    text = run.text
    if not text:
        return ""
    
    # 检测代码格式（Courier New 字体）
    is_code = run.font.name == "Courier New"
    
    # 删除线
    has_strike = run.font.strike if run.font.strike is not None else False
    
    # 粗体
    has_bold = run.bold if run.bold is not None else False
    
    # 斜体  
    has_italic = run.italic if run.italic is not None else False
    
    # 应用格式（注意顺序：代码 > 删除线 > 粗体/斜体）
    if is_code:
        text = f"`{text}`"
    elif has_strike:
        text = f"~~{text}~~"
    elif has_bold and has_italic:
        # 粗斜体组合
        text = f"***{text}***"
    elif has_bold:
        text = f"**{text}**"
    elif has_italic:
        text = f"*{text}*"
    
    return text


def _render_para_with_inline_format(para) -> str:
    """将段落的所有 Run 转换为带格式的 Markdown 文本"""
    parts = []
    for run in para.runs:
        formatted_text = _render_run_with_format(run)
        parts.append(formatted_text)
    return "".join(parts)


def _is_code_block(para) -> bool:
    """检测段落是否为代码块（所有 Run 都是 Courier New + No Spacing 样式）"""
    if not para.runs:
        return False
    
    # 检查样式
    style_name = para.style.name if para.style else ""
    if style_name != "No Spacing":
        return False
    
    # 检查所有非空 Run 是否都使用 Courier New
    non_empty_runs = [run for run in para.runs if run.text.strip()]
    if not non_empty_runs:
        return False
    
    return all(run.font.name == "Courier New" for run in non_empty_runs)


def _convert_table_to_markdown(table) -> str:
    """将 Word 表格转换为 Markdown 表格格式"""
    if not table.rows:
        return ""
    
    md_lines = []
    
    # 表头（第一行）
    header_cells = []
    for cell in table.rows[0].cells:
        # 对单元格内容也应用格式解析
        cell_text = []
        for para in cell.paragraphs:
            para_style = para.style.name if para.style else "Normal"
            para_text = _render_para_with_inline_format(para).strip()
            # 移除表头自动添加的粗体格式
            # 因为导出时会自动给表头加粗，但原始Markdown可能没有
            # 这里需要判断：如果所有文本都是粗体，则移除粗体标记
            if para_text.startswith("**") and para_text.endswith("**") and para_text.count("**") == 2:
                para_text = para_text[2:-2]
            if para_text:
                # 如果表头单元格使用自定义样式，添加注释
                if not _is_standard_style(para_style) and para_style not in ("Normal", "Table Grid"):
                    para_text = f"<!-- style: {para_style} -->{para_text}"
                cell_text.append(para_text)
        header_cells.append(" ".join(cell_text))
    
    md_lines.append("| " + " | ".join(header_cells) + " |")
    
    # 分隔行
    md_lines.append("| " + " | ".join(["---"] * len(header_cells)) + " |")
    
    # 数据行
    for row in table.rows[1:]:
        cells = []
        for cell in row.cells[:len(header_cells)]:  # 确保列数一致
            # 对单元格内容也应用格式解析
            cell_text = []
            for para in cell.paragraphs:
                para_style = para.style.name if para.style else "Normal"
                para_text = _render_para_with_inline_format(para).strip()
                
                if para_text:
                    # 如果单元格段落使用自定义样式，添加注释
                    if not _is_standard_style(para_style) and para_style not in ("Normal", "Table Grid"):
                        para_text = f"<!-- style: {para_style} -->{para_text}"
                    cell_text.append(para_text)
            cells.append(" ".join(cell_text))
        md_lines.append("| " + " | ".join(cells) + " |")
    
    return "\n".join(md_lines)


def _get_list_info(para) -> tuple[str, int]:
    """获取列表类型和深度
    返回：(list_type, depth)
    list_type: "bullet" 或 "number" 或 None
    depth: 1, 2, 3... 或 0
    """
    style_name = para.style.name if para.style else ""
    if not style_name:
        return None, 0
    
    # 直接根据样式名判断
    # 注意：python-docx 生成的列表可能使用 "List Bullet 2" 等样式
    if "List Bullet" in style_name:
        # 提取数字
        if "List Bullet 3" in style_name:
            return "bullet", 3
        elif "List Bullet 2" in style_name:
            return "bullet", 2
        else:
            return "bullet", 1
    
    elif "List Number" in style_name:
        if "List Number 3" in style_name:
            return "number", 3
        elif "List Number 2" in style_name:
            return "number", 2
        else:
            return "number", 1
    
    return None, 0


def _has_bottom_border(para) -> bool:
    """检测段落是否有底部边框（用于分隔线）"""
    try:
        pPr = para._p.pPr
        if pPr is None:
            return False
        pBdr = pPr.find(qn('w:pBdr'))
        if pBdr is None:
            return False
        bottom = pBdr.find(qn('w:bottom'))
        return bottom is not None
    except Exception:
        return False


def _docx_to_markdown(path: Path) -> str:
    """将 .docx 文件解析为完整 Markdown 文本（支持往返转换）"""
    doc = DocxDocument(str(path))
    
    # 提取所有图片及其位置信息
    from app.services.image_service import extract_images_from_word
    images = extract_images_from_word(doc)
    # 按段落索引建立映射，方便查找
    image_map = {}
    for img in images:
        para_idx = img['paragraph_index']
        if para_idx not in image_map:
            image_map[para_idx] = []
        image_map[para_idx].append(img)
    
    # 收集所有元素（段落和表格）并按文档顺序排列
    elements = []
    
    # 通过 XML body 遍历，保持原始顺序
    body = doc.element.body
    para_map = {p._element: p for p in doc.paragraphs}
    table_map = {t._element: t for t in doc.tables}
    
    for child in body:
        tag = child.tag
        if tag.endswith('p'):
            para = para_map.get(child)
            if para:
                elements.append(('para', para))
        elif tag.endswith('tbl'):
            table = table_map.get(child)
            if table:
                elements.append(('table', table))
    
    # 转换为 Markdown
    lines = []
    in_code_block = False
    code_block_lines = []
    
    # 记录实际的段落索引到元素索引的映射
    para_idx_in_elements = {}
    actual_para_idx = 0
    for elem_idx, (elem_type, elem) in enumerate(elements):
        if elem_type == "para":
            para_idx_in_elements[actual_para_idx] = elem_idx
            actual_para_idx += 1
    
    for elem_idx, (elem_type, elem) in enumerate(elements):
        if elem_type == "para":
            para = elem
            style_name = para.style.name if para.style else ""
            
            # 检测代码块开始/结束
            if _is_code_block(para):
                if not in_code_block:
                    in_code_block = True
                    code_block_lines = []
                code_block_lines.append(para.text)  # 使用原始文本，不要strip
                continue
            else:
                # 如果之前在代码块中，现在结束了
                if in_code_block:
                    in_code_block = False
                    lines.append("```")
                    # 移除首尾空行，但保留中间的空行
                    while code_block_lines and not code_block_lines[0].strip():
                        code_block_lines.pop(0)
                    while code_block_lines and not code_block_lines[-1].strip():
                        code_block_lines.pop()
                    lines.extend(code_block_lines)
                    lines.append("```")
                    code_block_lines = []
            
            # 带格式的文本
            text = _render_para_with_inline_format(para).strip()
            
            # 空段落
            if not text:
                # 检测是否是分隔线（只有底部边框的空段落）
                if _has_bottom_border(para):
                    lines.append("---")
                
                # 空段落也要检查是否有图片
                current_para_idx = None
                for para_idx, e_idx in para_idx_in_elements.items():
                    if e_idx == elem_idx:
                        current_para_idx = para_idx
                        break
                
                if current_para_idx is not None and current_para_idx in image_map:
                    from app.services.image_service import create_image_markdown
                    for img in image_map[current_para_idx]:
                        img_md = create_image_markdown(
                            img['data_url'],
                            img['style'],
                            alt=f"图片"
                        )
                        lines.append(img_md.strip())
                
                continue
            
            # 标题
            if style_name.startswith("Heading "):
                try:
                    level = int(style_name.split()[-1])
                except ValueError:
                    level = 1
                level = min(max(level, 1), 6)
                lines.append(f"{'#' * level} {text}")
            
            # 列表
            elif (list_info := _get_list_info(para))[0]:
                list_type, depth = list_info
                indent = "  " * (depth - 1)
                if list_type == "bullet":
                    lines.append(f"{indent}- {text}")
                else:  # number
                    lines.append(f"{indent}1. {text}")
            
            # 块引用
            elif style_name == "Quote":
                lines.append(f"> {text}")
            
            # 普通段落
            else:
                # 如果是自定义样式，添加样式注释（和文本在同一行）
                if not _is_standard_style(style_name) and style_name != "Normal":
                    lines.append(f"<!-- style: {style_name} -->{text}")
                else:
                    lines.append(text)
            
            # 段落处理完成后，检查是否有图片
            # 找到当前 elem_idx 对应的实际段落索引
            current_para_idx = None
            for para_idx, e_idx in para_idx_in_elements.items():
                if e_idx == elem_idx:
                    current_para_idx = para_idx
                    break
            
            if current_para_idx is not None and current_para_idx in image_map:
                # 在段落后输出图片
                from app.services.image_service import create_image_markdown
                for img in image_map[current_para_idx]:
                    img_md = create_image_markdown(
                        img['data_url'],
                        img['style'],
                        alt=f"图片"
                    )
                    lines.append(img_md.strip())
        
        elif elem_type == "table":
            # 如果之前在代码块中，先结束
            if in_code_block:
                in_code_block = False
                lines.append("```")
                lines.extend(code_block_lines)
                lines.append("```")
                code_block_lines = []
            
            table = elem
            table_md = _convert_table_to_markdown(table)
            if table_md:
                lines.append(table_md)
    
    # 处理未结束的代码块
    if in_code_block:
        lines.append("```")
        # 移除首尾空行
        while code_block_lines and not code_block_lines[0].strip():
            code_block_lines.pop(0)
        while code_block_lines and not code_block_lines[-1].strip():
            code_block_lines.pop()
        lines.extend(code_block_lines)
        lines.append("```")
    
    # 过滤空字符串，用双换行连接
    filtered_lines = [line for line in lines if line]
    result = "\n\n".join(filtered_lines).strip()
    
    # 修复列表项之间的额外换行
    # 递归替换，直到没有匹配项
    import re
    # 匹配模式：(可选缩进)(列表标记 - 或 1.)(内容)(双换行)(可选缩进)(列表标记)
    pattern = r'((?:^|\n)(  )?(?:-|\d+\.)\s+[^\n]+)\n\n(  )?(?=(?:-|\d+\.)\s+)'
    while True:
        new_result = re.sub(pattern, r'\1\n\3', result, flags=re.MULTILINE)
        if new_result == result:
            break
        result = new_result
    
    # 确保以单个换行符结尾（Markdown 标准）
    return result + "\n" if result else ""


async def _download_and_parse(file_url: str) -> str:
    """下载 Word 文档到临时文件，解析为 Markdown，自动清理临时文件。"""
    suffix = Path(file_url.split("?")[0]).suffix.lower() or ".docx"
    if suffix not in (".doc", ".docx"):
        raise DocumentParseError(f"不支持的文件格式: {suffix}")

    try:
        async with httpx.AsyncClient(timeout=60) as client:
            resp = await client.get(file_url)
            resp.raise_for_status()
    except httpx.HTTPError as exc:
        raise DocumentParseError(f"文件下载失败: {exc}") from exc

    with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
        tmp_path = Path(tmp.name)
        tmp_path.write_bytes(resp.content)

    try:
        content = _docx_to_markdown(tmp_path)
    except Exception as exc:
        raise DocumentParseError(f"Word 解析失败: {exc}") from exc
    finally:
        tmp_path.unlink(missing_ok=True)

    return content


# ------------------------------------------------------------------ #
# 服务类
# ------------------------------------------------------------------ #

class DocumentService:
    def __init__(self, db: AsyncSession) -> None:
        self.db = db

    async def create_document(self, data: CreateDocumentRequest) -> Document:
        content = await _download_and_parse(data.file_url)
        doc = Document(
            content=content,
            format="markdown",
            session_id=data.session_id,
            created_by=data.user_id,
        )
        self.db.add(doc)
        await self.db.commit()
        await self.db.refresh(doc)
        return doc

    async def get_document(self, document_id: str) -> Document:
        result = await self.db.execute(
            select(Document).where(Document.id == document_id)
        )
        doc = result.scalar_one_or_none()
        if doc is None:
            raise DocumentNotFoundError(document_id)
        return doc

    async def list_documents(
        self,
        user_id: str,
        page: int = 1,
        page_size: int = 20,
        session_id: str | None = None,
        sort_by: str = "updated_at",
        sort_order: str = "desc",
    ) -> tuple[list[Document], int]:
        query = select(Document).where(Document.created_by == user_id)

        if session_id:
            query = query.where(Document.session_id == session_id)

        sort_col = getattr(Document, sort_by, Document.updated_at)
        query = query.order_by(sort_col.asc() if sort_order == "asc" else sort_col.desc())

        count_q = select(func.count()).select_from(query.subquery())
        total: int = (await self.db.execute(count_q)).scalar_one()

        offset = (page - 1) * page_size
        result = await self.db.execute(query.offset(offset).limit(page_size))
        return list(result.scalars().all()), total

    async def update_document(self, document_id: str, data: UpdateDocumentRequest) -> Document:
        doc = await self.get_document(document_id)

        if data.content is not None:
            doc.content = data.content
        elif data.blocks is not None:
            doc.content = self._apply_block_updates(doc.content, data.blocks)

        doc.updated_at = datetime.now(timezone.utc)
        await self.db.commit()
        await self.db.refresh(doc)
        return doc

    async def delete_document(self, document_id: str) -> None:
        doc = await self.get_document(document_id)
        await self.db.delete(doc)
        await self.db.commit()

    async def delete_documents_by_session(self, session_id: str) -> int:
        """删除指定 sessionId 的所有文档，返回删除数量。"""
        result = await self.db.execute(
            select(Document).where(Document.session_id == session_id)
        )
        docs = list(result.scalars().all())
        
        for doc in docs:
            await self.db.delete(doc)
        
        await self.db.commit()
        return len(docs)

    @staticmethod
    def _apply_block_updates(content: str, blocks: list) -> str:
        """按 level+index 定位标题块并替换，其余内容不变。"""
        lines = content.split("\n")
        heading_pattern = re.compile(r"^(#{1,6})\s+")

        # 收集标题位置和 index
        level_counter: dict[int, int] = {}
        heading_info: list[tuple[int, int, int]] = []  # (line_idx, level, index)
        for i, line in enumerate(lines):
            m = heading_pattern.match(line)
            if m:
                level = len(m.group(1))
                idx = level_counter.get(level, 0)
                heading_info.append((i, level, idx))
                level_counter[level] = idx + 1

        # 切分为块
        split_points = [hi[0] for hi in heading_info] + [len(lines)]
        pre_end = split_points[0] if split_points else len(lines)
        chunks: list[str] = ["\n".join(lines[:pre_end])]

        key_to_chunk: dict[tuple[int, int], int] = {}
        for i, sp in enumerate(split_points[:-1]):
            end = split_points[i + 1]
            chunks.append("\n".join(lines[sp:end]))
            _, level, index = heading_info[i]
            key_to_chunk[(level, index)] = len(chunks) - 1

        # 替换
        for block in blocks:
            key = (block.level, block.index)
            if key in key_to_chunk:
                chunks[key_to_chunk[key]] = block.content

        return "\n".join(chunks).strip() + "\n"