#!/bin/bash

# Script to organize image and video files from Downloads and home directory
# Zips all media files with timestamp and removes originals
# Author: Generated for linux-configuration

# Set strict error handling
set -euo pipefail

# Defaults / flags
DRY_RUN=false
SAMPLE_LIMIT=20

# Simple usage helper
usage() {
        cat <<EOF
Usage: $(basename "$0") [--dry-run|-n] [--sample=N]

Options:
    -n, --dry-run     Analyze and print what would be archived without creating a zip or removing files.
            --sample=N    In dry-run, show up to N sample file paths for select extensions (default: $SAMPLE_LIMIT).
    -h, --help        Show this help.
EOF
}

# Define directories to scan
DOWNLOADS_DIR="$HOME/Downloads"
HOME_DIR="$HOME"
# Prefer a temp dir outside Downloads to avoid recursive re-inclusion; fall back to /tmp
TEMP_DIR="${XDG_CACHE_HOME:-$HOME/.cache}/media_organize_$$"

# Define image and video file extensions
# Keep common raster/image formats; exclude svg/ico by default (often project assets), can be re-added if needed
IMAGE_EXTENSIONS=("jpg" "jpeg" "png" "gif" "bmp" "tiff" "tif" "webp" "raw" "cr2" "nef" "orf" "arw" "dng" "heic" "heif")
VIDEO_EXTENSIONS=("mp4" "avi" "mkv" "mov" "wmv" "flv" "webm" "m4v" "3gp" "ogv" "mpg" "mpeg" "mts" "m2ts" "vob")

# Function to log messages with timestamp
log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1"
}

# Parse CLI flags early
while [[ ${1:-} =~ ^- ]]; do
    case "${1}" in
        -n|--dry-run)
            DRY_RUN=true
            shift
            ;;
        --sample=*)
            SAMPLE_LIMIT="${1#*=}"
            shift
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        *)
            echo "Unknown option: $1" >&2
            usage
            exit 1
            ;;
    esac
done

# Function to check if file has media extension
is_media_file() {
    local file="$1"
    local extension="${file##*.}"
    extension=$(echo "$extension" | tr '[:upper:]' '[:lower:]')
    
    # Check if it's an image
    for ext in "${IMAGE_EXTENSIONS[@]}"; do
        if [[ "$extension" == "$ext" ]]; then
            return 0
        fi
    done
    
    # Check if it's a video
    for ext in "${VIDEO_EXTENSIONS[@]}"; do
        if [[ "$extension" == "$ext" ]]; then
            return 0
        fi
    done
    
    return 1
}

# Function to find media files in a directory (non-recursive for home, avoid common system dirs)
find_media_files() {
    local search_dir="$1"
    local files=()
    # Directories to exclude under Downloads
    local -a EXCLUDES=(
        ".git" ".hg" ".svn" ".cache" "node_modules" "dist" "build" "out" "target" "coverage" "__pycache__" "venv" ".venv"
        # previous staging dirs created by this script
        ".media_organize_" "media_organize_"
    )
    
    if [[ "$search_dir" == "$HOME_DIR" ]]; then
        # For home directory, only check files directly in ~ (not subdirectories)
        # Exclude common system/config directories
        while IFS= read -r -d '' file; do
            local basename=$(basename "$file")
            # Skip hidden files and common system directories
            if [[ ! "$basename" =~ ^\. ]] && [[ -f "$file" ]]; then
                if is_media_file "$file"; then
                    files+=("$file")
                fi
            fi
        done < <(find "$search_dir" -maxdepth 1 -type f -print0 2>/dev/null)
    else
        # For Downloads, search recursively, pruning excluded directories
        # Build prune expression
        local prune_expr=()
        for ex in "${EXCLUDES[@]}"; do
            prune_expr+=( -name "$ex*" -o )
        done
        # Remove trailing -o
        unset 'prune_expr[${#prune_expr[@]}-1]'

        while IFS= read -r -d '' file; do
            if is_media_file "$file"; then
                files+=("$file")
            fi
        done < <(find "$search_dir" \( -type d \( ${prune_expr[@]} \) -prune \) -o -type f -print0 2>/dev/null)
    fi
    
    printf '%s\n' "${files[@]}"
}

