backup_module.py

"""
Backup Module - Export project ke Google Drive tanpa zip.
Backup SMART: cuma file penting (models, data, config, code).

Fitur:
- Backup selective (essential / models only / data only / full)
- Restore dari backup
- Import dataset dari Drive ke /content/data/
- List & clean backup lama
"""

import os
import json
import shutil
from datetime import datetime
from pathlib import Path
from typing import Optional, Set, List, Dict, Tuple

from rich.table import Table
from rich.panel import Panel
from rich.prompt import Prompt, Confirm
from rich.progress import (
    Progress, SpinnerColumn, TextColumn, BarColumn,
    TimeElapsedColumn, TimeRemainingColumn,
)
from rich import box

from utils import console, Theme, error_logger, debug_logger


# ======================================================================
# Konfigurasi
# ======================================================================
DEFAULT_SOURCE = Path("/content")
DEFAULT_DEST_BASE = Path("/content/drive/MyDrive/colab_backup")

# Folder/file yang SELALU di-skip (biar backup tidak bengkak)
SKIP_NAMES: Set[str] = {
    "drive",
    "sample_data",
    "__pycache__",
    ".ipynb_checkpoints",
    ".git",
    ".config",
    ".cache",
    "checkpoints",   # checkpoint training - bisa GB
    "temp",          # file temporer
    "logs",          # log files
    "backups",
    "exports",
}

SKIP_EXTENSIONS: Set[str] = {
    ".pyc", ".pyo", ".log", ".tmp",
}

# Mode backup -> item yang di-include
BACKUP_MODES = {
    "essential": {
        "label": "Essential",
        "desc": "config.json + *.py + models/ + data/",
        "files": ["config.json"],
        "globs": ["*.py"],
        "dirs": ["models", "data"],
    },
    "models_only": {
        "label": "Models only",
        "desc": "config.json + models/",
        "files": ["config.json"],
        "globs": [],
        "dirs": ["models"],
    },
    "data_only": {
        "label": "Data only",
        "desc": "data/ (dataset + KB + quizzes)",
        "files": [],
        "globs": [],
        "dirs": ["data"],
    },
    "full": {
        "label": "Full project",
        "desc": "Semua isi /content (kecuali folder besar)",
        "files": [],
        "globs": [],
        "dirs": [],
    },
}


# ======================================================================
# Helper
# ======================================================================
def _format_size(num_bytes: int) -> str:
    for unit in ["B", "KB", "MB", "GB", "TB"]:
        if num_bytes < 1024:
            return f"{num_bytes:.2f} {unit}"
        num_bytes /= 1024
    return f"{num_bytes:.2f} PB"


def _should_skip(path: Path, skip_names: Set[str], skip_exts: Set[str]) -> bool:
    if path.name in skip_names:
        return True
    if path.is_file() and path.suffix.lower() in skip_exts:
        return True
    return False


