#!/bin/bash

# Author: oneiricsky
# License: MIT
# Recursively moves audio files from /downloads to /music using metadata or best-guess parsing.

set -euo pipefail  # Exit on error, undefined vars, pipe failures

# Default configuration - modify these as needed
DOWNLOADS_DIR="/downloads"
MUSIC_DIR="/shared"
LOG_FILE="./music_organizer.log"

# Plex configuration
PLEX_URL="http://plex:32400"
PLEX_TOKEN="your-token-here"
PLEX_LIBRARY_SECTION="Music"
plex_refresh_override=""

readonly SCRIPT_NAME="$(basename "$0")"

# Performance: Pre-compile regex patterns
readonly YEAR_ALBUM_REGEX='^([0-9]{4})[[:space:]]*-[[:space:]]*(.+)$'
readonly ARTIST_ALBUM_REGEX='^([^-]+)[[:space:]]*-[[:space:]]*(.+)$'
readonly COMPLEX_TRACK_REGEX='^([0-9]{1,2})[[:space:]]*-[[:space:]]*([^-]+)[[:space:]]*-[[:space:]]*([^-]+)[[:space:]]*-[[:space:]]*([0-9]{1,2})[[:space:]]*[-\.]?[[:space:]]*(.+)$'
readonly TRACK_TITLE_REGEX='^([0-9]{1,2})[[:space:]]*-[[:space:]]*(.+)$'
readonly TRACK_DOT_TITLE_REGEX='^([0-9]{1,2})\.[[:space:]]*(.+)$'
readonly ARTIST_TRACK_TITLE_REGEX='^([^-]+)[[:space:]]*-[[:space:]]*([0-9]{1,2})[[:space:]]*-[[:space:]]*(.+)$'
readonly TRACK_ARTIST_TITLE_REGEX='^([0-9]{1,2})-([^-]+)-(.+)$'

# Audio file extensions
readonly -a AUDIO_EXTENSIONS=("mp3" "flac" "m4a" "wav" "ogg" "aac" "wma" "opus" "mp4" "webm")

# Statistics
declare -g processed_count=0 failed_count=0 skipped_count=0

log() {
    local level="${2:-INFO}"
    printf '[%s] [%s] %s\n' "$(date '+%Y-%m-%d %H:%M:%S')" "$level" "$1" | tee -a "$LOG_FILE"
}

error_exit() {
    log "$1" "ERROR" >&2
    exit 1
}

# Optimized sanitization with single pass
sanitize() {
    local input="$1"
    # Remove forbidden characters and normalize whitespace in one pass
    printf '%s' "$input" | \
        tr -d '<>:"|?*' | \
        tr '/\\' '_' | \
        sed 's/[[:space:]]\+/ /g; s/^[[:space:]]*//; s/[[:space:]]*$//'
}

pad_track() {
    local track="$1"
    # Extract only the numeric part before any slash
    track="${track%%/*}"
    track="${track//[^0-9]/}"
    
    if [[ "$track" =~ ^[0-9]+$ ]] && (( track > 0 )); then
        printf "%02d" "$track"
    else
        printf "00"
    fi
}

# Cache metadata tools availability
declare -g HAS_METAFLAC HAS_CURL HAS_RSYNC
check_tools() {
    command -v ffprobe &>/dev/null || error_exit "ffprobe required. Please install ffmpeg."
    HAS_METAFLAC=$(command -v metaflac &>/dev/null && echo "true" || echo "false")
    HAS_CURL=$(command -v curl &>/dev/null && echo "true" || echo "false")
    HAS_RSYNC=$(command -v rsync &>/dev/null && echo "true" || echo "false")
    
    if [[ "$HAS_METAFLAC" == "true" ]]; then
        log "metaflac found - will use for FLAC files"
    else
        log "metaflac not found - falling back to ffprobe for FLAC files" "WARN"
    fi
    if [[ "$HAS_RSYNC" == "true" ]]; then
        log "rsync found - will use for file moves"
    else
        log "rsync not found - will fallback to cp/rm for moving files" "WARN"
    fi
}

