app/utils/file_utils.py

"""
File utilities for Image Tagger application.
"""

import os
import time

def save_tags_to_file(image_path, all_tags, original_filename=None, custom_dir=None, overwrite=False):
    """
    Save tags to a text file in a dedicated 'saved_tags' folder or custom directory.
    
    Args:
        image_path: Path to the original image
        all_tags: List of all tags to save
        original_filename: Original filename if uploaded through Streamlit
        custom_dir: Custom directory to save tags to (if None, uses 'saved_tags' folder)
    
    Returns:
        Path to the saved file
    """
    # Determine the save directory
    if custom_dir and os.path.isdir(custom_dir):
        save_dir = custom_dir
    else:
        # Create a dedicated folder for saved tags in the app's root directory
        app_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        save_dir = os.path.join(app_dir, "saved_tags")
    
    # Ensure the directory exists
    os.makedirs(save_dir, exist_ok=True)
    
    # Determine the filename
    if original_filename:
        # For uploaded files, use original filename
        base_name = os.path.splitext(original_filename)[0]
    else:
        # For non-uploaded files, use the image path
        base_name = os.path.splitext(os.path.basename(image_path))[0]
    
    # Create the output path
    output_path = os.path.join(save_dir, f"{base_name}.txt")
    
    # If overwrite is False and file exists, add a timestamp to avoid overwriting
    if not overwrite and os.path.exists(output_path):
        timestamp = time.strftime("%Y%m%d-%H%M%S")
        output_path = os.path.join(save_dir, f"{base_name}_{timestamp}.txt")
    
    # Write the tags to file
    with open(output_path, 'w', encoding='utf-8') as f:
        if all_tags:
            # Add comma after each tag including the last one
            tag_text = ", ".join(all_tags) + ","
            f.write(tag_text)
    
    return output_path

def get_default_save_locations():
    """
    Get default save locations for tag files.
    
    Returns:
        List of default save locations
    """
    # App directory
    app_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    save_dir = os.path.join(app_dir, "saved_tags")
    
    # Common user directories
    desktop_dir = os.path.expanduser("~/Desktop")
    download_dir = os.path.expanduser("~/Downloads")
    documents_dir = os.path.expanduser("~/Documents")
    
    # List of default save locations
    save_locations = [
        save_dir,
        desktop_dir,
        download_dir,
        documents_dir,
    ]
    
    # Ensure directories exist
    for folder in save_locations:
        os.makedirs(folder, exist_ok=True)
    
    return save_locations

def apply_category_limits(result, category_limits):
    """
    Apply category limits to a result dictionary.
    
    Args:
        result: Result dictionary containing tags and all_tags
        category_limits: Dictionary mapping categories to their tag limits
                         (0 = exclude category, -1 = no limit/include all)
        
    Returns:
        Updated result dictionary with limits applied
    """
    if not category_limits or not result['success']:
        return result
    
    # Get the filtered tags
    filtered_tags = result['tags']
    
    # Apply limits to each category
    for category, cat_tags in list(filtered_tags.items()):
        # Get limit for this category, default to -1 (no limit)
        limit = category_limits.get(category, -1)
        
        if limit == 0:
            # Exclude this category entirely
            del filtered_tags[category]
        elif limit > 0 and len(cat_tags) > limit:
            # Limit to top N tags for this category
            filtered_tags[category] = cat_tags[:limit]
    
    # Regenerate all_tags list after applying limits
    all_tags = []
    for category, cat_tags in filtered_tags.items():
        for tag, _ in cat_tags:
            all_tags.append(tag)
    
    # Update the result with limited tags
    result['tags'] = filtered_tags
    result['all_tags'] = all_tags
    
    return result