src/easytranslate/utils/cloud_storage.py

"""
Cloud storage integration module for Google Colab and Google Drive.

Provides secure, authenticated storage operations for training results,
checkpoints, and logs. Supports automatic environment detection for
Colab vs local execution.
"""

from __future__ import annotations

import json
import logging
import os
import shutil
import time
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)


def is_colab_environment() -> bool:
    try:
        import google.colab
        return True
    except ImportError:
        return False


def mount_google_drive(mount_point: str = "/content/drive") -> bool:
    if not is_colab_environment():
        logger.info("Not running in Colab, skipping Google Drive mount")
        return False

    try:
        from google.colab import drive
        drive.mount(mount_point)
        logger.info("Google Drive mounted at %s", mount_point)
        return True
    except Exception as e:
        logger.error("Failed to mount Google Drive: %s", e)
        return False


def get_drive_path(base_path: str = "/content/drive/MyDrive/EasyTranslate") -> Optional[Path]:
    if not os.path.exists("/content/drive"):
        logger.warning("Google Drive not mounted, cannot resolve drive path")
        return None

    drive_path = Path(base_path)
    drive_path.mkdir(parents=True, exist_ok=True)
    return drive_path


def sync_checkpoints_to_drive(
    local_checkpoint_dir: str | Path,
    drive_base_path: str = "/content/drive/MyDrive/EasyTranslate",
    max_retries: int = 3,
) -> bool:
    local_dir = Path(local_checkpoint_dir)
    if not local_dir.exists():
        logger.warning("Local checkpoint directory does not exist: %s", local_dir)
        return False

    drive_path = get_drive_path(drive_base_path)
    if drive_path is None:
        return False

    drive_checkpoint_dir = drive_path / "checkpoints"
    drive_checkpoint_dir.mkdir(parents=True, exist_ok=True)

    success = True
    for ckpt_file in local_dir.glob("*.pt"):
        dest = drive_checkpoint_dir / ckpt_file.name
        for attempt in range(max_retries):
            try:
                shutil.copy2(ckpt_file, dest)
                logger.info("Synced checkpoint to Drive: %s", dest)
                break
            except Exception as e:
                logger.warning("Sync attempt %d/%d failed for %s: %s", attempt + 1, max_retries, ckpt_file.name, e)
                if attempt == max_retries - 1:
                    success = False
                time.sleep(2 ** attempt)

    return success


def sync_logs_to_drive(
    local_log_dir: str | Path,
    drive_base_path: str = "/content/drive/MyDrive/EasyTranslate",
    max_retries: int = 3,
) -> bool:
    local_dir = Path(local_log_dir)
    if not local_dir.exists():
        logger.warning("Local log directory does not exist: %s", local_dir)
        return False

    drive_path = get_drive_path(drive_base_path)
    if drive_path is None:
        return False

    drive_log_dir = drive_path / "logs"
    drive_log_dir.mkdir(parents=True, exist_ok=True)

    success = True
    for log_file in local_dir.glob("*"):
        if log_file.is_file():
            dest = drive_log_dir / log_file.name
            for attempt in range(max_retries):
                try:
                    shutil.copy2(log_file, dest)
                    break
                except Exception as e:
                    logger.warning("Log sync attempt %d/%d failed: %s", attempt + 1, max_retries, e)
                    if attempt == max_retries - 1:
                        success = False
                    time.sleep(2 ** attempt)

    return success


def save_training_summary_to_drive(
    summary: dict,
    drive_base_path: str = "/content/drive/MyDrive/EasyTranslate",
) -> bool:
    drive_path = get_drive_path(drive_base_path)
    if drive_path is None:
        return False

    summary_path = drive_path / "training_summary.json"
    try:
        with open(summary_path, "w", encoding="utf-8") as f:
            json.dump(summary, f, indent=2, ensure_ascii=False, default=str)
        logger.info("Training summary saved to Drive: %s", summary_path)
        return True
    except Exception as e:
        logger.error("Failed to save training summary to Drive: %s", e)
        return False


def sync_all_to_drive(
    checkpoint_dir: str | Path = "checkpoints",
    log_dir: str | Path = "logs",
    drive_base_path: str = "/content/drive/MyDrive/EasyTranslate",
) -> dict:
    results = {
        "checkpoints_synced": sync_checkpoints_to_drive(checkpoint_dir, drive_base_path),
        "logs_synced": sync_logs_to_drive(log_dir, drive_base_path),
    }

    summary_path = Path(checkpoint_dir) / "training_summary.json"
    if summary_path.exists():
        with open(summary_path, "r", encoding="utf-8") as f:
            summary = json.load(f)
        results["summary_saved"] = save_training_summary_to_drive(summary, drive_base_path)

    logger.info("Drive sync results: %s", results)
    return results


def setup_colab_environment() -> dict:
    env_info = {
        "is_colab": is_colab_environment(),
        "gpu_available": False,
        "gpu_info": None,
        "drive_mounted": False,
        "drive_path": None,
    }

    if env_info["is_colab"]:
        env_info["drive_mounted"] = mount_google_drive()
        env_info["drive_path"] = str(get_drive_path()) if env_info["drive_mounted"] else None

    try:
        import torch
        env_info["gpu_available"] = torch.cuda.is_available()
        if env_info["gpu_available"]:
            env_info["gpu_info"] = {
                "device_count": torch.cuda.device_count(),
                "device_name": torch.cuda.get_device_name(0),
                "device_capability": torch.cuda.get_device_capability(0),
            }
    except ImportError:
        pass

    logger.info("Environment setup: %s", {k: v for k, v in env_info.items() if k != "gpu_info"})
    return env_info