# Optimized metadata extraction - FLAC gets priority with metaflac
get_metadata() {
    local file="$1" field="$2"
    local result
    
    # For FLAC files, try metaflac first if available
    if [[ "${file,,}" == *.flac ]] && [[ "$HAS_METAFLAC" == "true" ]]; then
        local flac_field
        case "${field,,}" in
            artist) flac_field="ARTIST" ;;
            albumartist) flac_field="ALBUMARTIST" ;;
            album) flac_field="ALBUM" ;;
            title) flac_field="TITLE" ;;
            track) flac_field="TRACKNUMBER" ;;
            date|year) flac_field="DATE" ;;
            genre) flac_field="GENRE" ;;
            *) flac_field="${field^^}" ;;
        esac
        
        if result=$(metaflac --show-tag="$flac_field" "$file" 2>/dev/null | \
                   cut -d= -f2- | head -1 | sed 's/^[[:space:]]*//; s/[[:space:]]*$//'); then
            if [[ -n "$result" ]]; then
                printf '%s' "$result"
                return 0
            fi
        fi
        
        # Try alternative FLAC field names for track numbers
        if [[ "${field,,}" == "track" ]]; then
            for tf in "TRACK_NUMBER" "TRACK"; do
                if result=$(metaflac --show-tag="$tf" "$file" 2>/dev/null | \
                           cut -d= -f2- | head -1 | sed 's/^[[:space:]]*//; s/[[:space:]]*$//'); then
                    if [[ -n "$result" ]]; then
                        printf '%s' "$result"
                        return 0
                    fi
                fi
            done
        fi
        
        # Try alternative FLAC field names for albumartist
        if [[ "${field,,}" == "albumartist" ]]; then
            for af in "ALBUM_ARTIST" "AlbumArtist"; do
                if result=$(metaflac --show-tag="$af" "$file" 2>/dev/null | \
                           cut -d= -f2- | head -1 | sed 's/^[[:space:]]*//; s/[[:space:]]*$//'); then
                    if [[ -n "$result" ]]; then
                        printf '%s' "$result"
                        return 0
                    fi
                fi
            done
        fi
    fi
    
    # Fallback to ffprobe for all files (including FLAC if metaflac failed)
    local -a cmd_variants=("$field" "${field^^}")
    local -a locations=("format_tags" "stream_tags")
    
    for location in "${locations[@]}"; do
        for variant in "${cmd_variants[@]}"; do
            if result=$(ffprobe -v quiet -show_entries "$location=$variant" \
                       -of default=noprint_wrappers=1:nokey=1 "$file" 2>/dev/null | \
                       head -1 | sed 's/^[[:space:]]*//; s/[[:space:]]*$//'); then
                if [[ -n "$result" ]]; then
                    printf '%s' "$result"
                    return 0
                fi
            fi
        done
    done
    
    # Special handling for track numbers with ffprobe
    if [[ "${field,,}" == "track" ]]; then
        local -a track_fields=("tracknumber" "TRACKNUMBER" "track_number" "TRACK_NUMBER")
        for location in "${locations[@]}"; do
            for tf in "${track_fields[@]}"; do
                if result=$(ffprobe -v quiet -show_entries "$location=$tf" \
                           -of default=noprint_wrappers=1:nokey=1 "$file" 2>/dev/null | \
                           head -1 | sed 's/^[[:space:]]*//; s/[[:space:]]*$//'); then
                    if [[ -n "$result" ]]; then
                        printf '%s' "$result"
                        return 0
                    fi
                fi
            done
        done
    fi
    
    return 1
}

