ForcePilot/backend/package/yuxi/channels/adapters/feishu/tools/doc.py

from __future__ import annotations

from typing import Any

from ._validate import validate_params

try:
    import lark_oapi

    HAS_LARK_SDK = True
except ImportError:
    HAS_LARK_SDK = False
    lark_oapi = None  # type: ignore


@validate_params
async def create_doc(client: Any, title: str, folder_token: str = "") -> dict[str, Any]:
    if not HAS_LARK_SDK or not client:
        raise RuntimeError("SDK 不可用")

    try:
        request_body = lark_oapi.api.docx.v1.CreateDocumentRequestBody.builder().title(title).build()
        if folder_token:
            request_body.folder_token = folder_token

        request = lark_oapi.api.docx.v1.CreateDocumentRequest.builder().request_body(request_body).build()
        resp = await client.docx.v1.document.create(request)
        if not resp.success():
            raise RuntimeError(f"创建文档失败: {resp.msg}")

        data = resp.data if hasattr(resp, "data") else {}
        doc = data.get("document", {})
        return {
            "document_id": doc.get("document_id", ""),
            "title": doc.get("title", title),
            "url": doc.get("url", ""),
        }
    except Exception as e:
        raise RuntimeError(f"创建文档失败: {e}") from e


@validate_params
async def get_doc_content(client: Any, document_id: str) -> dict[str, Any]:
    if not HAS_LARK_SDK or not client:
        raise RuntimeError("SDK 不可用")

    try:
        request = lark_oapi.api.docx.v1.GetDocumentRequest.builder().document_id(document_id).build()
        resp = await client.docx.v1.document.get(request)
        if not resp.success():
            raise RuntimeError(f"获取文档失败: {resp.msg}")

        data = resp.data if hasattr(resp, "data") else {}
        doc = data.get("document", {})
        blocks = doc.get("blocks", [])
        return {
            "document_id": doc.get("document_id", document_id),
            "title": doc.get("title", ""),
            "block_count": len(blocks),
            "blocks": _extract_doc_blocks(blocks),
        }
    except Exception as e:
        raise RuntimeError(f"获取文档失败: {e}") from e


async def list_docs(client: Any, folder_token: str = "", page_size: int = 50) -> list[dict[str, Any]]:
    if not HAS_LARK_SDK or not client:
        return []

    docs: list[dict[str, Any]] = []
    page_token = ""
    try:
        while True:
            request = (
                lark_oapi.api.docx.v1.ListDocumentRequest.builder()
                .page_size(page_size)
                .page_token(page_token)
                .folder_token(folder_token)
                .build()
            )
            resp = await client.docx.v1.document.list(request)
            if not resp.success():
                break

            data = resp.data if hasattr(resp, "data") else {}
            items = data.get("items", [])
            for item in items:
                docs.append(
                    {
                        "document_id": item.get("document_id", ""),
                        "title": item.get("title", ""),
                        "url": item.get("url", ""),
                        "create_time": item.get("create_time", ""),
                        "edit_time": item.get("edit_time", ""),
                    }
                )

            page_token = data.get("page_token", "")
            if not page_token or not items:
                break
    except Exception as e:
        raise RuntimeError(f"获取文档列表失败: {e}") from e

    return docs


def _extract_doc_blocks(blocks: list) -> list[dict[str, Any]]:
    result: list[dict[str, Any]] = []
    for block in blocks:
        block_type = block.get("block_type", block.get("blockType", 0))
        extracted = {"block_type": block_type}

        text_content = ""
        for elem_type in ("text", "heading1", "heading2", "heading3", "heading4", "heading5", "bullet", "ordered"):
            elem = block.get(elem_type, {})
            if elem:
                elements = elem.get("elements", [])
                for e in elements:
                    text_run = e.get("text_run", {})
                    text_content += text_run.get("content", "")

        if text_content:
            extracted["content"] = text_content
            result.append(extracted)

    return result


@validate_params
async def update_doc_block(
    client: Any,
    document_id: str,
    block_id: str,
    content: str,
    block_type: str = "text",
) -> dict[str, Any]:
    if not HAS_LARK_SDK or not client:
        raise RuntimeError("SDK 不可用")

    try:
        text_elements = [{"text_run": {"content": content}}]
        block_data = {
            "block_id": block_id,
            block_type: {"elements": text_elements},
        }
        if block_type in ("heading1", "heading2", "heading3", "heading4", "heading5"):
            block_data[block_type]["style"] = {}

        request_body = (
            lark_oapi.api.docx.v1.UpdateDocumentBlockRequestBody.builder()
            .update_document_block_request(
                lark_oapi.api.docx.v1.UpdateDocumentBlockRequestBlock.builder().replace_block(block_data).build()
            )
            .build()
        )

        request = (
            lark_oapi.api.docx.v1.PatchDocumentBlockRequest.builder()
            .document_id(document_id)
            .block_id(block_id)
            .request_body(request_body)
            .build()
        )

        resp = await client.docx.v1.document_block.patch(request)
        if not resp.success():
            raise RuntimeError(f"更新文档块失败: {resp.msg}")

        data = resp.data if hasattr(resp, "data") else {}
        block = data.get("block", {}) or data
        return {
            "block_id": block.get("block_id", block_id),
            "document_id": document_id,
        }
    except Exception as e:
        raise RuntimeError(f"更新文档块失败: {e}") from e


@validate_params
async def append_doc_block(
    client: Any,
    document_id: str,
    parent_block_id: str,
    content: str,
    block_type: str = "text",
    index: int = -1,
) -> dict[str, Any]:
    if not HAS_LARK_SDK or not client:
        raise RuntimeError("SDK 不可用")

    try:
        text_elements = [{"text_run": {"content": content}}]
        children: list[dict[str, Any]] = [
            {
                "block_type": block_type,
                block_type: {"elements": text_elements},
            }
        ]

        request_body = (
            lark_oapi.api.docx.v1.CreateDocumentBlockChildrenRequestBody.builder()
            .children(children)
            .index(index if index >= 0 else -1)
            .build()
        )

        request = (
            lark_oapi.api.docx.v1.CreateDocumentBlockChildrenRequest.builder()
            .document_id(document_id)
            .block_id(parent_block_id)
            .request_body(request_body)
            .build()
        )

        resp = await client.docx.v1.document_block_children.create(request)
        if not resp.success():
            raise RuntimeError(f"追加文档块失败: {resp.msg}")

        data = resp.data if hasattr(resp, "data") else {}
        children_resp = data.get("children", [])
        block = children_resp[0] if children_resp else data
        return {
            "block_id": block.get("block_id", ""),
            "parent_block_id": parent_block_id,
            "document_id": document_id,
        }
    except Exception as e:
        raise RuntimeError(f"追加文档块失败: {e}") from e


@validate_params
async def delete_doc_block(
    client: Any,
    document_id: str,
    block_id: str,
) -> dict[str, Any]:
    if not HAS_LARK_SDK or not client:
        raise RuntimeError("SDK 不可用")

    try:
        request = (
            lark_oapi.api.docx.v1.DeleteDocumentBlockRequest.builder()
            .document_id(document_id)
            .block_id(block_id)
            .build()
        )

        resp = await client.docx.v1.document_block.delete(request)
        if not resp.success():
            raise RuntimeError(f"删除文档块失败: {resp.msg}")

        return {
            "deleted": True,
            "block_id": block_id,
            "document_id": document_id,
        }
    except Exception as e:
        raise RuntimeError(f"删除文档块失败: {e}") from e