ForcePilot/backend/package/yuxi/knowledge/indexing.py

import asyncio
import base64
import os
import re
import time
from pathlib import Path

import aiofiles
from docling.datamodel.base_models import InputFormat
from docling.document_converter import DocumentConverter
from langchain_community.document_loaders import (
    CSVLoader,
    JSONLoader,
    PyPDFLoader,
    TextLoader,
    UnstructuredHTMLLoader,
    UnstructuredMarkdownLoader,
    UnstructuredWordDocumentLoader,
)
from langchain_text_splitters import RecursiveCharacterTextSplitter
from markdownify import markdownify as md_convert

from yuxi.plugins.parser.zip_utils import process_zip_file as _process_zip_file
from yuxi.storage.minio import get_minio_client
from yuxi.utils import logger

SUPPORTED_FILE_EXTENSIONS: tuple[str, ...] = (
    ".txt",
    ".md",
    ".docx",
    ".html",
    ".htm",
    ".json",
    ".csv",
    ".xls",
    ".xlsx",
    ".pdf",
    ".pptx",
    ".jpg",
    ".jpeg",
    ".png",
    ".bmp",
    ".tiff",
    ".tif",
    ".zip",
)


def is_supported_file_extension(file_name: str | os.PathLike[str]) -> bool:
    """Check whether the given file path has a supported extension."""
    return Path(file_name).suffix.lower() in SUPPORTED_FILE_EXTENSIONS


# Docling 文档转换器（单例模式）
_docling_converter: DocumentConverter | None = None


def _get_docling_converter() -> DocumentConverter:
    """获取 Docling 文档转换器单例"""
    global _docling_converter
    if _docling_converter is None:
        _docling_converter = DocumentConverter(
            format_options={
                InputFormat.DOCX: None,
                InputFormat.XLSX: None,
                InputFormat.PPTX: None,
            }
        )
    return _docling_converter


def _resolve_image_storage_params(params: dict | None) -> tuple[str, str]:
    params = params or {}

    image_bucket = params.get("image_bucket") or "public"
    image_prefix = params.get("image_prefix")
    if image_prefix:
        normalized_prefix = str(image_prefix).strip("/")
        if normalized_prefix:
            return image_bucket, normalized_prefix

    db_id = params.get("db_id")
    if db_id:
        return image_bucket, f"{db_id}/kb-images"

    return image_bucket, "unknown/kb-images"


def _upload_image_to_minio(image_data: bytes, filename: str, bucket_name: str, object_prefix: str) -> str:
    """上传图片到 MinIO，返回 URL"""
    minio_client = get_minio_client()
    minio_client.ensure_bucket_exists(bucket_name)

    normalized_prefix = object_prefix.strip("/") or "unknown/kb-images"
    timestamp = int(time.time() * 1000000)
    suffix = Path(filename).suffix.lower()
    object_name = f"{normalized_prefix}/{timestamp}_{Path(filename).name}"
    content_type_map = {
        ".jpg": "image/jpeg",
        ".jpeg": "image/jpeg",
        ".png": "image/png",
        ".gif": "image/gif",
        ".webp": "image/webp",
        ".bmp": "image/bmp",
        ".tif": "image/tiff",
        ".tiff": "image/tiff",
    }
    content_type = content_type_map.get(suffix, "image/jpeg")

    result = minio_client.upload_file(
        bucket_name=bucket_name,
        object_name=object_name,
        data=image_data,
        content_type=content_type,
    )
    return result.url


def _parse_data_uri(data_uri: str) -> tuple[bytes, str]:
    """解析 data URI，返回 (image_data, mime_type)"""
    header, base64_data = data_uri.split(",", 1)
    mime_type = header.split(":")[1].split(";")[0]
    image_data = base64.b64decode(base64_data)
    return image_data, mime_type