# Optimized path parsing with pre-compiled regex
parse_from_path() {
    local path="$1"
    local filename dirname name_no_ext
    filename="$(basename "$path")"
    dirname="$(basename "$(dirname "$path")")"
    name_no_ext="${filename%.*}"

    local artist="" album="" title="" track="" year=""

    # Parse directory structure
    if [[ "$dirname" =~ $YEAR_ALBUM_REGEX ]]; then
        year="${BASH_REMATCH[1]}"
        album="${BASH_REMATCH[2]}"
    elif [[ "$dirname" =~ $ARTIST_ALBUM_REGEX ]]; then
        artist="${BASH_REMATCH[1]}"
        album="${BASH_REMATCH[2]}"
    else
        album="$dirname"
    fi

    # Parse filename with optimized regex matching
    if [[ "$name_no_ext" =~ $COMPLEX_TRACK_REGEX ]]; then
        artist="${BASH_REMATCH[2]}"
        album="${BASH_REMATCH[3]}"
        track="${BASH_REMATCH[4]}"
        title="${BASH_REMATCH[5]}"
    elif [[ "$name_no_ext" =~ $TRACK_TITLE_REGEX ]]; then
        track="${BASH_REMATCH[1]}"
        title="${BASH_REMATCH[2]}"
    elif [[ "$name_no_ext" =~ $TRACK_DOT_TITLE_REGEX ]]; then
        track="${BASH_REMATCH[1]}"
        title="${BASH_REMATCH[2]}"
    elif [[ "$name_no_ext" =~ $ARTIST_TRACK_TITLE_REGEX ]]; then
        [[ -z "$artist" ]] && artist="${BASH_REMATCH[1]}"
        track="${BASH_REMATCH[2]}"
        title="${BASH_REMATCH[3]}"
    elif [[ "$name_no_ext" =~ $TRACK_ARTIST_TITLE_REGEX ]]; then
        track="${BASH_REMATCH[1]}"
        [[ -z "$artist" ]] && artist="${BASH_REMATCH[2]}"
        title="${BASH_REMATCH[3]}"
    else
        title="$name_no_ext"
    fi

    # Clean up title by removing redundant artist/album prefixes
    if [[ -n "$artist" && -n "$title" ]]; then
        title="${title#"$artist"}"
        title="${title#*([[:space:]])-*([[:space:]])}"
    fi
    if [[ -n "$album" && -n "$title" ]]; then
        title="${title#"$album"}"
        title="${title#*([[:space:]])-*([[:space:]]*)}"
    fi

    # Output sanitized results
    printf "PARSED_ARTIST:%s\n" "$(sanitize "$artist")"
    printf "PARSED_ALBUM:%s\n" "$(sanitize "$album")"
    printf "PARSED_TITLE:%s\n" "$(sanitize "$title")"
    printf "PARSED_TRACK:%s\n" "$track"
    printf "PARSED_YEAR:%s\n" "$year"
}

