ForcePilot/backend/package/yuxi/channels/adapters/feishu/tools/drive.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


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

    files: list[dict[str, Any]] = []
    page_token = ""
    try:
        while True:
            request = (
                lark_oapi.api.drive.v1.ListFileRequest.builder()
                .page_size(page_size)
                .page_token(page_token)
                .folder_token(folder_token)
                .order_by(order_by)
                .direction(direction)
                .build()
            )
            resp = await client.drive.v1.file.list(request)
            if not resp.success():
                break

            data = resp.data if hasattr(resp, "data") else {}
            items = data.get("files", [])
            for item in items:
                files.append(
                    {
                        "token": item.get("token", ""),
                        "name": item.get("name", ""),
                        "type": item.get("type", ""),
                        "size": item.get("size", 0),
                        "url": item.get("url", ""),
                        "created_time": item.get("created_time", ""),
                        "modified_time": item.get("modified_time", ""),
                        "owner_id": item.get("owner_id", ""),
                    }
                )

            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 files


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

    try:
        request = lark_oapi.api.drive.v1.GetFileRequest.builder().file_token(file_token).build()
        resp = await client.drive.v1.file.get(request)
        if not resp.success():
            raise RuntimeError(f"获取文件详情失败: {resp.msg}")

        data = resp.data if hasattr(resp, "data") else {}
        return {
            "token": data.get("token", file_token),
            "name": data.get("name", ""),
            "type": data.get("type", ""),
            "size": data.get("size", 0),
            "url": data.get("url", ""),
            "created_time": data.get("created_time", ""),
            "modified_time": data.get("modified_time", ""),
            "owner_id": data.get("owner_id", ""),
        }
    except Exception as e:
        raise RuntimeError(f"获取文件详情失败: {e}") from e


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

    try:
        request_body = (
            lark_oapi.api.drive.v1.CreateFolderRequestBody.builder().name(name).folder_token(parent_token).build()
        )
        request = lark_oapi.api.drive.v1.CreateFolderRequest.builder().request_body(request_body).build()
        resp = await client.drive.v1.folder.create(request)
        if not resp.success():
            raise RuntimeError(f"创建文件夹失败: {resp.msg}")

        data = resp.data if hasattr(resp, "data") else {}
        return {
            "token": data.get("token", ""),
            "name": data.get("name", name),
            "url": data.get("url", ""),
        }
    except Exception as e:
        raise RuntimeError(f"创建文件夹失败: {e}") from e


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

    import os

    try:
        if not file_name:
            file_name = os.path.basename(file_path)

        file_size = os.path.getsize(file_path)

        with open(file_path, "rb") as f:
            if file_size <= 20 * 1024 * 1024:
                resp = await client.drive.v1.media.upload(
                    folder_token=folder_token,
                    file_name=file_name,
                    file=f,
                )
            else:
                resp = await client.drive.v1.media.upload_all(
                    folder_token=folder_token,
                    file_name=file_name,
                    file=f,
                )

        if not resp.success():
            raise RuntimeError(f"上传文件失败: {resp.msg}")

        data = resp.data if hasattr(resp, "data") else {}
        return {
            "token": data.get("file_token", ""),
            "name": file_name,
            "size": file_size,
        }
    except Exception as e:
        raise RuntimeError(f"上传文件失败: {e}") from e


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

    import os

    try:
        request = lark_oapi.api.drive.v1.DownloadFileRequest.builder().file_token(file_token).build()

        resp = await client.drive.v1.file.download(request)
        if not resp.success():
            raise RuntimeError(f"下载文件失败: {resp.msg}")

        content = resp.file.read() if hasattr(resp, "file") else resp.read()

        os.makedirs(os.path.dirname(save_path) or ".", exist_ok=True)
        with open(save_path, "wb") as f:
            f.write(content)

        return {
            "token": file_token,
            "saved_path": save_path,
            "size": len(content),
        }
    except Exception as e:
        raise RuntimeError(f"下载文件失败: {e}") from e