#!/usr/bin/env python3
"""
FULL AUTO BATCH ASSET GENERATION
Generates 126 assets with dual-style, background removal, and organization
"""

import json
import time
import os
from pathlib import Path
from PIL import Image
import google.generativeai as genai

# Configure Gemini
genai.configure(api_key=os.environ.get("GEMINI_API_KEY"))

def create_preview(image_path: Path, size=256):
    """Create preview version of image"""
    try:
        img = Image.open(image_path)
        preview = img.resize((size, size), Image.Resampling.LANCZOS)
        preview_path = image_path.parent / f"{image_path.stem}_preview_{size}x{size}.png"
        preview.save(preview_path, 'PNG', optimize=True)
        return preview_path
    except Exception as e:
        print(f"  ⚠️ Preview creation failed: {e}")
        return None

def generate_and_save(asset_name, prompt, style, target_dir, log_file):
    """Generate single image and save to proper location"""
    
    start_time = time.time()
    
    try:
        # Generate image filename
        filename = f"{asset_name}_{style}.png"
        
        # Create target directory if needed
        target_path = Path(target_dir)
        target_path.mkdir(parents=True, exist_ok=True)
        
        # Generate image
        print(f"  🎨 Generating: {filename}")
        log_file.write(f"{time.strftime('%H:%M:%S')} - Generating {filename}\n")
        log_file.flush()
        
        model = genai.GenerativeModel('gemini-2.0-flash-exp')
        response = model.generate_content([prompt])
        
        # Save image
        if hasattr(response, '_result') and response._result.candidates:
            image_data = response._result.candidates[0].content.parts[0].inline_data.data
            
            output_path = target_path / filename
            with open(output_path, 'wb') as f:
                f.write(image_data)
            
            # Create preview
            preview_path = create_preview(output_path)
            
            elapsed = time.time() - start_time
            print(f"  ✅ Saved: {filename} ({elapsed:.1f}s)")
            log_file.write(f"{time.strftime('%H:%M:%S')} - SUCCESS {filename} ({elapsed:.1f}s)\n")
            log_file.flush()
            
            return {
                'success': True,
                'file': str(output_path),
                'preview': str(preview_path) if preview_path else None,
                'time': elapsed
            }
        else:
            print(f"  ❌ Generation failed: No image data")
            log_file.write(f"{time.strftime('%H:%M:%S')} - FAILED {filename} - No image data\n")
            log_file.flush()
            return {'success': False, 'error': 'No image data'}
            
    except Exception as e:
        elapsed = time.time() - start_time
        print(f"  ❌ Error: {e}")
        log_file.write(f"{time.strftime('%H:%M:%S')} - ERROR {filename} - {e}\n")
        log_file.flush()
        return {'success': False, 'error': str(e), 'time': elapsed}