# Plex library refresh functionality
refresh_plex_library() {
    if [[ "$plex_refresh_override" != "true" ]]; then
        log "Plex refresh disabled"
        return 0
    fi
    
    if [[ -z "$PLEX_TOKEN" ]]; then
        log "Plex token not provided - skipping refresh" "WARN"
        return 1
    fi
    
    if [[ "$HAS_CURL" != "true" ]]; then
        log "curl not available - cannot refresh Plex library" "WARN"
        return 1
    fi
    
    local plex_sections_url="$PLEX_URL/library/sections?X-Plex-Token=$PLEX_TOKEN"
    local section_id=""
    
    # If specific section provided, use it; otherwise find music library
    if [[ -n "$PLEX_LIBRARY_SECTION" ]]; then
        if [[ "$PLEX_LIBRARY_SECTION" =~ ^[0-9]+$ ]]; then
            section_id="$PLEX_LIBRARY_SECTION"
        else
            # Find section by name
            log "Looking up Plex library section: $PLEX_LIBRARY_SECTION"
            local sections_xml
            if sections_xml=$(curl -s -f -m 10 "$plex_sections_url" 2>/dev/null); then
                section_id=$(echo "$sections_xml" | grep -i "title=\"$PLEX_LIBRARY_SECTION\"" | sed -n 's/.*key="\([0-9]*\)".*/\1/p' | head -1)
            fi
        fi
    else
        # Auto-discover music library
        log "Auto-discovering Plex music library"
        local sections_xml
        if sections_xml=$(curl -s -f -m 10 "$plex_sections_url" 2>/dev/null); then
            section_id=$(echo "$sections_xml" | grep -i 'type="artist"' | sed -n 's/.*key="\([0-9]*\)".*/\1/p' | head -1)
        fi
    fi
    
    if [[ -z "$section_id" ]]; then
        log "Could not find Plex music library section" "WARN"
        return 1
    fi
    
    log "Refreshing Plex library section: $section_id"
    local refresh_url="$PLEX_URL/library/sections/$section_id/refresh?X-Plex-Token=$PLEX_TOKEN"
    
    if curl -s -f -m 30 -X POST "$refresh_url" >/dev/null 2>&1; then
        log "SUCCESS: Plex library refresh triggered"
        return 0
    else
        log "ERROR: Failed to refresh Plex library" "ERROR"
        return 1
    fi
}

# Get Plex server info for validation
validate_plex_connection() {
    if [[ "$plex_refresh_override" == null ]]; then
        return 0
    fi
    
    if [[ -z "$PLEX_TOKEN" ]] || [[ "$HAS_CURL" != "true" ]]; then
        return 0
    fi
    
    log "Validating Plex connection..."
    local identity_url="$PLEX_URL/identity?X-Plex-Token=$PLEX_TOKEN"
    
    if curl -s -f -m 10 "$identity_url" >/dev/null 2>&1; then
        log "Plex connection validated"
        return 0
    else
        log "Failed to connect to Plex server - refresh will be skipped" "WARN"
        return 1
    fi
}

# Move file using rsync if available, fallback to cp/rm if not
move_file() {
    local file="$1"
    local filename extension target_dir target_path counter

    filename="$(basename "$file")"
    extension="${filename##*.}"

    log "Processing: $filename"

    # Declare metadata variables
    local album_artist artist album title track
    local got_metadata=false

    # Batch metadata extraction
    local -A metadata
    local -a fields=("albumartist" "artist" "album" "title" "track")
    
    for field in "${fields[@]}"; do
        if metadata["$field"]="$(get_metadata "$file" "$field")" && [[ -n "${metadata["$field"]}" ]]; then
            got_metadata=true
        fi
    done

    # Assign metadata to variables
    album_artist="${metadata[albumartist]}"
    artist="${metadata[artist]}"
    album="${metadata[album]}"
    title="${metadata[title]}"
    track="${metadata[track]}"

    if [[ "$got_metadata" == "true" ]]; then
        log "Extracted metadata - Artist: '$artist', Album: '$album', Title: '$title', Track: '$track', AlbumArtist: '$album_artist'"
    else
        log "No embedded metadata found, parsing from path and filename"
        local parsed_output parsed_year
        parsed_output="$(parse_from_path "$file")"
        
        # Extract parsed values efficiently
        artist="${artist:-$(printf '%s' "$parsed_output" | sed -n 's/PARSED_ARTIST://p')}"
        album="${album:-$(printf '%s' "$parsed_output" | sed -n 's/PARSED_ALBUM://p')}"
        title="${title:-$(printf '%s' "$parsed_output" | sed -n 's/PARSED_TITLE://p')}"
        track="${track:-$(printf '%s' "$parsed_output" | sed -n 's/PARSED_TRACK://p')}"
        parsed_year="$(printf '%s' "$parsed_output" | sed -n 's/PARSED_YEAR://p')"
        
        [[ -n "$parsed_year" ]] && album="($parsed_year) $album"
        log "Parsed - Artist: '$artist', Album: '$album', Title: '$title', Track: '$track'"
    fi

    # Set defaults and sanitize
    album_artist="${album_artist:-${artist:-Unknown Artist}}"
    album="${album:-Unknown Album}"
    title="${title:-${filename%.*}}"
    track="$(pad_track "${track:-00}")"

    album_artist="$(sanitize "$album_artist")"
    album="$(sanitize "$album")"
    title="$(sanitize "$title")"

    # Create target directory and filename
    target_dir="$MUSIC_DIR/$album_artist/$album"
    if ! mkdir -p "$target_dir"; then
        log "ERROR: Failed to create directory: $target_dir" "ERROR"
        return 1
    fi

    target_path="$target_dir/$track - $title.$extension"

    # Handle filename conflicts more efficiently
    counter=1
    while [[ -f "$target_path" ]]; do
        target_path="$target_dir/$track - ${title}_$counter.$extension"
        ((counter++))
        # Prevent infinite loops
        if (( counter > 999 )); then
            log "ERROR: Too many filename conflicts for: $filename" "ERROR"
            return 1
        fi
    done

    log "Target: $album_artist/$album/$(basename "$target_path")"

    # Use rsync if available, fallback to cp/rm
    if [[ "$HAS_RSYNC" == "true" ]]; then
        if rsync -a --remove-source-files "$file" "$target_path"; then
            # Remove original if rsync left an empty file (happens sometimes with --remove-source-files)
            [[ -f "$file" ]] && rm -f "$file"
            log "SUCCESS: Moved to $target_path (via rsync)"
            return 0
        else
            log "ERROR: rsync failed for $file" "ERROR"
            return 1
        fi
    else
        if cp -p "$file" "$target_path" && rm -f "$file"; then
            log "SUCCESS: Moved to $target_path (via cp/rm)"
            return 0
        else
            log "ERROR: cp/rm failed for $file" "ERROR"
            return 1
        fi
    fi
}