# ======================================================================
# BackupManager
# ======================================================================
class BackupManager:
    def __init__(
        self,
        source: Path = DEFAULT_SOURCE,
        dest_base: Path = DEFAULT_DEST_BASE,
    ):
        self.source = Path(source)
        self.dest_base = Path(dest_base)

    # ------------------------------------------------------------------
    def mount_drive(self) -> bool:
        if Path("/content/drive/MyDrive").exists():
            console.print(Theme.dim("✓ Google Drive sudah ter-mount"))
            return True
        try:
            from google.colab import drive
            console.print(Theme.info("Mounting Google Drive..."))
            console.print(Theme.dim("(Ikuti link otorisasi)"))
            drive.mount("/content/drive")
            return Path("/content/drive/MyDrive").exists()
        except ImportError:
            console.print(Theme.warning("⚠ Bukan Colab — lewati mount"))
            self.dest_base = Path("/content/colab_backup")
            self.dest_base.mkdir(parents=True, exist_ok=True)
            return True
        except Exception as e:
            console.print(Theme.error(f"✗ Gagal mount Drive: {e}"))
            return False

    # ------------------------------------------------------------------
    def _collect_files(self, mode: str) -> List[Path]:
        """Kumpulkan file berdasarkan mode backup."""
        files: List[Path] = []

        if mode == "full":
            for item in self.source.rglob("*"):
                if not item.is_file():
                    continue
                if _should_skip(item, SKIP_NAMES, SKIP_EXTENSIONS):
                    continue
                try:
                    item.resolve().relative_to(DEFAULT_DEST_BASE.resolve())
                    continue
                except (ValueError, OSError):
                    pass
                files.append(item)
            return files

        cfg = BACKUP_MODES.get(mode)
        if cfg is None:
            return files

        # File langsung
        for fname in cfg["files"]:
            f = self.source / fname
            if f.exists() and f.is_file():
                files.append(f)

        # Glob patterns
        for pattern in cfg["globs"]:
            for f in self.source.glob(pattern):
                if f.is_file():
                    files.append(f)

        # Folder
        for dname in cfg["dirs"]:
            d = self.source / dname
            if not d.exists():
                continue
            if d.is_file():
                files.append(d)
            else:
                for f in d.rglob("*"):
                    if not f.is_file():
                        continue
                    if _should_skip(f, SKIP_NAMES, SKIP_EXTENSIONS):
                        continue
                    files.append(f)

        # Dedup
        seen = set()
        unique = []
        for f in files:
            key = str(f)
            if key not in seen:
                seen.add(key)
                unique.append(f)
        return unique

    # ------------------------------------------------------------------
    def export(self, mode: str = "essential", dry_run: bool = False) -> Optional[Path]:
        """Backup dengan mode tertentu."""
        if mode not in BACKUP_MODES:
            console.print(Theme.error(f"Mode tidak dikenal: {mode}"))
            return None

        if not self.mount_drive():
            return None

        # Kumpulkan file
        console.print(Theme.info(f"Scanning files (mode: {mode})..."))
        files = self._collect_files(mode)

        if not files:
            console.print(Theme.warning("Tidak ada file untuk di-backup"))
            return None

        n_files = len(files)
        total_bytes = 0
        for f in files:
            try:
                total_bytes += f.stat().st_size
            except Exception:
                pass

        mode_info = BACKUP_MODES[mode]
        console.print(Panel(
            f"[cyan]Mode       :[/cyan] {mode_info['label']}\n"
            f"[cyan]Desc       :[/cyan] {mode_info['desc']}\n"
            f"[cyan]Files      :[/cyan] {n_files:,} files\n"
            f"[cyan]Total size :[/cyan] {_format_size(total_bytes)}\n"
            f"[cyan]Destination:[/cyan] {self.dest_base}",
            title="[bold yellow] BACKUP PREVIEW",
            style="yellow",
        ))

        if dry_run:
            console.print(Theme.dim("(dry run - tidak ada file yang di-copy)"))
            return None

        if not Confirm.ask("[yellow]Lanjutkan backup?", default=True):
            console.print(Theme.info("Dibatalkan"))
            return None

        # Buat folder backup
        ts = datetime.now().strftime("%Y%m%d_%H%M%S")
        dest = self.dest_base / f"{mode}_{ts}"
        dest.mkdir(parents=True, exist_ok=True)

        console.print(Theme.info(f"Backup ke: {dest}"))

        # Copy
        counter = {"copied": 0, "errors": 0, "bytes": 0}
        start = datetime.now()

        try:
            with Progress(
                SpinnerColumn(),
                TextColumn("[progress.description]{task.description}"),
                BarColumn(),
                TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
                TimeElapsedColumn(),
                TimeRemainingColumn(),
            ) as progress:
                task = progress.add_task("[cyan]Copying...", total=len(files))
                for f in files:
                    try:
                        rel = f.relative_to(self.source)
                    except ValueError:
                        progress.update(task, advance=1)
                        continue
                    target = dest / rel
                    target.parent.mkdir(parents=True, exist_ok=True)
                    try:
                        shutil.copy2(f, target)
                        counter["copied"] += 1
                        counter["bytes"] += f.stat().st_size
                    except (PermissionError, OSError) as e:
                        counter["errors"] += 1
                        debug_logger.debug(f"Skip {f}: {e}")
                    progress.update(task, advance=1)
        except KeyboardInterrupt:
            console.print(Theme.warning("\n⚠ Backup dihentikan"))
            return dest
        except Exception as e:
            error_logger.error(f"Backup error: {e}")
            console.print(Theme.error(f"✗ Backup gagal: {e}"))
            return None

        elapsed = (datetime.now() - start).total_seconds()

        # Tulis manifest
        manifest = {
            "mode": mode,
            "label": mode_info["label"],
            "created": datetime.now().isoformat(),
            "source": str(self.source),
            "files": counter["copied"],
            "bytes": counter["bytes"],
            "errors": counter["errors"],
            "elapsed_sec": elapsed,
        }
        try:
            with open(dest / "MANIFEST.json", "w") as f:
                json.dump(manifest, f, indent=2)
        except Exception:
            pass

        table = Table(title="Backup Result", box=box.ROUNDED)
        table.add_column("Item", style="cyan")
        table.add_column("Value", style="green")
        table.add_row("✓ File di-copy", f"{counter['copied']:,}")
        table.add_row("✗ Error", f"{counter['errors']:,}")
        table.add_row("📦 Total size", _format_size(counter['bytes']))
        table.add_row("⏱ Waktu", f"{elapsed:.1f} detik")
        table.add_row("📁 Lokasi", str(dest))
        console.print(table)
        console.print(Theme.success(f" ✓ Backup selesai: {dest.name}"))

        return dest

    # ------------------------------------------------------------------
    def restore(self, backup_path: Optional[Path] = None) -> bool:
        """Restore dari backup ke /content."""
        if backup_path is None:
            backups = self._list_backup_paths()
            if not backups:
                console.print(Theme.warning("Tidak ada backup tersedia"))
                return False

            console.print(Theme.info("Pilih backup yang mau di-restore:"))
            for i, b in enumerate(backups, 1):
                size = sum(f.stat().st_size for f in b.rglob("*") if f.is_file())
                mode = b.name.rsplit("_", 2)[0] if "_" in b.name else "?"
                print(f"  {i}. [{mode}] {b.name}  ({_format_size(size)})")

            pilih = Prompt.ask("[yellow]Nomor", default="1")
            try:
                backup_path = backups[int(pilih) - 1]
            except (ValueError, IndexError):
                console.print(Theme.error("Pilihan invalid"))
                return False

        backup_path = Path(backup_path)
        if not backup_path.exists():
            console.print(Theme.error(f"Backup tidak ada: {backup_path}"))
            return False

        console.print(Panel(
            f"[cyan]Backup :[/cyan] {backup_path.name}\n"
            f"[cyan]Target :[/cyan] {self.source}\n\n"
            "[yellow]⚠ File yang ada di target dengan nama sama akan DITIMPA[/yellow]",
            title="[bold yellow] RESTORE",
            style="yellow",
        ))

        if not Confirm.ask("[yellow]Lanjutkan restore?", default=True):
            console.print(Theme.info("Dibatalkan"))
            return False

        counter = {"copied": 0, "errors": 0, "bytes": 0}
        start = datetime.now()

        try:
            with Progress(
                SpinnerColumn(),
                TextColumn("[progress.description]{task.description}"),
                BarColumn(),
                TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
                TimeElapsedColumn(),
                TimeRemainingColumn(),
            ) as progress:
                all_files = [f for f in backup_path.rglob("*") if f.is_file()]
                task = progress.add_task(
                    "[cyan]Restoring...",
                    total=max(1, len(all_files)),
                )
                for f in all_files:
                    if f.name == "MANIFEST.json":
                        progress.update(task, advance=1)
                        continue
                    try:
                        rel = f.relative_to(backup_path)
                    except ValueError:
                        progress.update(task, advance=1)
                        continue
                    target = self.source / rel
                    target.parent.mkdir(parents=True, exist_ok=True)
                    try:
                        shutil.copy2(f, target)
                        counter["copied"] += 1
                        counter["bytes"] += f.stat().st_size
                    except (PermissionError, OSError) as e:
                        counter["errors"] += 1
                        debug_logger.debug(f"Skip {f}: {e}")
                    progress.update(task, advance=1)
        except Exception as e:
            error_logger.error(f"Restore error: {e}")
            console.print(Theme.error(f"✗ Restore gagal: {e}"))
            return False

        elapsed = (datetime.now() - start).total_seconds()

        table = Table(title="Restore Result", box=box.ROUNDED)
        table.add_column("Item", style="cyan")
        table.add_column("Value", style="green")
        table.add_row("✓ File di-restore", f"{counter['copied']:,}")
        table.add_row("✗ Error", f"{counter['errors']:,}")
        table.add_row("📦 Size", _format_size(counter['bytes']))
        table.add_row("⏱ Waktu", f"{elapsed:.1f} detik")
        table.add_row("📁 Target", str(self.source))
        console.print(table)

        self._verify_restore()
        return True

    # ------------------------------------------------------------------
    def _verify_restore(self) -> None:
        """Cek hasil restore — models, config, dataset."""
        console.print("\n" + Theme.header(" Verifikasi Restore:"))

        models_dir = self.source / "models"
        if models_dir.exists() and any(models_dir.iterdir()):
            console.print(Theme.success("✓ Folder /content/models/"))
            for m in models_dir.iterdir():
                if m.is_dir():
                    size = sum(f.stat().st_size for f in m.rglob("*") if f.is_file())
                    console.print(f"    - {m.name}  ({_format_size(size)})")
        else:
            console.print(Theme.warning("⚠ /content/models/ kosong"))

        cfg_path = self.source / "config.json"
        if cfg_path.exists():
            try:
                with open(cfg_path) as f:
                    cfg = json.load(f)
                models = list(cfg.get("models", {}).keys())
                current = cfg.get("current_model")
                console.print(Theme.success("✓ /content/config.json"))
                console.print(f"    Models: {models}")
                console.print(f"    Current: {current}")
            except Exception as e:
                console.print(Theme.warning(f"⚠ config.json rusak: {e}"))
        else:
            console.print(Theme.warning("⚠ /content/config.json tidak ada"))

        data_dir = self.source / "data"
        if data_dir.exists() and any(data_dir.iterdir()):
            console.print(Theme.success("✓ Folder /content/data/"))
            for f in data_dir.iterdir():
                if f.is_file():
                    console.print(f"    - {f.name}  ({_format_size(f.stat().st_size)})")
                elif f.is_dir():
                    n = sum(1 for _ in f.rglob("*") if _.is_file())
                    console.print(f"    - {f.name}/  ({n} files)")

    # ------------------------------------------------------------------
    def import_dataset_from_drive(self) -> None:
        """Import dataset dari Drive ke /content/data/."""
        console.print(Panel(
            Theme.header(" IMPORT DATASET DARI DRIVE"),
            title="IMPORT", style="bold yellow",
        ))

        search_dirs = [
            Path("/content/drive/MyDrive/colab_backup"),
            Path("/content/drive/MyDrive"),
        ]

        supported_ext = (
            ".txt", ".json", ".jsonl", ".csv", ".tsv",
            ".parquet", ".arrow", ".json.gz", ".jsonl.gz",
        )

        found_files: List[Path] = []
        seen = set()
        for d in search_dirs:
            if not d.exists():
                continue
            for ext in supported_ext:
                for f in d.rglob(f"*{ext}"):
                    if f.is_file() and f.stat().st_size > 0:
                        key = str(f.resolve())
                        if key not in seen:
                            seen.add(key)
                            found_files.append(f)
            if found_files:
                break

        if not found_files:
            console.print(Theme.warning("Tidak ada file dataset di Drive"))
            return

        console.print(Theme.info("File dataset di Drive:"))
        for i, f in enumerate(found_files, 1):
            size = f.stat().st_size
            console.print(f"  {i}. {f.name}  ({_format_size(size)})")
            console.print(f"     [dim]{f.parent}[/dim]")

        pilih = Prompt.ask("[yellow]Pilih nomor", default="1")
        try:
            src = found_files[int(pilih) - 1]
        except (ValueError, IndexError):
            console.print(Theme.error("Pilihan invalid"))
            return

        dest_dir = self.source / "data"
        dest_dir.mkdir(parents=True, exist_ok=True)
        dest = dest_dir / src.name

        try:
            shutil.copy2(src, dest)
            console.print(Theme.success(f"✓ Dataset di-copy ke {dest}"))
            console.print(Theme.dim(f"  Size: {_format_size(src.stat().st_size)}"))
        except Exception as e:
            console.print(Theme.error(f"✗ Gagal copy: {e}"))

    # ------------------------------------------------------------------
    def _list_backup_paths(self) -> List[Path]:
        if not self.dest_base.exists():
            return []
        return sorted(
            [d for d in self.dest_base.iterdir()
             if d.is_dir() and not d.name.startswith(".")],
            key=lambda d: d.stat().st_mtime,
            reverse=True,
        )

    def list_backups(self) -> List[Path]:
        backups = self._list_backup_paths()
        if not backups:
            console.print(Theme.warning("Belum ada backup"))
            return []

        table = Table(title=f"Backups di {self.dest_base}", box=box.ROUNDED)
        table.add_column("No", style="cyan")
        table.add_column("Nama", style="green")
        table.add_column("Mode", style="magenta")
        table.add_column("Size", style="yellow")
        table.add_column("Files", style="blue")
        table.add_column("Waktu", style="dim")

        for idx, b in enumerate(backups, 1):
            try:
                files = [f for f in b.rglob("*") if f.is_file()]
                size = sum(f.stat().st_size for f in files)
                mtime = datetime.fromtimestamp(b.stat().st_mtime).strftime("%Y-%m-%d %H:%M")
                mode = b.name.rsplit("_", 2)[0] if "_" in b.name else "?"
                table.add_row(
                    str(idx), b.name, mode,
                    _format_size(size), f"{len(files):,}", mtime,
                )
            except Exception:
                table.add_row(str(idx), b.name, "?", "?", "?", "?")

        console.print(table)
        return backups

    # ------------------------------------------------------------------
    def clean_old_backups(self, keep: int = 3) -> None:
        backups = self._list_backup_paths()
        to_delete = backups[keep:]

        if not to_delete:
            console.print(Theme.success(f"✓ Tidak ada backup lama (total: {len(backups)})"))
            return

        console.print(Theme.warning(f"Akan menghapus {len(to_delete)} backup:"))
        for old in to_delete:
            console.print(f"  - {old.name}")

        if not Confirm.ask(Theme.error(f"Hapus {len(to_delete)} backup?"), default=False):
            console.print(Theme.info("Dibatalkan"))
            return

        for old in to_delete:
            try:
                shutil.rmtree(old)
                console.print(Theme.success(f"  ✓ {old.name} dihapus"))
            except Exception as e:
                console.print(Theme.error(f"  ✗ {old.name}: {e}"))

    # ------------------------------------------------------------------
    def show_sizes(self) -> None:
        """Tampilkan ukuran folder penting di /content."""
        table = Table(title="Ukuran folder di /content", box=box.ROUNDED)
        table.add_column("Folder/File", style="cyan")
        table.add_column("Size", style="green")
        table.add_column("Files", style="blue")

        items = [
            "models", "data", "checkpoints", "logs", "temp",
            "cache", "backups", "config.json",
        ]
        for name in items:
            p = self.source / name
            if not p.exists():
                continue
            if p.is_file():
                table.add_row(name, _format_size(p.stat().st_size), "1")
            else:
                files = [f for f in p.rglob("*") if f.is_file()]
                size = sum(f.stat().st_size for f in files)
                table.add_row(name, _format_size(size), f"{len(files):,}")

        console.print(table)


