"""
Directory Scanner Module

Scans configured model directories and finds available model files.
"""

import os
import logging
from typing import List, Dict, Tuple

# Import folder_paths lazily - it may not be available until ComfyUI is initialized
try:
    import folder_paths
except ImportError:
    folder_paths = None
    logging.warning("Model Linker: folder_paths not available yet - will retry later")

# Model file extensions to look for
# This matches folder_paths.supported_pt_extensions
MODEL_EXTENSIONS = {'.ckpt', '.pt', '.pt2', '.bin', '.pth', '.safetensors', '.pkl', '.sft', '.onnx'}


def get_model_directories() -> Dict[str, Tuple[List[str], set]]:
    """
    Get all configured model directories from folder_paths.
    
    Returns:
        Dictionary mapping category name to (list of paths, set of extensions)
    """
    global folder_paths
    
    if folder_paths is None:
        # Try to import again
        try:
            import folder_paths as fp
            folder_paths = fp
        except ImportError:
            logging.error("Model Linker: folder_paths still not available")
            return {}
    
    return folder_paths.folder_names_and_paths.copy()


def scan_directory(directory: str, extensions: set, category: str) -> List[Dict[str, str]]:
    """
    Recursively scan a single directory for model files.
    
    Args:
        directory: Absolute path to directory to scan
        extensions: Set of file extensions to look for
        category: Model category name (e.g., 'checkpoints', 'loras')
        
    Returns:
        List of dictionaries with model information:
        {
            'filename': 'model.safetensors',
            'path': 'absolute/path/to/model.safetensors',
            'relative_path': 'subfolder/model.safetensors' or 'model.safetensors',
            'category': 'checkpoints',
            'base_directory': 'absolute/path/to/base'
        }
    """
    models = []
    
    if not os.path.exists(directory) or not os.path.isdir(directory):
        logging.debug(f"Directory does not exist or is not accessible: {directory}")
        return models
    
    try:
        # Get absolute path and normalize
        base_directory = os.path.abspath(directory)
        
        # Walk through directory recursively
        for root, dirs, files in os.walk(base_directory, followlinks=True):
            # Skip hidden directories
            dirs[:] = [d for d in dirs if not d.startswith('.')]
            
            for filename in files:
                # Check if file has a model extension
                file_ext = os.path.splitext(filename)[1].lower()
                
                # For categories with empty extension set, accept all files
                # Otherwise, check if extension matches
                if len(extensions) == 0 or file_ext in extensions or file_ext in MODEL_EXTENSIONS:
                    full_path = os.path.join(root, filename)
                    
                    # Calculate relative path from base directory
                    # IMPORTANT: Use OS-native path separators (backslashes on Windows)
                    # This matches ComfyUI's recursive_search format for get_filename_list
                    try:
                        relative_path = os.path.relpath(full_path, base_directory)
                        # DO NOT normalize - keep OS-native separators to match ComfyUI
                        # ComfyUI's get_filename_list uses os.path.relpath which returns
                        # backslashes on Windows, forward slashes on Unix
                    except ValueError:
                        # If paths are on different drives (Windows), use filename only
                        relative_path = filename
                    
                    models.append({
                        'filename': filename,
                        'path': full_path,
                        'relative_path': relative_path,
                        'category': category,
                        'base_directory': base_directory
                    })
    except (OSError, PermissionError) as e:
        logging.warning(f"Error scanning directory {directory}: {e}")
    
    return models


def scan_all_directories() -> List[Dict[str, str]]:
    """
    Scan all configured model directories and return list of available models.
    
    Returns:
        List of dictionaries with model information (same format as scan_directory)
    """
    all_models = []
    directories = get_model_directories()
    
    for category, (paths, extensions) in directories.items():
        # Skip categories that aren't typically model directories
        if category in ['custom_nodes', 'configs']:
            continue
            
        for directory_path in paths:
            try:
                models = scan_directory(directory_path, extensions, category)
                all_models.extend(models)
                logging.debug(f"Found {len(models)} models in {category}/{directory_path}")
            except Exception as e:
                logging.warning(f"Error scanning {category} directory {directory_path}: {e}")
    
    return all_models


def get_model_files() -> List[Dict[str, str]]:
    """
    Get list of all available model files with metadata.
    
    This is the main entry point for getting model files.
    
    Returns:
        List of model dictionaries (same format as scan_directory)
    """
    return scan_all_directories()