cs249r_book/binder/cli/commands/clean.py

"""
Clean command implementation for MLSysBook CLI.

Handles cleaning build artifacts and restoring configurations.
"""

import shutil
from pathlib import Path
from rich.console import Console
from rich.markup import escape as _rich_escape
from rich.panel import Panel
from rich.table import Table

from cli.core.artifacts import clean_build_artifacts
from cli.core.workspace import WORKSPACE_ROOT, remove_stale_runs

console = Console()

FORMATS = ("html", "pdf", "epub")


def _count_files(directory: Path) -> int:
    """Return the number of files below *directory*."""
    return sum(1 for path in directory.rglob("*") if path.is_file())


class CleanCommand:
    """Handles cleaning operations for the MLSysBook."""

    def __init__(self, config_manager, chapter_discovery):
        """Initialize clean command.

        Args:
            config_manager: ConfigManager instance
            chapter_discovery: ChapterDiscovery instance
        """
        self.config_manager = config_manager
        self.chapter_discovery = chapter_discovery

    def clean_all(self) -> bool:
        """Clean all build artifacts and restore configs.

        Returns:
            True if cleaning succeeded, False otherwise
        """
        console.print("[bold blue]🧹 MLSysBook Cleanup[/bold blue]")
        console.print("[dim]Cleaning build artifacts and restoring configurations...[/dim]\n")

        success = True

        # Clean build directories
        success &= self._clean_build_directories()

        # Clean temporary files
        success &= self._clean_temp_files()

        # Restore configurations
        success &= self._restore_configs()

        # Show current status
        self._show_cleanup_status()

        if success:
            console.print("\n[green]✅ Cleanup completed successfully![/green]")
        else:
            console.print("\n[yellow]⚠️ Cleanup completed with some issues[/yellow]")

        return success

    def _clean_build_directories(self) -> bool:
        """Clean build output directories, stopping at the first failure."""
        console.print("[blue]📁 Cleaning build directories...[/blue]")
        for format_type in FORMATS:
            if not self._remove_output_dir(format_type, indent="  "):
                return False
        return True

    def _remove_output_dir(self, format_type: str, indent: str = "") -> bool:
        """Delete the active output directory for *format_type* and report what was removed."""
        output_dir = self.config_manager.get_output_dir(format_type)
        label = format_type.upper()
        if not output_dir.exists():
            console.print(f"{indent}📁 {label} build directory not found (already clean): {output_dir}")
            return True
        file_count = _count_files(output_dir)
        try:
            shutil.rmtree(output_dir)
        except OSError as e:
            console.print(f"{indent}❌ Could not remove {label} build directory {output_dir}: {e}")
            return False
        console.print(f"{indent}✅ Cleaned {label} build ({file_count} files) from {output_dir}")
        return True

    def _clean_temp_files(self) -> bool:
        """Clean temporary files and caches."""
        console.print("[blue]🗑️ Cleaning temporary files...[/blue]")

        temp_patterns = [
            "**/.quarto",
            "**/.*_cache",
            "**/*_files",
            "**/.DS_Store",
            "**/Thumbs.db",
            "**/*.tmp",
            "**/*.log"
        ]

        cleaned_count = 0

        for pattern in temp_patterns:
            try:
                for temp_file in self.config_manager.book_dir.glob(pattern):
                    if temp_file.exists():
                        if temp_file.is_file():
                            temp_file.unlink()
                            cleaned_count += 1
                        elif temp_file.is_dir():
                            shutil.rmtree(temp_file)
                            cleaned_count += 1
            except OSError as e:
                console.print(f"  ⚠️ Could not remove some files matching {pattern} under "
                              f"{self.config_manager.book_dir}: {e}")

        console.print(f"  ✅ Cleaned {cleaned_count} temporary files/directories")
        return True

    def _restore_configs(self) -> bool:
        """Restore configuration files to clean state."""
        console.print("[blue]⚙️ Restoring configurations...[/blue]")

        try:
            # Remove the generated _quarto.yml; leave a hand-written one alone
            source = self.config_manager.active_config_source()
            if source:
                self.config_manager.active_config.unlink()
                console.print(f"  ✅ Removed generated _quarto.yml (copy of {source})")
            elif self.config_manager.active_config.exists():
                console.print("  📄 Active config was not generated by binder (not removing)")
            else:
                console.print("  📁 No active config found")

            return True

        except OSError as e:
            console.print(f"  ❌ Could not remove {self.config_manager.active_config}: {e}")
            return False

    def _show_cleanup_status(self) -> None:
        """Show current status after cleanup."""
        console.print("\n[blue]📊 Post-cleanup status:[/blue]")

        # Show the active config
        self.config_manager.show_active_config()

        # Show build directory status
        for format_type in FORMATS:
            output_dir = self.config_manager.get_output_dir(format_type)
            if output_dir.exists():
                console.print(f"  📁 {format_type.upper()}: {_count_files(output_dir)} files remaining in {output_dir}")
            else:
                console.print(f"  📁 {format_type.upper()}: Clean (no build directory)")

    def clean_format(self, format_type: str) -> bool:
        """Clean artifacts for a specific format.

        Args:
            format_type: Format to clean ('html', 'pdf', 'epub')

        Returns:
            True if cleaning succeeded, False otherwise
        """
        console.print(f"[blue]🧹 Cleaning {format_type.upper()} artifacts...[/blue]")
        return self._remove_output_dir(format_type)

    def clean_artifacts(self) -> bool:
        """Clean build artifacts using Binder-native cleanup primitives.

        This is the same cleanup that runs as a pre-commit hook.
        """
        clean_build_artifacts(self.config_manager.book_dir, console=console)
        return True

    def clean_workspaces(self) -> bool:
        """Remove build worktrees left by parallel or debug runs that were killed.

        Only runs whose binder process has exited are touched, so a build that
        is still running keeps its worktrees. Runs kept with
        ``--keep-workspaces`` are removed too.

        Returns:
            True (leftovers that cannot be removed are reported, not fatal).
        """
        repo_root = self.config_manager.book_dir.parent
        removed = remove_stale_runs(repo_root, include_kept=True)
        if removed:
            console.print(f"[green]✅ Removed {len(removed)} leftover build run(s) from {_rich_escape(str(WORKSPACE_ROOT))}[/green]")
            for run_dir in removed:
                console.print(f"[dim]  {_rich_escape(run_dir.name)}[/dim]")
        else:
            console.print("[green]✅ No leftover build worktrees[/green]")
        return True

    def print_help(self) -> None:
        """Print the `binder clean` command reference."""
        table = Table(show_header=True, header_style="bold cyan", box=None)
        table.add_column("Target", style="cyan", width=14)
        table.add_column("Description", style="white", width=58)
        table.add_row("(none)", "Clean HTML/PDF/EPUB build directories, temp files, and the generated _quarto.yml")
        table.add_row("html", "Clean the active HTML output directory")
        table.add_row("pdf", "Clean the active PDF output directory")
        table.add_row("epub", "Clean the active EPUB output directory")
        table.add_row("artifacts", "Clean generated build artifacts using the same native primitive used by pre-commit")
        table.add_row("workspaces", "Remove build worktrees left by finished --parallel or debug runs (killed or --keep-workspaces)")
        console.print(Panel(table, title=_rich_escape("binder clean [target]"), border_style="cyan"))
        console.print("[dim]Examples:[/dim]")
        console.print("  [cyan]./binder/binder clean[/cyan]            [dim]# full local cleanup[/dim]")
        console.print("  [cyan]./binder/binder clean html[/cyan]       [dim]# clean HTML output only[/dim]")
        console.print("  [cyan]./binder/binder clean artifacts[/cyan]  [dim]# generated-artifact cleanup[/dim]")
        console.print("  [cyan]./binder/binder clean workspaces[/cyan] [dim]# leftover build worktrees[/dim]")
        console.print()