"""Shared Google Drive storage — in-process MCP so Jarvis (and any AI on the
same credentials) can read/write/upload/remove files from one shared folder.

This is the storage twin of the shared MEMORY: a single "Jarvis Drive" folder on
Ahmed's Google account that every Jarvis instance / Claude can use to pass files
around (a CRM export, a doc, a screenshot, notes).

AUTH = SERVICE ACCOUNT (key in jarvis_drive_sa.json), NOT a user login. A service
account has no Drive of its own, so it can ONLY see the folders Ahmed explicitly
SHARES with it — the one "Jarvis Drive" folder. His personal Drive is not merely
code-blocked, it is structurally invisible: the SA cannot list or reach it at all.
The key also works on any machine (no per-user login, no token expiry), so this
is genuinely shared across devices. MCP_DRIVE=0 disables the whole server.

Tools (Claude sees them as mcp__drive__<name>):
  drive_list      list files in the shared folder
  drive_read      read a text file's contents (by name or id)
  drive_write     create/overwrite a text file from given content
  drive_upload    upload a local file into the folder
  drive_download  download a folder file to a local path
  drive_delete    remove (trash) a file (by name or id)
"""
from __future__ import annotations

import io
import mimetypes
import os
from pathlib import Path

_HERE = Path(__file__).resolve().parent.parent
_SA_KEY = _HERE / "jarvis_drive_sa.json"      # service-account key (gitignored)
_FOLDER_FILE = _HERE / "jarvis_drive_folder.txt"
# Full Drive scope, but on a SERVICE ACCOUNT it only reaches what Ahmed shares
# with it (the Jarvis Drive folder) — an SA has no Drive of its own to roam. The
# code is also folder-jailed below as defence in depth.
_SCOPES = ["https://www.googleapis.com/auth/drive"]


def enabled() -> bool:
    """On when the service-account key exists and it isn't explicitly disabled."""
    if os.environ.get("MCP_DRIVE", "1") == "0":
        return False
    return _SA_KEY.exists()


def _text(msg: str) -> dict:
    return {"content": [{"type": "text", "text": msg}]}


def _folder_id() -> str:
    return _FOLDER_FILE.read_text().strip()


def _service():
    """Build a Drive v3 client authenticated as the service account. The SA can
    only reach folders shared with it, so this can never touch Ahmed's own files."""
    from google.oauth2 import service_account
    from googleapiclient.discovery import build
    creds = service_account.Credentials.from_service_account_file(
        str(_SA_KEY), scopes=_SCOPES)
    return build("drive", "v3", credentials=creds, cache_discovery=False)


def _resolve(svc, name_or_id: str) -> dict | None:
    """Find a file by id or name — but ONLY inside the Jarvis Drive folder.
    This is the code jail: with the full-Drive token, a bare id could point at
    any file in Ahmed's Drive, so we REFUSE anything whose parent isn't our
    folder. Every read/write/delete goes through here, so nothing outside the
    shared folder is ever touched."""
    ref = str(name_or_id).strip()
    if not ref:
        return None
    fid = _folder_id()
    # try as an id, but only accept it if it lives in our folder
    try:
        f = svc.files().get(
            fileId=ref, fields="id,name,mimeType,size,parents").execute()
        return f if fid in (f.get("parents") or []) else None
    except Exception:  # noqa: BLE001 — not an id; fall through to name search
        pass
    q = (f"'{fid}' in parents and trashed=false and name='{ref}'")
    files = svc.files().list(q=q, spaces="drive",
                            fields="files(id,name,mimeType,size)"
                            ).execute().get("files", [])
    return files[0] if files else None


