#!/usr/bin/env python3
"""
Cleanup Script for Character Forge Deployment
==============================================
Licensed under GNU AGPL v3.0

This script removes all generated content, test files, and temporary data
before deploying to HuggingFace or committing to Git.

It will DELETE:
- All generated images
- Test outputs
- Log files
- Cache directories
- Temporary files

SAFE TO RUN: Only removes generated content, never source code.
"""

import os
import shutil
from pathlib import Path

def get_base_dir():
    """Get the base directory of the project."""
    return Path(__file__).parent.absolute()

def remove_directory(path):
    """Safely remove a directory and all its contents."""
    if path.exists() and path.is_dir():
        try:
            shutil.rmtree(path)
            print(f"[DELETED] {path}")
            return True
        except Exception as e:
            print(f"[ERROR] Failed to delete {path}: {e}")
            return False
    return False

def remove_file(path):
    """Safely remove a file."""
    if path.exists() and path.is_file():
        try:
            path.unlink()
            print(f"[DELETED] {path}")
            return True
        except Exception as e:
            print(f"[ERROR] Failed to delete {path}: {e}")
            return False
    return False

def find_and_remove_pattern(base_dir, pattern, file_type="file"):
    """Find and remove files or directories matching a pattern."""
    count = 0
    if file_type == "file":
        for path in base_dir.rglob(pattern):
            if path.is_file():
                if remove_file(path):
                    count += 1
    else:  # directory
        for path in base_dir.rglob(pattern):
            if path.is_dir():
                if remove_directory(path):
                    count += 1
    return count

def cleanup_outputs(base_dir):
    """Remove all output directories."""
    print("\n" + "="*70)
    print("CLEANING OUTPUT DIRECTORIES")
    print("="*70)

    output_dirs = [
        base_dir / "outputs",
        base_dir / "output",
        base_dir / "character_forge_image" / "outputs",
    ]

    count = 0
    for output_dir in output_dirs:
        if remove_directory(output_dir):
            count += 1

    print(f"\n[OK] Removed {count} output directories")
    return count

def cleanup_images(base_dir):
    """Remove all generated images."""
    print("\n" + "="*70)
    print("CLEANING GENERATED IMAGES")
    print("="*70)

    image_extensions = ["*.png", "*.jpg", "*.jpeg", "*.webp", "*.gif"]

    # Directories to preserve (docs, assets, etc.)
    preserve_dirs = ["docs", "assets", ".git"]

    total_count = 0
    for ext in image_extensions:
        for img_path in base_dir.rglob(ext):
            # Skip if in preserved directories
            if any(preserve in str(img_path) for preserve in preserve_dirs):
                print(f"[SKIP] Preserving {img_path}")
                continue

            if remove_file(img_path):
                total_count += 1

    print(f"\n[OK] Removed {total_count} image files")
    return total_count

def cleanup_logs(base_dir):
    """Remove all log files."""
    print("\n" + "="*70)
    print("CLEANING LOG FILES")
    print("="*70)

    count = find_and_remove_pattern(base_dir, "*.log", "file")
    print(f"\n[OK] Removed {count} log files")
    return count

def cleanup_cache(base_dir):
    """Remove cache directories."""
    print("\n" + "="*70)
    print("CLEANING CACHE DIRECTORIES")
    print("="*70)

    cache_patterns = ["__pycache__", ".library", ".cache"]

    total_count = 0
    for pattern in cache_patterns:
        count = find_and_remove_pattern(base_dir, pattern, "directory")
        total_count += count

    print(f"\n[OK] Removed {total_count} cache directories")
    return total_count

def cleanup_temp(base_dir):
    """Remove temporary files and directories."""
    print("\n" + "="*70)
    print("CLEANING TEMPORARY FILES")
    print("="*70)

    temp_patterns = ["*.tmp", "*.temp", "tmp", "temp"]

    total_count = 0
    for pattern in temp_patterns:
        if pattern.startswith("*"):
            count = find_and_remove_pattern(base_dir, pattern, "file")
        else:
            count = find_and_remove_pattern(base_dir, pattern, "directory")
        total_count += count

    print(f"\n[OK] Removed {total_count} temporary items")
    return total_count

def cleanup_test_files(base_dir):
    """Remove test output files and directories."""
    print("\n" + "="*70)
    print("CLEANING TEST FILES")
    print("="*70)

    # Remove test output directories
    test_dirs = [
        base_dir / "character_forge_image" / "outputs" / "test_female_tattoos",
        base_dir / "character_forge_image" / "outputs" / "test_flux_pipeline",
    ]

    count = 0
    for test_dir in test_dirs:
        if remove_directory(test_dir):
            count += 1

    print(f"\n[OK] Removed {count} test directories")
    return count

def get_directory_size(path):
    """Calculate total size of a directory in MB."""
    total_size = 0
    try:
        for dirpath, dirnames, filenames in os.walk(path):
            for filename in filenames:
                filepath = os.path.join(dirpath, filename)
                if os.path.exists(filepath):
                    total_size += os.path.getsize(filepath)
    except Exception as e:
        print(f"[ERROR] Could not calculate size: {e}")
        return 0

    return total_size / (1024 * 1024)  # Convert to MB

def main():
    """Main cleanup function."""
    print("="*70)
    print("CHARACTER FORGE - DEPLOYMENT CLEANUP")
    print("="*70)
    print("\nThis will remove all generated content, test files, and logs.")
    print("Source code will NOT be touched.")

    base_dir = get_base_dir()
    print(f"\nBase directory: {base_dir}")

    # Calculate initial size
    initial_size = get_directory_size(base_dir)
    print(f"Initial size: {initial_size:.2f} MB")

    # Perform cleanup
    stats = {
        "outputs": cleanup_outputs(base_dir),
        "images": cleanup_images(base_dir),
        "logs": cleanup_logs(base_dir),
        "cache": cleanup_cache(base_dir),
        "temp": cleanup_temp(base_dir),
        "tests": cleanup_test_files(base_dir),
    }

    # Calculate final size
    final_size = get_directory_size(base_dir)
    saved_size = initial_size - final_size

    # Summary
    print("\n" + "="*70)
    print("CLEANUP SUMMARY")
    print("="*70)
    print(f"Output directories removed: {stats['outputs']}")
    print(f"Image files removed: {stats['images']}")
    print(f"Log files removed: {stats['logs']}")
    print(f"Cache directories removed: {stats['cache']}")
    print(f"Temporary items removed: {stats['temp']}")
    print(f"Test directories removed: {stats['tests']}")
    print(f"\nInitial size: {initial_size:.2f} MB")
    print(f"Final size: {final_size:.2f} MB")
    print(f"Space saved: {saved_size:.2f} MB")

    print("\n" + "="*70)
    print("[SUCCESS] CLEANUP COMPLETE!")
    print("="*70)
    print("\nYour project is now clean and ready for:")
    print("  - Git commit")
    print("  - HuggingFace deployment")
    print("  - GitHub upload")
    print("\nThe .gitignore file will prevent these files from being")
    print("added again in the future.")

if __name__ == "__main__":
    main()