ForcePilot/backend/package/yuxi/channel/extensions/msteams/graph_upload.py

from __future__ import annotations

import logging

from .graph import MSTeamsGraphClient

logger = logging.getLogger(__name__)


async def upload_to_onedrive(
    graph_client: MSTeamsGraphClient,
    file_content: bytes,
    file_name: str,
    *,
    folder: str = "ForcePilot",
) -> dict | None:
    upload_path = f"/me/drive/root:/{folder}/{file_name}:/content"
    upload_url = f"/me/drive/root:/{folder}/{file_name}:/createUploadSession"

    try:
        session_data = await graph_client.fetch_json(
            upload_url,
            method="POST",
            body={
                "@microsoft.graph.conflictBehavior": "rename",
            },
        )
    except Exception as e:
        logger.warning("Failed to create OneDrive upload session for %s: %s", file_name, e)
        return None

    upload_url_dest = session_data.get("uploadUrl", "")
    if not upload_url_dest:
        try:
            result = await graph_client.fetch_json(upload_path, method="PUT", body=file_content)
            return await _create_share_link(graph_client, result.get("id", ""), file_name)
        except Exception as e:
            logger.warning("Failed to upload %s to OneDrive: %s", file_name, e)
            return None

    chunk_size = 327680
    total_len = len(file_content)
    offset = 0

    while offset < total_len:
        end = min(offset + chunk_size, total_len)
        chunk = file_content[offset:end]
        headers = {
            "Content-Length": str(len(chunk)),
            "Content-Range": f"bytes {offset}-{end - 1}/{total_len}",
        }
        try:
            import httpx

            token = await graph_client._get_token()
            async with httpx.AsyncClient(timeout=60.0) as client:
                resp = await client.put(
                    upload_url_dest,
                    content=chunk,
                    headers={"Authorization": f"Bearer {token}", **headers},
                )
                resp.raise_for_status()
                if end == total_len:
                    result = resp.json()
                    return await _create_share_link(graph_client, result.get("id", ""), file_name)
        except Exception as e:
            logger.warning("Failed to upload chunk for %s: %s", file_name, e)
            return None

        offset = end

    return None


async def upload_to_sharepoint(
    graph_client: MSTeamsGraphClient,
    file_content: bytes,
    file_name: str,
    site_id: str,
    *,
    folder: str = "General",
) -> dict | None:
    path = f"/sites/{site_id}/drive/root:/{folder}/{file_name}:/content"
    try:
        await graph_client.fetch_json(path, method="PUT", body=file_content)
        item_path = f"/sites/{site_id}/drive/root:/{folder}/{file_name}"
        item_data = await graph_client.fetch_json(item_path)
        return await _create_share_link(graph_client, item_data.get("id", ""), file_name, site_id=site_id)
    except Exception as e:
        logger.warning("Failed to upload %s to SharePoint: %s", file_name, e)
        return None


async def _create_share_link(
    graph_client: MSTeamsGraphClient,
    item_id: str,
    file_name: str,
    site_id: str | None = None,
) -> dict | None:
    if not item_id:
        return None

    if site_id:
        share_path = f"/sites/{site_id}/drive/items/{item_id}/createLink"
    else:
        share_path = f"/me/drive/items/{item_id}/createLink"

    try:
        result = await graph_client.fetch_json(
            share_path,
            method="POST",
            body={
                "type": "view",
                "scope": "organization",
            },
        )
        web_url = result.get("link", {}).get("webUrl", "")
        return {
            "file_name": file_name,
            "item_id": item_id,
            "web_url": web_url,
            "share_url": web_url,
        }
    except Exception as e:
        logger.warning("Failed to create share link for %s: %s", file_name, e)

    try:
        item_path = f"/me/drive/items/{item_id}"
        item = await graph_client.fetch_json(item_path)
        web_url = item.get("webUrl", "")
        return {
            "file_name": file_name,
            "item_id": item_id,
            "web_url": web_url,
            "share_url": web_url,
        }
    except Exception:
        pass

    return None


async def get_file_content(graph_client: MSTeamsGraphClient, item_id: str) -> bytes | None:
    path = f"/me/drive/items/{item_id}/content"
    try:
        import httpx

        token = await graph_client._get_token()
        url = f"https://graph.microsoft.com/v1.0{path}"
        async with httpx.AsyncClient(timeout=60.0) as client:
            resp = await client.get(url, headers={"Authorization": f"Bearer {token}"})
            resp.raise_for_status()
            return resp.content
    except Exception as e:
        logger.warning("Failed to get file content for %s: %s", item_id, e)
        return None