from __future__ import annotations

import asyncio

from mcp.server.fastmcp import FastMCP
from mrp import Agent

from mrp_mcp.state import SessionState


def register(mcp: FastMCP, agent: Agent, state: SessionState) -> None:
    @mcp.tool()
    async def mrp_upload_blob(
        file_path: str,
        content_type: str = "application/octet-stream",
    ) -> str:
        """Upload a file to share with other agents.

        The file is stored on the relay and you get a blob ID that you can
        include when sending messages (via structured_body with attachments).

        Args:
            file_path: Path to the file to upload. Max 100 MiB.
            content_type: MIME type of the file (e.g. "image/png", "application/pdf").

        Returns:
            The blob ID and metadata. Use this blob_id in mrp_send's structured_body.
        """

        def _upload():
            with open(file_path, "rb") as f:
                data = f.read()
            return agent.upload(data, content_type=content_type)

        blob = await asyncio.to_thread(_upload)

        return (
            f"File uploaded successfully.\n"
            f"Blob ID: {blob.blob_id}\n"
            f"Size: {blob.size} bytes\n"
            f"Hash: {blob.hash}\n"
            f"Expires: {blob.expires_at.isoformat()}\n\n"
            f"To send this file, include the blob_id as an attachment when sending a message."
        )

    @mcp.tool()
    async def mrp_download_blob(
        blob_id: str,
        save_to: str,
    ) -> str:
        """Download a file (blob) from the relay.

        Use this to download files attached to messages you've received.

        Args:
            blob_id: The blob ID from a message's attachments.
            save_to: Local file path to save the downloaded file.

        Returns:
            Confirmation with file path and size.
        """

        def _download():
            data, ct = agent.download(blob_id)
            with open(save_to, "wb") as f:
                f.write(data)
            return data, ct

        data, content_type_result = await asyncio.to_thread(_download)

        return (
            f"File downloaded successfully.\n"
            f"Saved to: {save_to}\n"
            f"Size: {len(data)} bytes\n"
            f"Content type: {content_type_result}"
        )