# Optimized file discovery with better error handling
build_find_expression() {
    local -a find_args=()
    
    for ext in "${AUDIO_EXTENSIONS[@]}"; do
        find_args+=("-iname" "*.${ext}" "-o")
    done
    
    # Remove the last "-o"
    unset 'find_args[-1]'
    
    printf '%s ' "${find_args[@]}"
}

process_files_recursive() {
    log "Starting recursive music organization process"
    log "Searching for audio files in: $DOWNLOADS_DIR (recursively)"
    log "Target directory: $MUSIC_DIR"

    if [[ ! -d "$DOWNLOADS_DIR" ]]; then
        log "ERROR: Downloads directory does not exist: $DOWNLOADS_DIR" "ERROR"
        return 1
    fi

    local find_expression file_count=0
    find_expression="$(build_find_expression)"
    
    # Count files first for progress tracking
    log "Counting audio files..."
    if ! file_count=$(find "$DOWNLOADS_DIR" -type f \( "${find_expression}" \) 2>/dev/null | wc -l); then
        log "ERROR: Failed to search for files in $DOWNLOADS_DIR" "ERROR"
        return 1
    fi
    
    if (( file_count == 0 )); then
        log "No audio files found in $DOWNLOADS_DIR"
        return 0
    fi
    
    log "Found $file_count audio files to process"

    # Process files with progress tracking
    local current_file=0
    while IFS= read -r -d '' file; do
        ((current_file++))
        log "Progress: $current_file/$file_count"
        
        if move_file "$file"; then
            ((processed_count++))
        else
            ((failed_count++))
        fi
    done < <(find "$DOWNLOADS_DIR" -type f \( "${find_expression}" \) -print0 2>/dev/null)

    log "Process completed: $processed_count files processed, $failed_count failed, $skipped_count skipped"
    
    # Trigger Plex library refresh if files were processed
    if (( processed_count > 0 )) && [[ "$plex_refresh_override" == "true" ]]; then
        log "Files were processed - attempting Plex library refresh"
        refresh_plex_library
    fi
    
    # Return appropriate exit code
    if (( failed_count > 0 )); then
        return 1
    fi
    return 0
}

