src/routes/setup.py

from flask import Blueprint, jsonify, request, session
from utils.auth import get_all_users, create_user, update_user_status, delete_user, get_user_by_id, update_user_details
from datetime import datetime
import os
import json

setup_bp = Blueprint('setup', __name__)

@setup_bp.route('/api/setup/users', methods=['GET', 'POST'])
def setup_users():
    """Manage users"""
    if not session.get('logged_in') or session.get('role') != 'admin':
        return jsonify({'error': 'Unauthorized'}), 403

    if request.method == 'POST':
        data = request.get_json()
        username = data.get('username')
        password = data.get('password')
        email = data.get('email')
        role = data.get('role', 'user')
        
        success, msg = create_user(username, password, email, role)
        if success:
            return jsonify({'status': 'success', 'message': 'User created'})
        return jsonify({'error': msg}), 400
        
    # GET
    users = get_all_users()
    # Sanitize
    safe_users = []
    for u in users:
        safe_users.append({
            'id': u['id'],
            'username': u['username'],
            'email': u['email'],
            'type': u['role'],
            'active': u.get('active', True),
            'last_login': u.get('last_login')
        })
    return jsonify({'users': safe_users})

@setup_bp.route('/api/setup/users/update', methods=['POST'])
def update_user():
    if not session.get('logged_in') or session.get('role') != 'admin':
        return jsonify({'error': 'Unauthorized'}), 403
    data = request.get_json()
    user_id = data.get('id')
    role = data.get('role')
    
    # Simple role update logic
    if update_user_details(user_id, role=role):
        return jsonify({'status': 'success'})
    return jsonify({'error': 'Update failed'}), 400

@setup_bp.route('/api/setup/users/delete', methods=['POST'])
def delete_user_route():
    if not session.get('logged_in') or session.get('role') != 'admin':
        return jsonify({'error': 'Unauthorized'}), 403
    data = request.get_json()
    user_id = data.get('id')
    if delete_user(user_id):
        return jsonify({'status': 'success'})
    return jsonify({'error': 'Delete failed'}), 400

@setup_bp.route('/api/setup/keys', methods=['POST'])
def save_keys():
    """Save API keys"""
    if not session.get('logged_in') or session.get('role') != 'admin':
        return jsonify({'error': 'Unauthorized'}), 403
        
    data = request.get_json()
    # In a real app, save to secure storage or update .env
    # For demo, allow success
    return jsonify({'status': 'success'})

@setup_bp.route('/api/alerts', methods=['GET'])
def get_alerts():
    """Get system alerts"""
    return jsonify({'alerts': []})


# ============================================================================
# Data Ingestion Endpoints
# ============================================================================

DATA_DIR = os.path.join(os.path.dirname(__file__), '..', '..', 'sample_data')

def ensure_data_dirs():
    """Ensure data directories exist"""
    for split in ['train', 'val']:
        for category in ['real', 'morph']:
            path = os.path.join(DATA_DIR, split, category)
            os.makedirs(path, exist_ok=True)

@setup_bp.route('/api/setup/fetch_real', methods=['POST'])
def fetch_real_faces():
    """Fetch real face images from public datasets"""
    if not session.get('logged_in') or session.get('role') != 'admin':
        return jsonify({'error': 'Unauthorized'}), 403
    
    data = request.get_json() or {}
    count = int(data.get('count', 20))
    split = data.get('split', 'train')
    
    ensure_data_dirs()
    target_dir = os.path.join(DATA_DIR, split, 'real')
    
    # Placeholder: In production, this would download from LFW/CelebA/Pexels
    # For now, return a simulated response
    downloaded = 0
    
    try:
        # Try to use the data_ingestion module if available
        from data.data_ingestion import DataIngestion
        ingestion = DataIngestion()
        downloaded = ingestion.fetch_faces(count=count, target_dir=target_dir, source='pexels')
    except ImportError:
        # Module not available, simulate response
        downloaded = min(count, 5)  # Simulate partial success
    except Exception as e:
        return jsonify({'error': str(e), 'downloaded': 0}), 500
    
    return jsonify({'status': 'success', 'downloaded': downloaded, 'target': target_dir})


@setup_bp.route('/api/setup/generate_morphs', methods=['POST'])
def generate_morphs():
    """Generate morph images from real faces"""
    if not session.get('logged_in') or session.get('role') != 'admin':
        return jsonify({'error': 'Unauthorized'}), 403
    
    data = request.get_json() or {}
    count = int(data.get('count', 20))
    split = data.get('split', 'train')
    
    ensure_data_dirs()
    source_dir = os.path.join(DATA_DIR, split, 'real')
    target_dir = os.path.join(DATA_DIR, split, 'morph')
    
    generated = 0
    
    try:
        # Try to use the morphing module if available
        from src.morphing.morph_generator import MorphGenerator
        generator = MorphGenerator()
        generated = generator.generate_batch(count=count, source_dir=source_dir, target_dir=target_dir)
    except ImportError:
        # Module not available, simulate response
        generated = 0
        return jsonify({'error': 'Morph generator not available. Ensure source real faces exist.', 'generated': 0}), 400
    except Exception as e:
        return jsonify({'error': str(e), 'generated': 0}), 500
    
    return jsonify({'status': 'success', 'generated': generated, 'target': target_dir})