def run_batch_generation(manifest_path="BATCH_GENERATION_MANIFEST.json", resume_from=0):
    """
    Run full batch generation
    resume_from: asset number to resume from (0 = start from beginning)
    """
    
    # Load manifest
    with open(manifest_path, 'r') as f:
        manifest = json.load(f)
    
    total_assets = manifest['total_assets']
    
    print("=" * 70)
    print("🚀 FULL AUTO BATCH ASSET GENERATION")
    print("=" * 70)
    print(f"\n📊 Total assets: {total_assets}")
    print(f"📦 Batches: {len(manifest['batches'])}")
    print(f"⏱️  Estimated time: {total_assets * 15 / 60:.0f} minutes")
    print(f"🔄 Resume from: Asset #{resume_from}")
    print("\n" + "=" * 70)
    
    # Create logs directory
    log_dir = Path("logs")
    log_dir.mkdir(exist_ok=True)
    
    # Open log file
    log_filename = log_dir / f"batch_generation_{time.strftime('%Y%m%d_%H%M%S')}.log"
    log_file = open(log_filename, 'w')
    
    log_file.write(f"BATCH GENERATION LOG - {time.strftime('%Y-%m-%d %H:%M:%S')}\n")
    log_file.write(f"Total assets: {total_assets}\n")
    log_file.write(f"Resume from: {resume_from}\n")
    log_file.write("=" * 70 + "\n\n")
    log_file.flush()
    
    # Statistics
    stats = {
        'total': 0,
        'success': 0,
        'failed': 0,
        'skipped': 0,
        'total_time': 0
    }
    
    asset_counter = 0
    
    try:
        # Process each batch
        for batch in manifest['batches']:
            print(f"\n{'='*70}")
            print(f"📦 BATCH: {batch['name']} ({batch['category']})")
            print(f"{'='*70}")
            
            log_file.write(f"\n{'='*70}\n")
            log_file.write(f"BATCH: {batch['name']}\n")
            log_file.write(f"{'='*70}\n\n")
            log_file.flush()
            
            # Process each asset in batch
            for asset in batch['assets']:
                # Generate styleA
                asset_counter += 1
                
                if asset_counter <= resume_from:
                    stats['skipped'] += 1
                    print(f"\n⏭️  [{asset_counter}/{total_assets}] Skipping: {asset['name']}_styleA")
                    continue
                
                stats['total'] += 1
                print(f"\n📸 [{asset_counter}/{total_assets}] {asset['name']}_styleA")
                
                result = generate_and_save(
                    asset['name'],
                    asset['styleA_prompt'],
                    'styleA',
                    asset['target_dir'],
                    log_file
                )
                
                if result['success']:
                    stats['success'] += 1
                    stats['total_time'] += result.get('time', 0)
                else:
                    stats['failed'] += 1
                
                # Small delay to avoid rate limits
                time.sleep(2)
                
                # Generate styleB
                asset_counter += 1
                
                if asset_counter <= resume_from:
                    stats['skipped'] += 1
                    print(f"\n⏭️  [{asset_counter}/{total_assets}] Skipping: {asset['name']}_styleB")
                    continue
                
                stats['total'] += 1
                print(f"\n📸 [{asset_counter}/{total_assets}] {asset['name']}_styleB")
                
                result = generate_and_save(
                    asset['name'],
                    asset['styleB_prompt'],
                    'styleB',
                    asset['target_dir'],
                    log_file
                )
                
                if result['success']:
                    stats['success'] += 1
                    stats['total_time'] += result.get('time', 0)
                else:
                    stats['failed'] += 1
                
                # Progress update
                progress = (asset_counter / total_assets) * 100
                avg_time = stats['total_time'] / stats['success'] if stats['success'] > 0 else 15
                remaining = (total_assets - asset_counter) * avg_time / 60
                
                print(f"\n📊 Progress: {progress:.1f}% | Success: {stats['success']} | Failed: {stats['failed']} | ETA: {remaining:.0f} min")
                
                # Small delay
                time.sleep(2)
        
        # Final summary
        print("\n" + "=" * 70)
        print("✅ BATCH GENERATION COMPLETE!")
        print("=" * 70)
        print(f"\n📊 FINAL STATISTICS:")
        print(f"  Total processed: {stats['total']}")
        print(f"  ✅ Success: {stats['success']}")
        print(f"  ❌ Failed: {stats['failed']}")
        print(f"  ⏭️  Skipped: {stats['skipped']}")
        print(f"  ⏱️  Total time: {stats['total_time'] / 60:.1f} minutes")
        print(f"  📈 Success rate: {stats['success']/stats['total']*100:.1f}%")
        print(f"\n📝 Log file: {log_filename}")
        
        log_file.write(f"\n{'='*70}\n")
        log_file.write(f"GENERATION COMPLETE\n")
        log_file.write(f"{'='*70}\n")
        log_file.write(f"Total processed: {stats['total']}\n")
        log_file.write(f"Success: {stats['success']}\n")
        log_file.write(f"Failed: {stats['failed']}\n")
        log_file.write(f"Success rate: {stats['success']/stats['total']*100:.1f}%\n")
        
    except KeyboardInterrupt:
        print(f"\n\n⚠️  INTERRUPTED at asset #{asset_counter}")
        print(f"To resume: python3 scripts/batch_generation_runner.py --resume {asset_counter}")
        
        log_file.write(f"\n\nINTERRUPTED at asset #{asset_counter}\n")
        
    except Exception as e:
        print(f"\n\n❌ CRITICAL ERROR: {e}")
        log_file.write(f"\n\nCRITICAL ERROR: {e}\n")
    
    finally:
        log_file.close()

def main():
    import argparse
    
    parser = argparse.ArgumentParser(description='Batch asset generation')
    parser.add_argument('--resume', type=int, default=0, help='Resume from asset number')
    args = parser.parse_args()
    
    run_batch_generation(resume_from=args.resume)

if __name__ == "__main__":
    main()