def _convert_with_docling(file_path: Path, params: dict | None = None) -> str:
    """
    使用 Docling 将 docx/xlsx/pptx 转换为 Markdown

    Args:
        file_path: 文件路径
        params: 参数，可包含 image_bucket/image_prefix

    Returns:
        Markdown 字符串
    """
    params = params or {}
    image_bucket, image_prefix = _resolve_image_storage_params(params)

    converter = _get_docling_converter()
    result = converter.convert(file_path)

    if result.status.name != "SUCCESS":
        raise RuntimeError(f"Docling 转换失败: {result.status}")

    doc = result.document

    # 提取图片并上传到 MinIO
    if hasattr(doc, "pictures") and doc.pictures:
        image_refs: list[tuple[str, bytes]] = []

        for pic in doc.pictures:
            if hasattr(pic, "image") and hasattr(pic.image, "uri"):
                uri = str(pic.image.uri)
                if uri.startswith("data:"):
                    image_data, mime_type = _parse_data_uri(uri)
                    timestamp = int(time.time() * 1000000)  # 微秒级时间戳
                    filename = f"image_{timestamp}.{mime_type.split('/')[-1]}"
                    image_refs.append((filename, image_data))

        # 上传图片并收集 URL
        image_urls: list[str] = []
        for filename, image_data in image_refs:
            try:
                url = _upload_image_to_minio(image_data, filename, image_bucket, image_prefix)
                image_urls.append(f"![{filename}]({url})")
            except Exception as e:
                logger.error(f"上传图片失败 {filename}: {e}")
                image_urls.append(f"[图片: {filename}]")

        # 导出 Markdown
        markdown = doc.export_to_markdown()

        # 替换 <!-- image --> 占位符为图片 URL
        # Docling 使用 <!-- image --> 作为占位符
        for url in reversed(image_urls):
            markdown = re.sub(r"<!--\s*image\s*-->", url, markdown, count=1)

        return markdown

    # 无图片时直接导出
    return doc.export_to_markdown()


def _convert_docx_with_python_docx(file_path: Path) -> str:
    """使用 python-docx 解析 DOCX（Docling 失败时兜底）"""
    from docx import Document

    document = Document(str(file_path))
    blocks: list[str] = []

    for para in document.paragraphs:
        text = para.text.strip()
        if text:
            blocks.append(text)

    for table in document.tables:
        rows: list[list[str]] = []
        for row in table.rows:
            cells = [cell.text.strip().replace("\n", " ") for cell in row.cells]
            if any(cells):
                rows.append(cells)

        if not rows:
            continue

        header = rows[0]
        blocks.append(f"| {' | '.join(header)} |")
        blocks.append(f"| {' | '.join(['---'] * len(header))} |")

        for row in rows[1:]:
            normalized_row = row + [""] * (len(header) - len(row))
            blocks.append(f"| {' | '.join(normalized_row[: len(header)])} |")

        blocks.append("")

    return "\n\n".join(blocks).strip()


def chunk_with_parser(file_path, params=None):
    """
    使用文件解析器将文件切分成固定大小的块

    Args:
        file_path: 文件路径
        params: 参数
    """
    params = params or {}
    chunk_size = int(params.get("chunk_size", 500))
    chunk_overlap = int(params.get("chunk_overlap", 100))

    file_type = Path(file_path).suffix.lower()

    # 选择合适的加载器
    if file_type in [".txt"]:
        loader = TextLoader(file_path)

    elif file_type in [".md"]:
        loader = UnstructuredMarkdownLoader(file_path)

    elif file_type in [".docx", ".doc"]:
        loader = UnstructuredWordDocumentLoader(file_path)

    elif file_type in [".html", ".htm"]:
        loader = UnstructuredHTMLLoader(file_path)

    elif file_type in [".json"]:
        loader = JSONLoader(file_path, jq_schema=".")

    elif file_type in [".csv"]:
        loader = CSVLoader(file_path)

    else:
        raise ValueError(f"不支持的文件类型: {file_type}")

    # 加载文档
    docs = loader.load()

    # 创建文本分割器
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        separators=["\n\n", "\n", ".", " ", ""],
    )

    # 分割文档
    nodes = text_splitter.split_documents(docs)

    # 添加序号信息到metadata
    for i, node in enumerate(nodes):
        if node.metadata is None:
            node.metadata = {}
        node.metadata["chunk_idx"] = i

    return nodes


