#!/usr/bin/env python3
"""
Complete asset organization for ALL asset directories:
1. Remove sample/test/placeholder files
2. Add 256x256 preview versions
3. Organize into subfolders with proper naming

Works on ALL directories in assets/images/
"""

import os
import shutil
from pathlib import Path
from PIL import Image
from collections import defaultdict

# Files to delete (junk/samples/tests)
DELETE_PATTERNS = [
    '*sample*', '*test*', '*placeholder*', '*_old*', 
    '*temp*', '*backup*', '*copy*', '*duplicate*'
]

def should_delete(filename: str) -> bool:
    """Check if file should be deleted"""
    lower = filename.lower()
    patterns = ['sample', 'test', 'placeholder', '_old', 'temp', 'backup', 'copy', 'duplicate']
    return any(p in lower for p in patterns)

def is_valid_image(filepath: Path) -> bool:
    """Check if image is valid (not corrupted, not too small)"""
    try:
        img = Image.open(filepath)
        width, height = img.size
        # Too small = probably junk
        if width < 16 or height < 16:
            return False
        # Check if image can be loaded
        img.verify()
        return True
    except:
        return False

def clean_junk_files(directory: Path, dry_run: bool = False):
    """Remove sample/test/junk files"""
    
    deleted = []
    
    for png_file in directory.rglob('*.png'):
        if should_delete(png_file.name):
            deleted.append(png_file.name)
            if not dry_run:
                png_file.unlink()
                print(f"  🗑️  Deleted junk: {png_file.name}")
        elif not is_valid_image(png_file):
            deleted.append(png_file.name)
            if not dry_run:
                png_file.unlink()
                print(f"  🗑️  Deleted invalid: {png_file.name}")
    
    return deleted

def add_preview_version(asset_path: Path, dry_run: bool = False):
    """Add 256x256 preview if missing"""
    
    # Skip if already a preview/sprite
    if 'preview' in asset_path.stem or 'sprite' in asset_path.stem:
        return False
    
    # Skip tiny files
    if any(size in asset_path.stem for size in ['32x32', '16x16', '24x24', '32x64']):
        return False
    
    try:
        # Create preview filename
        preview_filename = f"{asset_path.stem}_preview_256x256.png"
        preview_path = asset_path.parent / preview_filename
        
        # Skip if exists
        if preview_path.exists():
            return False
        
        # Load and resize
        img = Image.open(asset_path)
        preview = img.resize((256, 256), Image.Resampling.LANCZOS)
        
        if not dry_run:
            preview.save(preview_path, 'PNG', optimize=True)
            print(f"    ✅ Created preview: {preview_filename}")
        
        return True
        
    except Exception as e:
        print(f"    ❌ Error creating preview for {asset_path.name}: {e}")
        return False

def get_base_asset_name(filename: str) -> str:
    """Extract base asset name"""
    stem = filename.replace('.png', '')
    
    # Remove style suffix
    if '_styleA' in stem:
        stem = stem.split('_styleA')[0]
    elif '_styleB' in stem:
        stem = stem.split('_styleB')[0]
    
    # Remove size suffix
    for size in ['_1024x1024', '_256x256', '_32x32', '_16x16', '_32x64', '_24x24']:
        stem = stem.replace(size, '')
    
    # Remove type suffix
    for suffix in ['_preview', '_sprite']:
        stem = stem.replace(suffix, '')
    
    return stem

