src/colab_utils.py

"""Helpers for running notebooks on Google Colab Pro+.

Each Colab notebook in ``notebooks/colab/`` calls :func:`setup_colab` in its
first cell to bootstrap the environment (mount Drive, install missing pip
deps, symlink data files into the cloned repo, configure git for auto-push).

After training cells complete and the user has eyeballed the output,
they manually run a final cell calling :func:`push_results` to commit
``results/*.json`` and push to ``main``.

The notebook ``.ipynb`` file itself is pushed via Colab's built-in
**File → Save a copy in GitHub** (uses Colab's own GitHub OAuth, no PAT
needed).
"""

from __future__ import annotations

import logging
import subprocess
import sys
from pathlib import Path

logger = logging.getLogger(__name__)

DRIVE_MOUNT = Path("/content/drive")
DEFAULT_DRIVE_ROOT = DRIVE_MOUNT / "MyDrive" / "thesis-cyberbullying"
DEFAULT_BRANCH = "main"
EXTRA_PIP_DEPS = ("Sastrawi==1.0.1", "sentencepiece", "gensim>=4.0.0")


def is_colab() -> bool:
    """True when running inside a Google Colab kernel."""
    return "google.colab" in sys.modules


def _run(cmd: list[str], cwd: str | Path | None = None, capture: bool = False):
    """Run a shell command, raising on non-zero exit."""
    return subprocess.run(
        cmd, cwd=str(cwd) if cwd else None, check=True,
        capture_output=capture, text=True,
    )


def mount_drive() -> Path:
    """Mount Google Drive at ``/content/drive`` (idempotent). Return ``MyDrive`` path."""
    if not is_colab():
        return DRIVE_MOUNT / "MyDrive"
    if not (DRIVE_MOUNT / "MyDrive").exists():
        from google.colab import drive  # noqa: WPS433
        drive.mount(str(DRIVE_MOUNT))
    return DRIVE_MOUNT / "MyDrive"


def install_deps(pkgs: tuple[str, ...] = EXTRA_PIP_DEPS) -> None:
    """``pip install -q`` packages Colab doesn't preinstall (Sastrawi, sentencepiece)."""
    if not is_colab():
        return
    _run([sys.executable, "-m", "pip", "install", "-q", *pkgs])


def link_data_from_drive(
    drive_root: Path = DEFAULT_DRIVE_ROOT,
    repo_root: Path | None = None,
) -> list[str]:
    """Symlink ``drive_root/data/*`` into ``repo_root/data/raw/``.

    Idempotent - existing symlinks/files are skipped.
    """
    if repo_root is None:
        repo_root = Path.cwd()
    drive_data = drive_root / "data"
    if not drive_data.exists():
        raise FileNotFoundError(
            f"Expected data files at {drive_data}. "
            "Upload Dataset_Komentar_Tiktok.csv, cc.id.300.bin, "
            "and slang JSON dicts there. See README.md § Colab workflow."
        )
    target = repo_root / "data" / "raw"
    target.mkdir(parents=True, exist_ok=True)
    linked: list[str] = []
    for src in sorted(drive_data.iterdir()):
        link = target / src.name
        if link.exists() or link.is_symlink():
            continue
        link.symlink_to(src)
        linked.append(src.name)
    return linked


def configure_git(
    github_pat: str,
    repo_root: Path | None = None,
    user_email: str = "colab@cyberbullying-detection",
    user_name: str = "Colab Auto-Push",
) -> None:
    """Configure git for auto-push from Colab using a Personal Access Token."""
    if repo_root is None:
        repo_root = Path.cwd()
    _run(["git", "config", "user.email", user_email], cwd=repo_root)
    _run(["git", "config", "user.name", user_name], cwd=repo_root)
    res = _run(["git", "remote", "get-url", "origin"], cwd=repo_root, capture=True)
    url = res.stdout.strip()
    if "@github.com" not in url:
        new_url = url.replace("https://github.com/", f"https://{github_pat}@github.com/")
        _run(["git", "remote", "set-url", "origin", new_url], cwd=repo_root)


