#!/usr/bin/env python3
"""
Google Colab Environment Setup Script for EasyTranslate.

This script is designed to be run as the first cell in a Colab notebook.
It handles:
1. GPU verification and CUDA setup
2. Repository cloning from GitHub
3. Dependency installation
4. Google Drive mounting
5. Environment variable configuration

Usage in Colab:
    !wget -q https://raw.githubusercontent.com/your-org/UCAS-EasyTranslate/main/scripts/setup_colab.py
    %run setup_colab.py
"""

import os
import subprocess
import sys
from pathlib import Path

REPO_URL = "https://github.com/your-org/UCAS-EasyTranslate.git"
REPO_DIR = Path("/content/UCAS-EasyTranslate")
DRIVE_MOUNT_POINT = "/content/drive"
DRIVE_BASE = "/content/drive/MyDrive/EasyTranslate"


def check_gpu():
    try:
        import torch
        if torch.cuda.is_available():
            print(f"[OK] GPU detected: {torch.cuda.get_device_name(0)}")
            print(f"     CUDA version: {torch.version.cuda}")
            print(f"     GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")
            return True
        else:
            print("[WARN] No GPU detected. Training will be very slow.")
            return False
    except ImportError:
        print("[WARN] PyTorch not found. Installing...")
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "torch"])
        return check_gpu()


def clone_repository():
    if REPO_DIR.exists():
        print(f"[INFO] Repository exists at {REPO_DIR}, pulling latest...")
        os.chdir(REPO_DIR)
        subprocess.check_call(["git", "pull", "origin", "main"])
    else:
        print(f"[INFO] Cloning repository from {REPO_URL}...")
        subprocess.check_call(["git", "clone", REPO_URL, str(REPO_DIR)])
    os.chdir(REPO_DIR)
    sys.path.insert(0, str(REPO_DIR / "src"))
    print(f"[OK] Repository ready at {REPO_DIR}")


def install_dependencies():
    req_file = REPO_DIR / "requirements-colab.txt"
    if req_file.exists():
        print("[INFO] Installing Colab dependencies...")
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "-r", str(req_file)])
    else:
        print("[INFO] Installing from requirements.txt...")
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "-r", str(REPO_DIR / "requirements.txt")])

    print("[INFO] Installing EasyTranslate package...")
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "-e", str(REPO_DIR)])
    print("[OK] All dependencies installed")


def mount_drive():
    try:
        from google.colab import drive
        drive.mount(DRIVE_MOUNT_POINT)
        if os.path.exists(DRIVE_MOUNT_POINT):
            os.makedirs(DRIVE_BASE, exist_ok=True)
            print(f"[OK] Google Drive mounted at {DRIVE_MOUNT_POINT}")
            print(f"     Base path: {DRIVE_BASE}")
            return True
    except Exception as e:
        print(f"[WARN] Google Drive mount failed: {e}")
    return False


def setup_environment():
    os.environ["TOKENIZERS_PARALLELISM"] = "false"
    os.environ["WANDB_MODE"] = os.environ.get("WANDB_MODE", "offline")
    os.environ["PYTHONHASHSEED"] = "42"
    print("[OK] Environment variables configured")


def main():
    print("=" * 60)
    print("  EasyTranslate — Colab Environment Setup")
    print("=" * 60)
    print()

    check_gpu()
    print()

    clone_repository()
    print()

    install_dependencies()
    print()

    mount_drive()
    print()

    setup_environment()
    print()

    print("=" * 60)
    print("  Setup Complete! Ready for training.")
    print("=" * 60)


if __name__ == "__main__":
    main()