def build_server():
    import asyncio
    from claude_agent_sdk import tool, create_sdk_mcp_server
    from googleapiclient.http import MediaFileUpload, MediaIoBaseUpload

    @tool("drive_list",
          "List the files in the shared Jarvis Drive folder (name, type, size, "
          "last modified). Use for 'what's in the drive/storage', 'list my "
          "shared files', 'what did I save there'.", {})
    async def drive_list(args: dict) -> dict:
        def work():
            svc = _service()
            q = f"'{_folder_id()}' in parents and trashed=false"
            files = svc.files().list(
                q=q, spaces="drive", orderBy="modifiedTime desc",
                fields="files(id,name,mimeType,size,modifiedTime)"
                ).execute().get("files", [])
            return files
        try:
            files = await asyncio.to_thread(work)
            if not files:
                return _text("The shared Drive folder is empty.")
            lines = []
            for f in files:
                sz = f.get("size")
                sz = f"{int(sz)//1024}KB" if sz else "—"
                lines.append(f"- {f['name']}  ({sz}, {f['modifiedTime'][:10]}) "
                             f"[id {f['id']}]")
            return _text(f"{len(files)} file(s) in Jarvis Drive:\n"
                         + "\n".join(lines))
        except Exception as e:  # noqa: BLE001
            return _text(f"drive_list failed: {str(e)[:200]}")

    @tool("drive_read",
          "Read a text file's contents from the shared folder. `file` = the "
          "file name or its id. Works for plain text / csv / json / markdown "
          "and Google Docs (exported as text).", {"file": str})
    async def drive_read(args: dict) -> dict:
        def work():
            svc = _service()
            f = _resolve(svc, args.get("file", ""))
            if not f:
                return None, "not found"
            mime = f.get("mimeType", "")
            if mime.startswith("application/vnd.google-apps"):
                data = svc.files().export(
                    fileId=f["id"], mimeType="text/plain").execute()
            else:
                data = svc.files().get_media(fileId=f["id"]).execute()
            return f, (data.decode("utf-8", "ignore")
                       if isinstance(data, bytes) else str(data))
        try:
            f, body = await asyncio.to_thread(work)
            if f is None:
                return _text("drive_read: no such file in the shared folder.")
            snip = body[:6000] + (" …[truncated]" if len(body) > 6000 else "")
            return _text(f"{f['name']}:\n{snip}")
        except Exception as e:  # noqa: BLE001
            return _text(f"drive_read failed: {str(e)[:200]}")

    @tool("drive_write",
          "Create (or overwrite) a text file in the shared folder from content "
          "you provide. `name` = the file name (e.g. 'riyadh-leads.csv'), "
          "`content` = the full text to store. Use to hand a file to another AI/"
          "device, e.g. write a CRM or notes for Jarvis to pick up later.",
          {"name": str, "content": str})
    async def drive_write(args: dict) -> dict:
        name = str(args.get("name", "")).strip()
        content = str(args.get("content", ""))
        if not name:
            return _text("drive_write failed: need a file name.")

        def work():
            svc = _service()
            mime = mimetypes.guess_type(name)[0] or "text/plain"
            media = MediaIoBaseUpload(
                io.BytesIO(content.encode("utf-8")), mimetype=mime,
                resumable=False)
            existing = _resolve(svc, name)
            if existing and existing.get("name") == name:
                svc.files().update(fileId=existing["id"], media_body=media
                                   ).execute()
                return "updated", existing["id"]
            meta = {"name": name, "parents": [_folder_id()]}
            out = svc.files().create(body=meta, media_body=media, fields="id"
                                     ).execute()
            return "created", out["id"]
        try:
            verb, fid = await asyncio.to_thread(work)
            return _text(f"{verb} {name} in Jarvis Drive (id {fid}).")
        except Exception as e:  # noqa: BLE001
            return _text(f"drive_write failed: {str(e)[:200]}")

    @tool("drive_upload",
          "Upload a LOCAL file (by path on this Mac) into the shared folder. "
          "`local_path` = absolute path; `name` = optional name to store it as "
          "(defaults to the file's own name). Use for 'put this file on the "
          "drive', 'upload X to storage'.",
          {"local_path": str, "name": str})
    async def drive_upload(args: dict) -> dict:
        lp = os.path.expanduser(str(args.get("local_path", "")).strip())
        if not lp or not os.path.isfile(lp):
            return _text(f"drive_upload failed: no file at {lp!r}.")
        name = str(args.get("name", "")).strip() or os.path.basename(lp)

        def work():
            svc = _service()
            mime = mimetypes.guess_type(lp)[0] or "application/octet-stream"
            media = MediaFileUpload(lp, mimetype=mime, resumable=True)
            existing = _resolve(svc, name)
            if existing and existing.get("name") == name:
                svc.files().update(fileId=existing["id"], media_body=media
                                   ).execute()
                return "updated", existing["id"]
            meta = {"name": name, "parents": [_folder_id()]}
            out = svc.files().create(body=meta, media_body=media, fields="id"
                                     ).execute()
            return "uploaded", out["id"]
        try:
            verb, fid = await asyncio.to_thread(work)
            return _text(f"{verb} {name} to Jarvis Drive (id {fid}).")
        except Exception as e:  # noqa: BLE001
            return _text(f"drive_upload failed: {str(e)[:200]}")

    @tool("drive_download",
          "Download a file from the shared folder to a LOCAL path on this Mac. "
          "`file` = name or id; `local_path` = where to save it. Use to pull a "
          "shared file down so a local tool (or Jarvis) can work with it.",
          {"file": str, "local_path": str})
    async def drive_download(args: dict) -> dict:
        dest = os.path.expanduser(str(args.get("local_path", "")).strip())
        if not dest:
            return _text("drive_download failed: need a local_path.")

        def work():
            from googleapiclient.http import MediaIoBaseDownload
            svc = _service()
            f = _resolve(svc, args.get("file", ""))
            if not f:
                return None
            buf = io.FileIO(dest, "wb")
            dl = MediaIoBaseDownload(buf, svc.files().get_media(fileId=f["id"]))
            done = False
            while not done:
                _, done = dl.next_chunk()
            buf.close()
            return f["name"]
        try:
            name = await asyncio.to_thread(work)
            if not name:
                return _text("drive_download: no such file in the shared folder.")
            return _text(f"Downloaded {name} → {dest}.")
        except Exception as e:  # noqa: BLE001
            return _text(f"drive_download failed: {str(e)[:200]}")

    @tool("drive_delete",
          "Remove (move to trash) a file from the shared folder. `file` = name "
          "or id. Use for 'delete X from the drive', 'remove that file'.",
          {"file": str})
    async def drive_delete(args: dict) -> dict:
        def work():
            svc = _service()
            f = _resolve(svc, args.get("file", ""))
            if not f:
                return None
            svc.files().update(fileId=f["id"], body={"trashed": True}).execute()
            return f["name"]
        try:
            name = await asyncio.to_thread(work)
            if not name:
                return _text("drive_delete: no such file in the shared folder.")
            return _text(f"Deleted {name} from Jarvis Drive.")
        except Exception as e:  # noqa: BLE001
            return _text(f"drive_delete failed: {str(e)[:200]}")

    return create_sdk_mcp_server(
        name="drive", version="1.0.0",
        tools=[drive_list, drive_read, drive_write, drive_upload,
               drive_download, drive_delete],
    )