#!/usr/bin/env python3
"""
Google Drive backup for Alpha MC — uses a Service Account (headless, no browser needed).

Setup (one-time):
  1. Go to https://console.cloud.google.com
  2. Create a project → Enable "Google Drive API"
  3. IAM & Admin → Service Accounts → Create service account
  4. Create a JSON key for the service account → download it
  5. In Google Drive, create a folder called "Minecraft-Bedrock-Backups"
     and share it with the service account email (e.g. mc@project.iam.gserviceaccount.com)
  6. Base64-encode the JSON key:
       base64 -w 0 your-key.json
  7. Add the result as HF Space Secret named: GDRIVE_SA_KEY

Usage:
  python3 google_drive_backup.py <path_to_backup.tar.gz>
"""

import os
import sys
import json
import base64
import logging
from pathlib import Path

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [DRIVE] %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S"
)
log = logging.getLogger(__name__)

BACKUP_FOLDER_NAME = "Minecraft-Bedrock-Backups"
MAX_BACKUPS = 10          # oldest files deleted after this limit
SA_KEY_ENV  = "GDRIVE_SA_KEY"   # HF Secret name (base64-encoded service account JSON)
SA_KEY_FILE = "/tmp/gdrive_sa.json"   # written at runtime, never committed


def load_service_account_credentials():
    """
    Loads credentials from the GDRIVE_SA_KEY environment variable.
    The variable must contain the base64-encoded service account JSON key.
    Returns a google.oauth2.service_account.Credentials object.
    """
    from google.oauth2 import service_account

    raw = os.environ.get(SA_KEY_ENV, "").strip()
    if not raw:
        raise EnvironmentError(
            f"Environment variable '{SA_KEY_ENV}' is not set.\n"
            "  1. Follow the setup steps at the top of google_drive_backup.py\n"
            "  2. Add the base64 key as an HF Space Secret named GDRIVE_SA_KEY"
        )

    # Decode base64 → JSON
    try:
        key_json = base64.b64decode(raw).decode("utf-8")
        key_dict = json.loads(key_json)
    except Exception as e:
        raise ValueError(f"Failed to decode GDRIVE_SA_KEY — make sure it's base64-encoded JSON: {e}")

    # Write to temp file (google-auth library needs a file path)
    with open(SA_KEY_FILE, "w") as f:
        json.dump(key_dict, f)

    scopes = ["https://www.googleapis.com/auth/drive"]
    creds = service_account.Credentials.from_service_account_file(SA_KEY_FILE, scopes=scopes)
    return creds


def get_or_create_folder(service, folder_name):
    """Find or create the backup folder in Drive, return its ID."""
    query = (
        f"name='{folder_name}' "
        f"and mimeType='application/vnd.google-apps.folder' "
        f"and trashed=false"
    )
    results = service.files().list(q=query, spaces="drive", fields="files(id, name)").execute()
    folders = results.get("files", [])

    if folders:
        log.info(f"Using existing Drive folder: {folder_name} (ID: {folders[0]['id']})")
        return folders[0]["id"]

    log.info(f"Creating new Drive folder: {folder_name}")
    metadata = {
        "name": folder_name,
        "mimeType": "application/vnd.google-apps.folder"
    }
    folder = service.files().create(body=metadata, fields="id").execute()
    return folder.get("id")


def prune_old_backups(service, folder_id):
    """Delete oldest backups beyond MAX_BACKUPS."""
    results = service.files().list(
        q=f"'{folder_id}' in parents and trashed=false",
        orderBy="createdTime desc",
        fields="files(id, name, createdTime)"
    ).execute()
    files = results.get("files", [])

    if len(files) > MAX_BACKUPS:
        to_delete = files[MAX_BACKUPS:]
        for f in to_delete:
            service.files().delete(fileId=f["id"]).execute()
            log.info(f"Deleted old backup: {f['name']}")
    else:
        log.info(f"Backup count: {len(files)}/{MAX_BACKUPS} — no pruning needed.")


def upload_backup(backup_path: str):
    from googleapiclient.discovery import build
    from googleapiclient.http import MediaFileUpload

    path = Path(backup_path)
    if not path.exists():
        raise FileNotFoundError(f"Backup file not found: {backup_path}")

    size_mb = path.stat().st_size / 1024 / 1024
    log.info(f"Uploading {path.name} ({size_mb:.1f} MB) to Google Drive...")

    creds   = load_service_account_credentials()
    service = build("drive", "v3", credentials=creds)

    folder_id = get_or_create_folder(service, BACKUP_FOLDER_NAME)

    media = MediaFileUpload(
        str(path),
        mimetype="application/gzip",
        resumable=True,      # resumable upload handles large files reliably
        chunksize=5 * 1024 * 1024  # 5 MB chunks
    )
    file_metadata = {"name": path.name, "parents": [folder_id]}
    request = service.files().create(body=file_metadata, media_body=media, fields="id, name, size")

    # Show upload progress
    response = None
    while response is None:
        status, response = request.next_chunk()
        if status:
            pct = int(status.progress() * 100)
            log.info(f"  Upload progress: {pct}%")

    log.info(f"Upload complete! File ID: {response.get('id')}")
    prune_old_backups(service, folder_id)

    # Clean up temp credentials file
    try:
        os.remove(SA_KEY_FILE)
    except Exception:
        pass


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: python3 google_drive_backup.py <path_to_backup.tar.gz>")
        sys.exit(1)

    try:
        upload_backup(sys.argv[1])
    except EnvironmentError as e:
        log.error(str(e))
        sys.exit(2)
    except Exception as e:
        log.error(f"Upload failed: {e}")
        sys.exit(1)