def organize_directory(directory: Path, dry_run: bool = False):
    """
    Organize single directory:
    1. Clean junk files
    2. Add preview versions
    3. Organize into subfolders
    """
    
    if not directory.exists() or not directory.is_dir():
        return 0
    
    print(f"\n{'='*70}")
    print(f"📂 ORGANIZING: {directory.relative_to('assets/images')}")
    print(f"{'='*70}")
    
    # Step 1: Clean junk
    deleted = clean_junk_files(directory, dry_run)
    if deleted:
        print(f"  🗑️  Cleaned {len(deleted)} junk files")
    
    # Step 2: Add preview versions
    preview_count = 0
    for png_file in directory.glob('*.png'):
        if png_file.is_file():
            if add_preview_version(png_file, dry_run):
                preview_count += 1
    
    if preview_count > 0:
        print(f"  🖼️  Created {preview_count} preview versions")
    
    # Step 3: Delete 'originals/' folder if exists
    originals_path = directory / 'originals'
    if originals_path.exists():
        if not dry_run:
            shutil.rmtree(originals_path)
        print(f"  🗑️  Deleted duplicate folder: originals/")
    
    # Step 4: Group files by base asset name
    asset_groups = defaultdict(list)
    
    for png_file in directory.glob('*.png'):
        if png_file.is_file():
            base_name = get_base_asset_name(png_file.name)
            asset_groups[base_name].append(png_file)
    
    if not asset_groups:
        print("  ✓ No assets to organize (already done or empty)")
        return 0
    
    # Step 5: Organize into subfolders
    organized = 0
    for asset_name, files in sorted(asset_groups.items()):
        # Organize if there's original + preview, or multiple related files
        has_preview = any('preview' in f.stem for f in files)
        has_sprite = any(any(s in f.stem for s in ['32x32', '16x16']) for f in files)
        
        # Only organize if we have multiple versions (original + preview/sprite)
        if len(files) < 2 and not (has_preview or has_sprite):
            continue
        
        asset_folder = directory / asset_name
        
        if not dry_run:
            asset_folder.mkdir(exist_ok=True)
        
        print(f"\n  📁 {asset_name}/ ({len(files)} files)")
        
        for file_path in files:
            # Determine new filename
            if 'preview' in file_path.stem:
                new_name = f"{asset_name}_preview_256x256.png"
            elif any(s in file_path.stem for s in ['32x32', '16x16', '32x64', '24x24']):
                # Extract size
                for s in ['32x32', '16x16', '32x64', '24x24']:
                    if s in file_path.stem:
                        new_name = f"{asset_name}_sprite_{s}.png"
                        break
            else:
                # Original - assume 1024x1024
                new_name = f"{asset_name}_1024x1024.png"
            
            new_path = asset_folder / new_name
            
            if not dry_run:
                file_path.rename(new_path)
            
            print(f"    → {new_name}")
        
        organized += 1
    
    return organized

def main():
    import argparse
    
    parser = argparse.ArgumentParser(description='Organize ALL assets with cleanup')
    parser.add_argument('--dry-run', action='store_true', help='Show what would be done')
    args = parser.parse_args()
    
    print("=" * 70)
    print("🧹 COMPLETE ASSET ORGANIZATION + CLEANUP")
    print("=" * 70)
    
    if args.dry_run:
        print("\n⚠️  DRY RUN MODE - No changes will be made\n")
    
    print("\nThis will:")
    print("  1. 🗑️  Delete sample/test/placeholder files")
    print("  2. 🖼️  Add 256x256 preview versions")
    print("  3. 📁 Organize into subfolders")
    print("  4. 🏷️  Rename to proper format")
    print("=" * 70)
    
    # Get all directories
    base_dir = Path('assets/images')
    
    # Skip these directories
    skip_dirs = {'demo_originals_with_white_bg'}
    
    total_organized = 0
    
    for subdir in sorted(base_dir.iterdir()):
        if subdir.is_dir() and subdir.name not in skip_dirs:
            count = organize_directory(subdir, dry_run=args.dry_run)
            total_organized += count
    
    print("\n" + "=" * 70)
    if not args.dry_run:
        print(f"✅ COMPLETE! Organized {total_organized} asset groups")
    else:
        print(f"✅ DRY RUN COMPLETE! Would organize {total_organized} asset groups")
    print("=" * 70)

if __name__ == "__main__":
    main()