def chunk_text(text, params=None):
    """
    将文本切分成固定大小的块
    """
    params = params or {}
    chunk_size = int(params.get("chunk_size", 500))
    chunk_overlap = int(params.get("chunk_overlap", 100))

    # 创建文本分割器
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size, chunk_overlap=chunk_overlap, separators=["\n\n", "\n", ".", " ", ""]
    )

    # 分割文档
    nodes = text_splitter.split_text(text)

    # 添加序号信息到metadata
    nodes = [{"text": node, "metadata": {"chunk_idx": i}} for i, node in enumerate(nodes)]
    return nodes


def chunk(text_or_path, params=None):
    raise NotImplementedError("chunk is deprecated, use chunk_with_parser or chunk_text instead")


def pdfreader(file_path, params=None):
    """读取PDF文件并返回text文本"""
    if isinstance(file_path, str):
        file_path = Path(file_path)

    assert file_path.exists(), "File not found"
    assert file_path.suffix.lower() == ".pdf", "File format not supported"

    # 使用LangChain的PDF加载器
    loader = PyPDFLoader(str(file_path))
    docs = loader.load()

    # 简单的拼接起来之后返回纯文本
    text = "\n\n".join([d.page_content for d in docs])
    return text


def plainreader(file_path):
    """读取普通文本文件并返回text文本"""
    assert os.path.exists(file_path), "File not found"

    # 使用LangChain的文本加载器
    loader = TextLoader(str(file_path))
    docs = loader.load()
    text = "\n\n".join([d.page_content for d in docs])
    return text


def parse_pdf(file, params=None):
    """
    解析PDF文件，支持多种OCR方式

    Args:
        file: PDF文件路径
        params: 参数字典，包含enable_ocr设置

    Returns:
        str: 解析得到的文本

    Raises:
        DocumentProcessorException: 处理失败时抛出
    """
    from yuxi.plugins.parser.base import DocumentProcessorException
    from yuxi.plugins.parser.factory import DocumentProcessorFactory

    params = params or {}
    opt_ocr = params.get("enable_ocr", "disable")

    if opt_ocr == "disable":
        return pdfreader(file, params=params)

    image_bucket, image_prefix = _resolve_image_storage_params(params)
    params.setdefault("image_bucket", image_bucket)
    params.setdefault("image_prefix", image_prefix)

    try:
        return DocumentProcessorFactory.process_file(opt_ocr, file, params)

    except DocumentProcessorException as e:
        logger.error(f"文档处理失败: {e.service_name} - {str(e)}")
        raise
    except Exception as e:
        logger.error(f"PDF 解析失败: {str(e)}")
        raise DocumentProcessorException(f"PDF解析失败: {str(e)}", opt_ocr, "parsing_failed")


def parse_image(file, params=None):
    """
    解析图像文件，支持多种OCR方式

    Args:
        file: 图像文件路径
        params: 参数字典，包含enable_ocr设置

    Returns:
        str: 解析得到的文本

    Raises:
        DocumentProcessorException: 处理失败时抛出
        ValueError: 图像文件禁用OCR时抛出
    """
    from yuxi.plugins.parser.base import DocumentProcessorException
    from yuxi.plugins.parser.factory import DocumentProcessorFactory

    params = params or {}
    opt_ocr = params.get("enable_ocr", "disable")

    # 图像文件必须使用 OCR,不能禁用
    if opt_ocr == "disable":
        raise ValueError(
            "图像文件必须启用OCR才能提取文本内容。"
            "请选择OCR方式 (rapid_ocr/mineru_ocr/mineru_official/pp_structure_v3_ocr) 或移除该文件。"
        )

    image_bucket, image_prefix = _resolve_image_storage_params(params)
    params.setdefault("image_bucket", image_bucket)
    params.setdefault("image_prefix", image_prefix)

    try:
        return DocumentProcessorFactory.process_file(opt_ocr, file, params)

    except DocumentProcessorException as e:
        logger.error(f"图像处理失败: {e.service_name} - {str(e)}")
        raise
    except Exception as e:
        logger.error(f"图像解析失败: {str(e)}")
        raise DocumentProcessorException(f"图像解析失败: {str(e)}", opt_ocr, "parsing_failed")