def setup_colab(
    drive_root: Path = DEFAULT_DRIVE_ROOT,
    github_pat: str | None = None,
) -> dict:
    """One-shot bootstrap. Call this in the first cell of every Colab notebook.

    Returns a paths dict with ``colab``, ``repo_root``, ``drive_root``,
    ``fasttext_path``, ``dataset_csv``, ``linked``.
    """
    if not is_colab():
        return {"colab": False, "repo_root": str(Path.cwd())}

    mount_drive()
    install_deps()

    repo_root = Path.cwd()
    linked = link_data_from_drive(drive_root, repo_root)

    if github_pat:
        configure_git(github_pat, repo_root)

    fasttext_path = drive_root / "data" / "cc.id.300.bin"
    dataset_csv = repo_root / "data" / "raw" / "Dataset_Komentar_Tiktok.csv"
    # Model checkpoints persist to Drive so they survive Colab session resets
    # and sync to the Windows host via Drive Desktop. App.py loads from the
    # local repo `models/` - copy/symlink from `<drive_root>/models/` after sync.
    models_drive_dir = drive_root / "models"
    models_drive_dir.mkdir(parents=True, exist_ok=True)

    return {
        "colab": True,
        "repo_root": str(repo_root),
        "drive_root": str(drive_root),
        "fasttext_path": str(fasttext_path),
        "dataset_csv": str(dataset_csv),
        "models_dir": str(models_drive_dir),
        "linked": linked,
    }


def resolve_models_dir(paths: dict | None = None) -> Path:
    """Return the directory where model checkpoints should be saved.

    In Colab, returns ``<drive_root>/models/`` (from the ``setup_colab`` paths
    dict) so checkpoints survive session resets and sync to the host via Drive
    Desktop. Locally, returns the repo's ``models/`` directory.
    """
    if paths and paths.get("colab") and paths.get("models_dir"):
        d = Path(paths["models_dir"])
    else:
        from src.paths import MODELS_DIR
        d = MODELS_DIR
    d.mkdir(parents=True, exist_ok=True)
    return d


def push_results(
    commit_message: str,
    paths: list[str] | None = None,
    branch: str = DEFAULT_BRANCH,
    repo_root: Path | None = None,
) -> bool:
    """Stage, commit, and push specified paths to ``branch``.

    Order: stage → commit → fetch+rebase → push. The fetch+rebase happens AFTER
    commit so the working tree is clean (notebook outputs that overwrote
    tracked files would otherwise cause ``pull --rebase`` to refuse with
    "you have unstaged changes" / fatal exit 128).

    Returns True if a push happened, False if there was nothing to commit.
    """
    if not is_colab():
        print("Not in Colab - push_results is a no-op locally.")
        return False
    if repo_root is None:
        repo_root = Path.cwd()
    paths = paths or ["results/"]

    # 1. Stage
    _run(["git", "add", *paths], cwd=repo_root)

    # 2. Commit (if anything to commit). Skip the rest only if there's truly
    #    nothing local - a no-op return is fine; nothing was changed.
    diff = subprocess.run(
        ["git", "diff", "--cached", "--quiet"], cwd=str(repo_root),
    )
    if diff.returncode == 0:
        print("No changes to commit.")
        return False
    _run(["git", "commit", "-m", commit_message], cwd=repo_root)

    # 3. Fetch + rebase onto remote (handles parallel pushes from other notebooks).
    #    With a clean tree (we just committed), this won't fail on dirty-tree.
    try:
        _run(["git", "fetch", "origin", branch], cwd=repo_root)
        _run(["git", "rebase", f"origin/{branch}"], cwd=repo_root)
    except subprocess.CalledProcessError as exc:
        print(
            f"WARN: rebase onto origin/{branch} failed ({exc}). "
            "Continuing - will rely on push to detect non-fast-forward."
        )

    # 4. Push
    _run(["git", "push", "origin", branch], cwd=repo_root)
    print(f"Pushed: {commit_message}")
    return True