show_help() {
    cat << EOF
Usage: $SCRIPT_NAME [OPTIONS]

Recursively organizes music files from source directory to target directory
using metadata extraction and intelligent filename parsing.

OPTIONS:
    -h, --help          Show this help message and exit
    -v, --verbose       Enable verbose logging
    -n, --dry-run       Show what would be done without moving files
    -p, --plex-refresh  Enable Plex library refresh after processing
    
CONFIGURATION:
    Edit the script variables at the top to customize:
    
    DOWNLOADS_DIR      Source directory (default: $DOWNLOADS_DIR)
    MUSIC_DIR          Target directory (default: $MUSIC_DIR)  
    LOG_FILE           Log file path (default: $LOG_FILE)
    
    PLEX_URL               Plex server URL (default: $PLEX_URL)
    PLEX_TOKEN             Plex authentication token (required for refresh)
    PLEX_LIBRARY_SECTION   Plex library section ID or name (auto-detect if empty)

EXAMPLES:
    $SCRIPT_NAME                           # Use default configuration
    $SCRIPT_NAME --plex-refresh            # Enable Plex refresh for this run
    $SCRIPT_NAME --dry-run                 # Preview what would be done
    
PLEX CONFIGURATION:
    To enable automatic Plex library refresh:
    1. Edit script: Set PLEX_TOKEN="your_plex_token"
    2. Optionally edit PLEX_URL if not using default (localhost:32400)
    3. Optionally set PLEX_LIBRARY_SECTION to specific library ID/name
    
    Get your Plex token: https://support.plex.tv/articles/204059436/
    
METADATA EXTRACTION:
    - FLAC files: Uses metaflac first, falls back to ffprobe
    - Other formats: Uses ffprobe
    
SUPPORTED FORMATS:
    ${AUDIO_EXTENSIONS[*]}

EOF
}

main() {
    # Parse command line arguments
    local dry_run=false verbose=false
    
    while (( $# > 0 )); do
        case $1 in
            -h|--help)
                show_help
                exit 0
                ;;
            -v|--verbose)
                verbose=true
                shift
                ;;
            -n|--dry-run)
                dry_run=true
                log "DRY RUN MODE - No files will be moved" "WARN"
                shift
                ;;
            -p|--plex-refresh)
                plex_refresh_override="true"
                shift
                ;;
            *)
                log "Unknown option: $1" "ERROR"
                show_help >&2
                exit 1
                ;;
        esac
    done

    # Initialize
    check_tools
    mkdir -p "$DOWNLOADS_DIR" "$MUSIC_DIR"
    
    # Ensure log directory exists
    local log_dir
    log_dir="$(dirname "$LOG_FILE")"
    if [[ ! -d "$log_dir" ]]; then
        mkdir -p "$log_dir" || error_exit "Cannot create log directory: $log_dir"
    fi

    log "=== $SCRIPT_NAME started ==="
    log "Source: $DOWNLOADS_DIR"
    log "Target: $MUSIC_DIR"
    log "Log: $LOG_FILE"
    log "Plex refresh: $plex_refresh_override"

    # Validate Plex connection only if enabled
    if [[ "$plex_refresh_override" == "true" ]]; then
        validate_plex_connection
    fi

    # Main processing
    if process_files_recursive; then
        log "=== $SCRIPT_NAME completed successfully ==="
        exit 0
    else
        log "=== $SCRIPT_NAME completed with errors ===" "ERROR"
        exit 1
    fi
}

# Trap to ensure cleanup on exit
cleanup() {
    local exit_code=$?
    if (( exit_code != 0 )); then
        log "Script terminated with exit code: $exit_code" "ERROR"
    fi
}
trap cleanup EXIT

# Only run main if script is executed directly
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi