#!/usr/bin/env python3
"""
Fix User Structure Script
This script reorganizes existing flat image structures into the proper angle-based subdirectories
that the training function expects.
"""

import os
import shutil
import json

def fix_user_structure(user_id):
    """Fix the image structure for a specific user"""
    user_folder = os.path.join("registered_faces", user_id)
    
    if not os.path.exists(user_folder):
        print(f"❌ User folder {user_folder} not found!")
        return False
    
    # Check if already has proper structure
    angle_dirs = ['front', 'left', 'right', 'up', 'down']
    has_proper_structure = all(os.path.isdir(os.path.join(user_folder, angle)) for angle in angle_dirs)
    
    if has_proper_structure:
        print(f"✅ User {user_id} already has proper structure")
        return True
    
    print(f"🔧 Fixing structure for User {user_id}...")
    
    # Get all jpg files in the user folder
    image_files = [f for f in os.listdir(user_folder) if f.lower().endswith('.jpg')]
    
    if not image_files:
        print(f"❌ No images found for User {user_id}")
        return False
    
    print(f"📸 Found {len(image_files)} images to reorganize")
    
    # Create angle subdirectories
    for angle in angle_dirs:
        angle_path = os.path.join(user_folder, angle)
        os.makedirs(angle_path, exist_ok=True)
    
    # Calculate images per angle
    images_per_angle = len(image_files) // len(angle_dirs)
    print(f"📊 Will organize {images_per_angle} images per angle")
    
    # Move images to appropriate subdirectories
    for i, filename in enumerate(sorted(image_files)):
        # Determine which angle directory this image belongs to
        angle_index = min(i // images_per_angle, len(angle_dirs) - 1)
        angle_dir = angle_dirs[angle_index]
        
        # Create new filename
        new_filename = f'{i % images_per_angle}.jpg'
        
        # Source and destination paths
        src_path = os.path.join(user_folder, filename)
        dst_path = os.path.join(user_folder, angle_dir, new_filename)
        
        # Move the file
        shutil.move(src_path, dst_path)
        print(f"  📁 Moved {filename} → {angle_dir}/{new_filename}")
    
    print(f"✅ User {user_id} structure fixed successfully!")
    return True

def fix_all_users():
    """Fix the structure for all users"""
    print("🔧 Fixing image structure for all users...")
    
    if not os.path.exists("registered_faces"):
        print("❌ registered_faces directory not found!")
        return
    
    success_count = 0
    total_users = 0
    
    for user_id in os.listdir("registered_faces"):
        user_folder = os.path.join("registered_faces", user_id)
        if os.path.isdir(user_folder):
            total_users += 1
            if fix_user_structure(user_id):
                success_count += 1
    
    print(f"\n📊 Summary:")
    print(f"  Total users processed: {total_users}")
    print(f"  Successfully fixed: {success_count}")
    print(f"  Failed: {total_users - success_count}")
    
    if success_count > 0:
        print(f"\n🔄 Now retraining the model with fixed structure...")
        try:
            from livestream import train_face_recognizer
            train_face_recognizer()
            print("✅ Model retrained successfully!")
        except Exception as e:
            print(f"❌ Error retraining model: {e}")

def show_user_status():
    """Show current status of all users"""
    print("📊 Current User Status:")
    print("=" * 60)
    
    if not os.path.exists("registered_faces"):
        print("❌ registered_faces directory not found!")
        return
    
    for user_id in os.listdir("registered_faces"):
        user_folder = os.path.join("registered_faces", user_id)
        if os.path.isdir(user_folder):
            # Check structure
            angle_dirs = ['front', 'left', 'right', 'up', 'down']
            has_proper_structure = all(os.path.isdir(os.path.join(user_folder, angle)) for angle in angle_dirs)
            
            # Count images
            total_images = 0
            if has_proper_structure:
                for angle in angle_dirs:
                    angle_path = os.path.join(user_folder, angle)
                    if os.path.exists(angle_path):
                        total_images += len([f for f in os.listdir(angle_path) if f.lower().endswith('.jpg')])
            else:
                total_images = len([f for f in os.listdir(user_folder) if f.lower().endswith('.jpg')])
            
            # Get user name
            user_name = "Unknown"
            if os.path.exists("user_data.json"):
                try:
                    with open("user_data.json", "r") as f:
                        user_data = json.load(f)
                        if str(user_id) in user_data:
                            user_name = user_data[str(user_id)]["name"]
                except:
                    pass
            
            status = "✅ Proper" if has_proper_structure else "❌ Flat"
            print(f"User {user_id:>6} | {user_name:<25} | Images: {total_images:>3} | Structure: {status}")

def main():
    """Main menu"""
    print("Fix User Structure Tool")
    print("=" * 30)
    
    while True:
        print("\nOptions:")
        print("1. Show current user status")
        print("2. Fix specific user structure")
        print("3. Fix all users' structure")
        print("4. Exit")
        
        choice = input("\nSelect an option (1-4): ").strip()
        
        if choice == "1":
            show_user_status()
        
        elif choice == "2":
            user_id = input("Enter user ID to fix: ").strip()
            if user_id:
                fix_user_structure(user_id)
        
        elif choice == "3":
            fix_all_users()
        
        elif choice == "4":
            print("Exiting...")
            break
        
        else:
            print("Invalid choice. Please try again.")

if __name__ == "__main__":
    main()