async def parse_pdf_async(file, params=None):
    return await asyncio.to_thread(parse_pdf, file, params=params)


async def parse_image_async(file, params=None):
    return await asyncio.to_thread(parse_image, file, params=params)


async def process_file_to_markdown(file_path: str, params: dict | None = None) -> str:
    """
    将不同类型的文件转换为markdown格式 - 支持本地文件和MinIO文件

    Args:
        file_path: 文件路径或MinIO URL
        params: 处理参数，对于ZIP文件需要包含 db_id

    Returns:
        markdown格式内容

    Note:
        对于ZIP文件，会在params中保存处理结果供调用方使用：
        - params['_zip_images_info']: 图片信息列表
        - params['_zip_content_hash']: 内容哈希值
    """
    import os
    import tempfile

    # 检测是否是MinIO URL
    from yuxi.knowledge.utils.kb_utils import is_minio_url

    if is_minio_url(file_path):
        # 从MinIO下载文件到临时位置
        logger.debug(f"Downloading file from MinIO: {file_path}")

        # 从MinIO URL中提取文件名
        if "?" in file_path:
            file_path_clean = file_path.split("?")[0]
        else:
            file_path_clean = file_path

        original_filename = file_path_clean.split("/")[-1]

        # 创建临时文件
        with tempfile.NamedTemporaryFile(delete=False, suffix=Path(original_filename).suffix) as temp_file:
            temp_path = temp_file.name

        try:
            # 使用通用函数解析MinIO URL并下载文件
            from yuxi.knowledge.utils.kb_utils import parse_minio_url
            from yuxi.storage.minio.client import get_minio_client

            # 解析MinIO URL获取bucket_name和object_name
            bucket_name, object_name = parse_minio_url(file_path)

            # 获取MinIO客户端并下载文件
            minio_client = get_minio_client()
            file_content = await minio_client.adownload_file(bucket_name, object_name)

            # 写入临时文件
            async with aiofiles.open(temp_path, "wb") as f:
                await f.write(file_content)

            logger.debug(f"File downloaded to temp path: {temp_path}")

            # 使用临时文件路径
            actual_file_path = temp_path

        except Exception as e:
            # 清理临时文件
            if os.path.exists(temp_path):
                os.unlink(temp_path)
            logger.error(f"Failed to download file from MinIO: {e}")
            raise ValueError(f"无法从MinIO下载文件: {e}")
    else:
        # 本地文件
        actual_file_path = file_path

    try:
        file_path_obj = Path(actual_file_path)
        file_ext = file_path_obj.suffix.lower()
        original_filename = file_path_obj.name

        if file_ext == ".pdf":
            # 使用 OCR 处理 PDF
            text = await parse_pdf_async(str(file_path_obj), params=params)
            result = f"{text}"

        elif file_ext in [".txt", ".md"]:
            # 直接读取文本文件
            with open(file_path_obj, encoding="utf-8") as f:
                content = f.read()
            result = f"{content}"

        elif file_ext == ".docx":
            # 优先使用 Docling，失败时回退到 python-docx
            try:
                result = _convert_with_docling(file_path_obj, params=params)
            except Exception as e:
                logger.warning(f"Docling 解析 DOCX 失败，回退到 python-docx: {file_path_obj.name}, {e}")
                result = _convert_docx_with_python_docx(file_path_obj)

        elif file_ext == ".pptx":
            # 使用 Docling 处理 pptx
            result = _convert_with_docling(file_path_obj, params=params)

        elif file_ext == ".doc":
            # 旧版 .doc 文件仍使用原有解析方式
            from langchain_community.document_loaders import UnstructuredWordDocumentLoader

            loader = UnstructuredWordDocumentLoader(str(file_path_obj))
            docs = loader.load()
            result = "\n".join(doc.page_content for doc in docs).strip()

        elif file_ext in [".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif"]:
            # 使用 OCR 处理图片
            text = await parse_image_async(str(file_path_obj), params=params)
            result = f"{text}"

        elif file_ext in [".html", ".htm"]:
            # 使用 BeautifulSoup 处理 HTML 文件
            from markdownify import markdownify as md

            with open(file_path_obj, encoding="utf-8") as f:
                content = f.read()
            text = md(content, heading_style="ATX")
            result = f"{text}"

        elif file_ext == ".csv":
            # 处理 CSV 文件
            import pandas as pd

            df = pd.read_csv(file_path_obj)
            # 将每一行数据与表头组合成独立的表格
            markdown_content = ""

            for index, row in df.iterrows():
                # 创建包含表头和当前行的小表格
                row_df = pd.DataFrame([row], columns=df.columns)
                markdown_table = row_df.to_markdown(index=False)
                markdown_content += f"{markdown_table}\n\n"

            result = markdown_content.strip()

        elif file_ext in [".xls", ".xlsx"]:
            # 使用 Docling 处理 Excel 文件
            result = _convert_with_docling(file_path_obj, params=params)

        elif file_ext == ".json":
            # 处理 JSON 文件
            import json

            async with aiofiles.open(file_path_obj, encoding="utf-8") as f:
                content = await f.read()
            data = json.loads(content)
            # 将 JSON 数据格式化为 markdown 代码块
            json_str = json.dumps(data, ensure_ascii=False, indent=2)
            result = f"```json\n{json_str}\n```"

        elif file_ext == ".zip":
            image_bucket, image_prefix = _resolve_image_storage_params(params)
            zip_result = await _process_zip_file(
                str(file_path_obj),
                image_bucket=image_bucket,
                image_prefix=image_prefix,
            )

            # 将处理结果保存到params中供调用方使用
            if params is not None:
                params["_zip_images_info"] = zip_result["images_info"]
                params["_zip_content_hash"] = zip_result["content_hash"]
                params["_zip_image_bucket"] = image_bucket
                params["_zip_image_prefix"] = image_prefix

            result = zip_result["markdown_content"]

        else:
            # 尝试作为文本文件读取
            raise ValueError(f"Unsupported file type: {file_ext}")

    except Exception:
        # 清理临时文件
        if is_minio_url(file_path) and os.path.exists(actual_file_path):
            try:
                os.unlink(actual_file_path)
                logger.debug(f"Cleaned up temp file: {actual_file_path}")
            except Exception as cleanup_e:
                logger.warning(f"Failed to clean up temp file {actual_file_path}: {cleanup_e}")
        raise

    finally:
        # 清理临时文件
        if is_minio_url(file_path) and os.path.exists(actual_file_path):
            try:
                os.unlink(actual_file_path)
                logger.debug(f"Cleaned up temp file: {actual_file_path}")
            except Exception as e:
                logger.warning(f"Failed to clean up temp file {actual_file_path}: {e}")

    return result


async def process_url_to_markdown(url: str, params: dict | None = None) -> str:
    """
    Fetch a URL and convert its content to Markdown.

    Args:
        url: The URL to fetch.
        params: Optional parameters (unused, kept for API compatibility).

    Returns:
        The Markdown content of the URL.
    """
    logger.info(f"Fetching URL: {url}")

    try:
        import httpx

        # 使用异步 HTTP 客户端获取页面
        async with httpx.AsyncClient(timeout=30.0, follow_redirects=True) as client:
            response = await client.get(url, headers={"User-Agent": "Mozilla/5.0"})
            response.raise_for_status()
            html_content = response.text

        # 使用 readability 提取正文 HTML
        from readability import Document

        doc = Document(html_content)
        body_html = doc.summary()

        # 转换为 Markdown
        markdown_content = md_convert(body_html, heading_style="atx")

        logger.info(f"Successfully converted URL to Markdown: {url}")
        return markdown_content

    except httpx.HTTPError as e:
        logger.error(f"Failed to fetch URL {url}: {e}")
        raise ValueError(f"Failed to fetch URL: {e}")
    except Exception as e:
        logger.error(f"Failed to process URL {url}: {e}")
        raise ValueError(f"Failed to process URL: {e}")