# ======================================================================
# Menu interaktif
# ======================================================================
def backup_menu() -> None:
    console.print(Panel(
        Theme.header(" BACKUP & RESTORE"),
        title="BACKUP", style="bold yellow",
    ))

    bm = BackupManager()
    drive_ready = Path("/content/drive/MyDrive").exists()
    status = "[green]CONNECTED" if drive_ready else "[yellow]NOT MOUNTED"
    console.print(f"  [cyan]Google Drive:[/cyan] {status}")

    bm.show_sizes()

    console.print()
    console.print(Theme.info("Pilihan:"))
    console.print("  [green]1. Backup ESSENTIAL (models + data + code) — direkomendasikan[/green]")
    console.print("  [green]2. Backup MODELS saja (paling kecil)[/green]")
    console.print("  [green]3. Backup DATA saja (dataset + KB + quizzes)[/green]")
    console.print("  [green]4. Backup FULL project[/green]")
    console.print("  [green]5. Restore dari backup[/green]")
    console.print("  [green]6. Import dataset dari Drive[/green]")
    console.print("  [green]7. List backups[/green]")
    console.print("  [green]8. Hapus backup lama[/green]")
    console.print("  [green]9. Kembali[/green]")

    choice = Prompt.ask(
        "[yellow]Pilih",
        choices=['1', '2', '3', '4', '5', '6', '7', '8', '9'],
    )

    try:
        if choice == '1':
            bm.export(mode="essential")
        elif choice == '2':
            bm.export(mode="models_only")
        elif choice == '3':
            bm.export(mode="data_only")
        elif choice == '4':
            console.print(Theme.warning("⚠ Full backup bisa sangat besar (GB+)"))
            if Confirm.ask("[yellow]Lanjutkan?", default=False):
                bm.export(mode="full")
        elif choice == '5':
            bm.restore()
        elif choice == '6':
            bm.import_dataset_from_drive()
        elif choice == '7':
            bm.list_backups()
        elif choice == '8':
            keep_str = Prompt.ask("[cyan]Simpan berapa backup terbaru?", default="3")
            try:
                keep = max(1, int(keep_str))
            except ValueError:
                keep = 3
            bm.clean_old_backups(keep=keep)
        elif choice == '9':
            return
    except Exception as e:
        error_logger.error(f"Backup menu error: {e}")
        console.print(Theme.error(f"Error: {e}"))

    input("\n[yellow]Press Enter...[/yellow]")