cluster4npu/utils/folder_dialog.py

"""
Folder selection utilities using PyQt5 as primary, tkinter as fallback
"""

import os

def select_folder(title="Select Folder", initial_dir="", must_exist=True):
    """
    Open a folder selection dialog using PyQt5 (preferred) or tkinter (fallback)
    
    Args:
        title (str): Dialog window title
        initial_dir (str): Initial directory to open
        must_exist (bool): Whether the folder must already exist
        
    Returns:
        str: Selected folder path, or empty string if cancelled
    """
    # Try PyQt5 first (more reliable on macOS)
    try:
        from PyQt5.QtWidgets import QApplication, QFileDialog
        import sys
        
        # Create QApplication if it doesn't exist
        app = QApplication.instance()
        if app is None:
            app = QApplication(sys.argv)
        
        # Set initial directory
        if not initial_dir:
            initial_dir = os.getcwd()
        elif not os.path.exists(initial_dir):
            initial_dir = os.getcwd()
        
        # Open folder selection dialog
        folder_path = QFileDialog.getExistingDirectory(
            None,
            title,
            initial_dir,
            QFileDialog.ShowDirsOnly | QFileDialog.DontResolveSymlinks
        )
        
        return folder_path if folder_path else ""
        
    except ImportError:
        print("PyQt5 not available, trying tkinter...")
        
        # Fallback to tkinter
        try:
            import tkinter as tk
            from tkinter import filedialog
            
            # Create a root window but keep it hidden
            root = tk.Tk()
            root.withdraw()  # Hide the main window
            root.attributes('-topmost', True)  # Bring dialog to front
            
            # Set initial directory
            if not initial_dir:
                initial_dir = os.getcwd()
            
            # Open folder selection dialog
            folder_path = filedialog.askdirectory(
                title=title,
                initialdir=initial_dir,
                mustexist=must_exist
            )
            
            # Destroy the root window
            root.destroy()
            
            return folder_path if folder_path else ""
            
        except ImportError:
            print("tkinter also not available")
            return ""
        except Exception as e:
            print(f"Error opening tkinter folder dialog: {e}")
            return ""
    
    except Exception as e:
        print(f"Error opening PyQt5 folder dialog: {e}")
        return ""

def select_assets_folder(initial_dir=""):
    """
    Specialized function for selecting Assets folder with validation
    
    Args:
        initial_dir (str): Initial directory to open
        
    Returns:
        dict: Result with 'path', 'valid', and 'message' keys
    """
    folder_path = select_folder(
        title="Select Assets Folder (containing Firmware/ and Models/)",
        initial_dir=initial_dir
    )
    
    if not folder_path:
        return {'path': '', 'valid': False, 'message': 'No folder selected'}
    
    # Validate folder structure
    validation_result = validate_assets_folder_structure(folder_path)
    
    return {
        'path': folder_path,
        'valid': validation_result['valid'],
        'message': validation_result['message'],
        'details': validation_result.get('details', {})
    }

def validate_assets_folder_structure(folder_path):
    """
    Validate that a folder has the expected Assets structure
    
    Expected structure:
    Assets/
    ├── Firmware/
    │   ├── KL520/
    │   │   ├── fw_scpu.bin
    │   │   └── fw_ncpu.bin
    │   └── KL720/
    │       ├── fw_scpu.bin
    │       └── fw_ncpu.bin
    └── Models/
        ├── KL520/
        │   └── model.nef
        └── KL720/
            └── model.nef
    
    Args:
        folder_path (str): Path to validate
        
    Returns:
        dict: Validation result with 'valid', 'message', and 'details' keys
    """
    if not os.path.exists(folder_path):
        return {'valid': False, 'message': 'Folder does not exist'}
    
    if not os.path.isdir(folder_path):
        return {'valid': False, 'message': 'Path is not a directory'}
    
    details = {}
    issues = []
    
    # Check for Firmware and Models folders
    firmware_path = os.path.join(folder_path, 'Firmware')
    models_path = os.path.join(folder_path, 'Models')
    
    has_firmware = os.path.exists(firmware_path) and os.path.isdir(firmware_path)
    has_models = os.path.exists(models_path) and os.path.isdir(models_path)
    
    details['has_firmware_folder'] = has_firmware
    details['has_models_folder'] = has_models
    
    if not has_firmware:
        issues.append("Missing 'Firmware' folder")
    
    if not has_models:
        issues.append("Missing 'Models' folder")
    
    if not (has_firmware and has_models):
        return {
            'valid': False,
            'message': f"Invalid folder structure: {', '.join(issues)}",
            'details': details
        }
    
    # Check for series subfolders
    expected_series = ['KL520', 'KL720', 'KL630', 'KL730', 'KL540']
    
    firmware_series = []
    models_series = []
    
    try:
        firmware_dirs = [d for d in os.listdir(firmware_path) 
                        if os.path.isdir(os.path.join(firmware_path, d))]
        firmware_series = [d for d in firmware_dirs if d in expected_series]
        
        models_dirs = [d for d in os.listdir(models_path)
                      if os.path.isdir(os.path.join(models_path, d))]
        models_series = [d for d in models_dirs if d in expected_series]
        
    except Exception as e:
        return {
            'valid': False,
            'message': f"Error reading folder contents: {e}",
            'details': details
        }
    
    details['firmware_series'] = firmware_series
    details['models_series'] = models_series
    
    # Find common series (have both firmware and models)
    common_series = list(set(firmware_series) & set(models_series))
    details['available_series'] = common_series
    
    if not common_series:
        return {
            'valid': False,
            'message': "No series found with both firmware and models folders",
            'details': details
        }
    
    # Check for actual files in series folders
    series_with_files = []
    for series in common_series:
        has_files = False
        
        # Check firmware files
        fw_series_path = os.path.join(firmware_path, series)
        if os.path.exists(fw_series_path):
            fw_files = [f for f in os.listdir(fw_series_path) 
                       if f.endswith('.bin')]
            if fw_files:
                has_files = True
        
        # Check model files
        model_series_path = os.path.join(models_path, series)
        if os.path.exists(model_series_path):
            model_files = [f for f in os.listdir(model_series_path) 
                          if f.endswith('.nef')]
            if model_files and has_files:
                series_with_files.append(series)
    
    details['series_with_files'] = series_with_files
    
    if not series_with_files:
        return {
            'valid': False,
            'message': "No series found with actual firmware and model files",
            'details': details
        }
    
    return {
        'valid': True,
        'message': f"Valid Assets folder with {len(series_with_files)} series: {', '.join(series_with_files)}",
        'details': details
    }

# Example usage
if __name__ == "__main__":
    print("Testing folder selection...")
    
    # Test basic folder selection
    folder = select_folder("Select any folder")
    print(f"Selected: {folder}")
    
    # Test Assets folder selection with validation
    result = select_assets_folder()
    print(f"Assets folder result: {result}")