#!/usr/bin/env python3
"""
Clean up stale HuggingFace cache artifacts from previous failed downloads.

This script removes incomplete marker files and temporary files from the
HuggingFace cache directory. These artifacts can accumulate from interrupted
or failed downloads and may interfere with future downloads.
"""

import os
import sys
import time
from pathlib import Path
from typing import List

try:
    from huggingface_hub import constants

    HF_HUB_AVAILABLE = True
except ImportError:
    print("Warning: huggingface_hub not available")
    HF_HUB_AVAILABLE = False


def get_hf_cache_dir() -> str:
    """Get the HuggingFace cache directory."""
    if HF_HUB_AVAILABLE:
        return constants.HF_HUB_CACHE

    # Fallback to environment variable or default
    hf_home = os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface"))
    return os.path.join(hf_home, "hub")


def find_stale_artifacts(cache_dir: str) -> List[Path]:
    """
    Find stale artifact files in the HuggingFace cache.

    Args:
        cache_dir: HuggingFace cache directory

    Returns:
        List of paths to stale artifact files
    """
    cache_path = Path(cache_dir)

    if not cache_path.exists():
        return []

    # Not "**/*.lock": another container may hold a download lock for 30+ min;
    # unlinking it gives the next acquirer a fresh inode and both proceed.
    patterns = [
        "**/*.incomplete",  # Incomplete download markers
        "**/*.tmp",  # Temporary files
    ]

    # Another container may still be appending to its blobs/*.incomplete.
    min_stale_age_seconds = 2 * 60 * 60

    stale_files = []
    now = time.time()
    for pattern in patterns:
        for path in cache_path.glob(pattern):
            try:
                if now - path.stat().st_mtime < min_stale_age_seconds:
                    continue
            except OSError:
                continue  # vanished mid-scan
            stale_files.append(path)

    return stale_files


def cleanup_artifacts(artifacts: List[Path]) -> tuple[int, int]:
    """
    Remove stale artifact files.

    Args:
        artifacts: List of file paths to remove

    Returns:
        Tuple of (successful_removals, failed_removals)
    """
    successful = 0
    failed = 0

    for file_path in artifacts:
        try:
            file_path.unlink()
            print(f"  Removed: {file_path}")
            successful += 1
        except Exception as e:
            print(f"  Warning: Could not remove {file_path}: {e}")
            failed += 1

    return successful, failed


def main() -> int:
    """
    Main cleanup logic.

    Returns:
        Always returns 0 (cleanup is best-effort and should not fail CI)
    """
    print("=" * 70)
    print("HuggingFace Cache Cleanup")
    print("=" * 70)

    # Get cache directory
    cache_dir = get_hf_cache_dir()
    print(f"Cache directory: {cache_dir}")

    if not os.path.exists(cache_dir):
        print("Cache directory does not exist - nothing to clean")
        return 0

    print("-" * 70)

    # Find stale artifacts
    print("Scanning for stale artifacts...")
    stale_artifacts = find_stale_artifacts(cache_dir)

    if not stale_artifacts:
        print("✓ No stale cache artifacts found")
        return 0

    # Clean up artifacts
    print(f"Found {len(stale_artifacts)} stale artifact(s) to remove:")
    successful, failed = cleanup_artifacts(stale_artifacts)

    print("-" * 70)

    # Summary
    if failed > 0:
        print(f"⚠ Cleaned up {successful} file(s), {failed} removal(s) failed")
    else:
        print(f"✓ Successfully cleaned up {successful} stale file(s)")

    # Always return 0 - cleanup failures should not fail CI
    return 0


if __name__ == "__main__":
    try:
        exit_code = main()
        sys.exit(exit_code)
    except KeyboardInterrupt:
        print("\nInterrupted by user")
        sys.exit(0)
    except Exception as e:
        print(f"ERROR: Unexpected error during cleanup: {e}")
        import traceback

        traceback.print_exc()
        # Still return 0 - cleanup failures should not fail CI
        sys.exit(0)