# Function to create timestamped zip archive
create_media_archive() {
    local files=("$@")
    
    if [[ ${#files[@]} -eq 0 ]]; then
        log "No media files found to archive."
        return 0
    fi
    
    # Create timestamp for archive name
    local timestamp=$(date '+%Y%m%d_%H%M%S')
    local archive_name="media_archive_${timestamp}.zip"
    local archive_path="$DOWNLOADS_DIR/$archive_name"
    
    # Create temporary directory (fallback to /tmp if needed)
    if ! mkdir -p "$TEMP_DIR" 2>/dev/null; then
        TEMP_DIR="/tmp/media_organize_$$"
        mkdir -p "$TEMP_DIR"
    fi

    # Ensure temp dir is cleaned up on function return; trap unsets itself after running
    trap 'rm -rf "$TEMP_DIR" 2>/dev/null || true; trap - RETURN' RETURN
    
    log "Found ${#files[@]} media files to archive."
    log "Creating archive: $archive_path"
    
    # Copy files to temp directory maintaining relative structure
    local successfully_copied=()
    local copy_errors=0
    
    for file in "${files[@]}"; do
        if [[ -f "$file" ]]; then
            local relative_path=""
            if [[ "$file" == "$DOWNLOADS_DIR"* ]]; then
                relative_path="downloads/${file#$DOWNLOADS_DIR/}"
            else
                relative_path="home/${file#$HOME_DIR/}"
            fi
            
            local temp_file="$TEMP_DIR/$relative_path"
            local temp_dir=$(dirname "$temp_file")
            
            mkdir -p "$temp_dir"
            # Check readability first to provide a clearer error
            if [[ ! -r "$file" ]]; then
                log "WARNING: Cannot read $file (permission denied?)"
                ((copy_errors++))
                continue
            fi

            # Attempt copy and capture any error for logging
            local cp_err
            if cp_err=$(cp -p "$file" "$temp_file" 2>&1); then
                successfully_copied+=("$file")
            else
                # Surface the cp error so the user can see the reason
                log "WARNING: Failed to copy $file -> $cp_err"
                # Special hint for space issues
                if echo "$cp_err" | grep -qi "No space left on device"; then
                    log "HINT: Not enough free space to stage files. Using $TEMP_DIR. Free up space or change TEMP_DIR."
                fi
                ((copy_errors++))
            fi
        fi
    done
    
    if [[ ${#successfully_copied[@]} -eq 0 ]]; then
        log "ERROR: No files were successfully copied to temp directory."
        rm -rf "$TEMP_DIR"
        return 1
    fi
    
    if [[ $copy_errors -gt 0 ]]; then
        log "WARNING: $copy_errors files failed to copy."
    fi
    
    # Create zip archive with maximum compression
    log "Creating zip archive with ${#successfully_copied[@]} files..."
    cd "$TEMP_DIR"
    if zip -9 -r "$archive_path" . 2>&1; then
        log "Successfully created archive with ${#successfully_copied[@]} files."
        
        # Verify the zip file was actually created and is not empty
        if [[ ! -f "$archive_path" ]]; then
            log "ERROR: Archive file was not created at $archive_path"
            rm -rf "$TEMP_DIR"
            return 1
        fi
        
        local archive_size=$(stat -c%s "$archive_path" 2>/dev/null || echo "0")
        if [[ "$archive_size" -eq 0 ]]; then
            log "ERROR: Archive file is empty"
            rm -rf "$TEMP_DIR"
            return 1
        fi
        
        # Remove original files only if zip was successful
        local removed_count=0
        local remove_errors=0
        
        log "Starting to remove ${#successfully_copied[@]} original files..."
        
        # Temporarily disable strict error handling for file removal
        set +e
        
        for file in "${successfully_copied[@]}"; do
            if [[ -f "$file" ]]; then
                if rm "$file" 2>/dev/null; then
                    removed_count=$((removed_count + 1))
                    log "Removed: $(basename "$file")"
                else
                    remove_errors=$((remove_errors + 1))
                    log "ERROR: Failed to remove $(basename "$file")"
                fi
            else
                log "WARNING: File no longer exists: $(basename "$file")"
            fi
        done
        
        # Re-enable strict error handling
        set -e
        
        log "Successfully removed $removed_count original files."
        if [[ $remove_errors -gt 0 ]]; then
            log "WARNING: Failed to remove $remove_errors files."
        fi
        log "Archive size: $(du -h "$archive_path" | cut -f1)"
        
    # Cleanup temp directory (trap will also attempt, which is safe)
    rm -rf "$TEMP_DIR"
        
        # Return success only if we removed files or there were no files to remove
        if [[ $removed_count -gt 0 ]] || [[ ${#successfully_copied[@]} -eq 0 ]]; then
            return 0
        else
            log "ERROR: Failed to remove any files after successful archive creation."
            return 1
        fi
    else
        log "ERROR: Failed to create archive. Original files preserved."
        log "Zip command failed."
    rm -rf "$TEMP_DIR"
        return 1
    fi
}

# Main execution
main() {
    log "Starting media file organization..."
    
    # Check if required directories exist
    if [[ ! -d "$DOWNLOADS_DIR" ]]; then
        log "ERROR: Downloads directory not found: $DOWNLOADS_DIR"
        exit 1
    fi
    
    if [[ ! -d "$HOME_DIR" ]]; then
        log "ERROR: Home directory not found: $HOME_DIR"
        exit 1
    fi
    
    # Check if zip command is available
    if ! command -v zip >/dev/null 2>&1; then
        log "ERROR: zip command not found. Please install zip package."
        exit 1
    fi
    
    # Find all media files
    log "Scanning for media files..."
    local all_files=()
    
    # Find files in Downloads directory
    log "Scanning Downloads directory..."
    while IFS= read -r file; do
        [[ -n "$file" ]] && all_files+=("$file")
    done < <(find_media_files "$DOWNLOADS_DIR")
    
    # Find files in home directory (only direct files, not subdirectories)
    log "Scanning home directory (root level only)..."
    while IFS= read -r file; do
        [[ -n "$file" ]] && all_files+=("$file")
    done < <(find_media_files "$HOME_DIR")
    
    if $DRY_RUN; then
        log "Dry-run mode: summarizing what would be archived."
        if [[ ${#all_files[@]} -eq 0 ]]; then
            log "No media files found to organize."
            exit 0
        fi

        # Count by extension
        declare -A ext_counts=()
        # Count by top-level directory under Downloads
        declare -A dir_counts=()
        # Sample paths for suspect extensions
        declare -A samples_ts=()
        declare -A samples_svg=()
        declare -A samples_ico=()

        for f in "${all_files[@]}"; do
            # Extension
            ext="${f##*.}"
            ext="${ext,,}"
            ((ext_counts["$ext"]++)) || true

            # Top directory under Downloads
            if [[ "$f" == "$DOWNLOADS_DIR"/* ]]; then
                rel="${f#"$DOWNLOADS_DIR"/}"
                topdir="${rel%%/*}"
                [[ "$topdir" == "$rel" ]] && topdir="."
                ((dir_counts["$topdir"]++)) || true
            else
                ((dir_counts["~"]++)) || true
            fi

            # Samples for suspect extensions
            case "$ext" in
                ts)
                    if [[ ${#samples_ts[@]} -lt $SAMPLE_LIMIT ]]; then samples_ts["$f"]=1; fi
                    ;;
                svg)
                    if [[ ${#samples_svg[@]} -lt $SAMPLE_LIMIT ]]; then samples_svg["$f"]=1; fi
                    ;;
                ico)
                    if [[ ${#samples_ico[@]} -lt $SAMPLE_LIMIT ]]; then samples_ico["$f"]=1; fi
                    ;;
            esac
        done

        echo ""
        echo "Summary by extension (top 20):"
        for k in "${!ext_counts[@]}"; do
            printf "%8d %s\n" "${ext_counts[$k]}" "$k"
        done | sort -nr | head -n 20

        echo ""
        echo "Top contributing directories under Downloads (top 20):"
        for k in "${!dir_counts[@]}"; do
            printf "%8d %s\n" "${dir_counts[$k]}" "$k"
        done | sort -nr | head -n 20

        echo ""
        if [[ ${#samples_ts[@]} -gt 0 ]]; then
            echo "Sample .ts files (TypeScript vs Transport Stream) up to $SAMPLE_LIMIT:"
            for p in "${!samples_ts[@]}"; do echo "  $p"; done | sort
            echo ""
        fi
        if [[ ${#samples_svg[@]} -gt 0 ]]; then
            echo "Sample .svg files up to $SAMPLE_LIMIT:"
            for p in "${!samples_svg[@]}"; do echo "  $p"; done | sort
            echo ""
        fi
        if [[ ${#samples_ico[@]} -gt 0 ]]; then
            echo "Sample .ico files up to $SAMPLE_LIMIT:"
            for p in "${!samples_ico[@]}"; do echo "  $p"; done | sort
            echo ""
        fi

        echo "Total files that would be archived: ${#all_files[@]}"
        echo "(Use: $(basename "$0") --dry-run --sample=50 to see more examples.)"
        exit 0
    fi

    # Create archive if files found
    if [[ ${#all_files[@]} -gt 0 ]]; then
        create_media_archive "${all_files[@]}"
        log "Media organization completed successfully."
    else
        log "No media files found to organize."
    fi
}

# Run main function
main "$@"