"""Utilities for managing scratch space used by frame sampling."""

import logging
import shutil
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)


class ScratchManager:
    """Create and clean per-video frame scratch directories."""

    def __init__(self, base_path: str = "/scratch", video_id: Optional[str] = None, auto_cleanup: bool = True):
        self.base_path = Path(base_path)
        self.video_id = video_id or "unknown"
        self.auto_cleanup = auto_cleanup
        self.frame_dir = self.base_path / self.video_id / "frames"

    def ensure_frame_dir(self) -> Path:
        """Create the per-video scratch directory if it does not exist."""
        self.frame_dir.mkdir(parents=True, exist_ok=True)
        return self.frame_dir

    def cleanup(self) -> None:
        """Remove the frame scratch directory only."""
        if self.auto_cleanup and self.frame_dir.exists():
            shutil.rmtree(self.frame_dir)
            logger.info("Cleaned scratch frames for %s", self.video_id)

    def cleanup_all(self) -> None:
        """Remove the entire per-video scratch directory (frames + crops)."""
        video_root = self.base_path / self.video_id
        if self.auto_cleanup and video_root.exists():
            shutil.rmtree(video_root)
            logger.info("Cleaned all scratch data for %s", self.video_id)

    def usage_bytes(self) -> int:
        """Return the total size of the scratch directory in bytes."""
        if not self.frame_dir.exists():
            return 0
        return sum(path.stat().st_size for path in self.frame_dir.rglob("*") if path.is_file())