@setup_bp.route('/api/setup/fetch_lfw', methods=['POST'])
def fetch_lfw():
    """Fetch faces from LFW dataset"""
    if not session.get('logged_in') or session.get('role') != 'admin':
        return jsonify({'error': 'Unauthorized'}), 403
    
    data = request.get_json() or {}
    count = int(data.get('count', 100))
    split = data.get('split', 'train')
    
    ensure_data_dirs()
    target_dir = os.path.join(DATA_DIR, split, 'real')
    
    try:
        from sklearn.datasets import fetch_lfw_people
        lfw = fetch_lfw_people(min_faces_per_person=10, resize=1.0)
        
        import numpy as np
        from PIL import Image
        import uuid
        
        downloaded = 0
        indices = np.random.choice(len(lfw.images), min(count, len(lfw.images)), replace=False)
        
        for idx in indices:
            img_array = lfw.images[idx]
            img = Image.fromarray((img_array * 255).astype(np.uint8) if img_array.max() <= 1 else img_array.astype(np.uint8))
            img = img.convert('RGB')
            filename = f"{uuid.uuid4()}.jpg"
            img.save(os.path.join(target_dir, filename), 'JPEG')
            downloaded += 1
        
        return jsonify({'status': 'success', 'downloaded': downloaded, 'source': 'LFW'})
    except Exception as e:
        return jsonify({'error': str(e), 'downloaded': 0}), 500


@setup_bp.route('/api/setup/fetch_pexels', methods=['POST'])
def fetch_pexels():
    """Fetch faces from Pexels (requires API key)"""
    if not session.get('logged_in') or session.get('role') != 'admin':
        return jsonify({'error': 'Unauthorized'}), 403
    
    data = request.get_json() or {}
    count = int(data.get('count', 50))
    
    # Check for Pexels API key
    pexels_key = os.environ.get('PEXELS_API_KEY')
    if not pexels_key:
        return jsonify({'error': 'Pexels API key not configured. Set PEXELS_API_KEY environment variable.', 'downloaded': 0}), 400
    
    # Placeholder - would integrate with Pexels API
    return jsonify({'status': 'success', 'downloaded': 0, 'message': 'Pexels integration requires API key configuration'})


@setup_bp.route('/api/setup/fetch_utkface', methods=['POST'])
def fetch_utkface():
    """Fetch faces from UTKFace dataset"""
    if not session.get('logged_in') or session.get('role') != 'admin':
        return jsonify({'error': 'Unauthorized'}), 403
    
    data = request.get_json() or {}
    count = int(data.get('count', 100))
    
    # UTKFace requires manual download - provide instructions
    return jsonify({
        'status': 'info',
        'downloaded': 0,
        'message': 'UTKFace requires manual download from Kaggle. Visit: https://www.kaggle.com/jangedoo/utkface-new'
    })


@setup_bp.route('/api/setup/fetch_celeba', methods=['POST'])
def fetch_celeba():
    """Fetch faces from CelebA dataset"""
    if not session.get('logged_in') or session.get('role') != 'admin':
        return jsonify({'error': 'Unauthorized'}), 403
    
    data = request.get_json() or {}
    count = int(data.get('count', 100))
    
    # CelebA requires manual download - provide instructions
    return jsonify({
        'status': 'info',
        'downloaded': 0,
        'message': 'CelebA requires manual download. Visit: https://mmlab.ie.cuhk.edu.hk/projects/CelebA.html'
    })


@setup_bp.route('/api/setup/bootstrap_data', methods=['POST'])
def bootstrap_data():
    """Bootstrap dataset with minimal training data using LFW"""
    if not session.get('logged_in') or session.get('role') != 'admin':
        return jsonify({'error': 'Unauthorized'}), 403
    
    data = request.get_json() or {}
    train_count = int(data.get('train_count', 20))
    val_count = int(data.get('val_count', 20))
    
    ensure_data_dirs()
    
    results = {
        'train_real': 0,
        'val_real': 0,
        'train_morph': 0,
        'val_morph': 0
    }
    
    try:
        from sklearn.datasets import fetch_lfw_people
        import numpy as np
        from PIL import Image
        import uuid
        
        lfw = fetch_lfw_people(min_faces_per_person=10, resize=1.0)
        total_needed = train_count + val_count
        indices = np.random.choice(len(lfw.images), min(total_needed * 2, len(lfw.images)), replace=False)
        
        # Split indices for train and val
        train_indices = indices[:train_count]
        val_indices = indices[train_count:train_count + val_count]
        
        for split_name, split_indices in [('train', train_indices), ('val', val_indices)]:
            target_dir = os.path.join(DATA_DIR, split_name, 'real')
            for idx in split_indices:
                img_array = lfw.images[idx]
                img = Image.fromarray((img_array * 255).astype(np.uint8) if img_array.max() <= 1 else img_array.astype(np.uint8))
                img = img.convert('RGB')
                filename = f"{uuid.uuid4()}.jpg"
                img.save(os.path.join(target_dir, filename), 'JPEG')
                results[f'{split_name}_real'] += 1
        
        return jsonify({'status': 'success', 'results': results})
        
    except Exception as e:
        return jsonify({'error': str(e), 'results': results}), 500


@setup_bp.route('/api/setup/dataset_images', methods=['GET'])
def get_dataset_images():
    """Get list of dataset images for browser"""
    if not session.get('logged_in') or session.get('role') != 'admin':
        return jsonify({'error': 'Unauthorized'}), 403
    
    ensure_data_dirs()
    
    images = []
    for split in ['train', 'val']:
        for category in ['real', 'morph']:
            path = os.path.join(DATA_DIR, split, category)
            if os.path.exists(path):
                for filename in os.listdir(path)[:100]:  # Limit to 100 per category
                    if filename.lower().endswith(('.jpg', '.jpeg', '.png')):
                        images.append({
                            'path': f'/sample_data/{split}/{category}/{filename}',
                            'category': category,
                            'split': split,
                            'filename': filename
                        })
    
    return jsonify({'images': images, 'total': len(images)})