#!/bin/bash

# Set the directory containing sound files, typically found in /data/sound in the games "contents"
SOUND_DIR=""

# Set alternative paths to external binaries (s3p_extract, 2dxBuild (apart of 2dxTools), ifstools, and ffmpeg)
# Otherwise, default to /usr/bin
S3P_EXTRACT_BIN="${S3P_EXTRACT_BIN:-s3p_extract}"
BUILD_2DX_BIN="${TWO_DX_BUILD_BIN:-2dxBuild}"
IFSTOOLS_BIN="${IFSTOOLS_BIN:-ifstools}"
FFMPEG_BIN="${FFMPEG_BIN:-ffmpeg}"

# Print basic usage if required args are missing
usage() {
    echo "Usage: $0 [-s </path/to/sound/dir>] [-e </path/to/s3p_extract>] [-b </path/to/2dxBuild>] [-i </path/to/ifstools>] [-f </path/to/ffmpeg>]"
    exit 1
}

# Options for sound directory and custom paths
while getopts ":s:e:b:i:f:" opt; do
    case $opt in
        s) SOUND_DIR="$OPTARG" ;;
        e) S3P_EXTRACT_BIN="$OPTARG" ;;
        b) BUILD_2DX_BIN="$OPTARG" ;;
        i) IFSTOOLS_BIN="$OPTARG" ;;
        f) FFMPEG_BIN="$OPTARG" ;;
        *) usage ;;
    esac
done

# Check if $SOUND_DIR is set
if [ -z "$SOUND_DIR" ]; then
    echo "Error: SOUND_DIR is not set. Use -s to specify the sound directory."
    usage
fi

# Extract .ifs file and descend into its directory
extract_ifs() {
    local ifs_file="$1"
    echo "Extracting $ifs_file..."

    # Ensure the path is correct by making it absolute if necessary
    if [[ "$ifs_file" != /* ]]; then
        ifs_file="$ifs_file"
    fi

    "IFSTOOLS_BIN" --super-disable "$ifs_file" || { echo "Failed to extract $ifs_file"; return 1; } # super-disable arg required for processing p0 .ifs without their parent .ifs
}

# ifstools outputs its files into a nested directory, with its parent dir being the song id, with the _ifs or -p0_ifs suffix
# This needs to be cleaned up before contents are moved into $SOUND_DIR, which is why this silly function exists
cleanup_ifs() {
    for dir in *_ifs *-p0_ifs; do
        if [ -d "$dir" ]; then
            if [[ "$dir" == *"-p0_ifs" ]]; then

                # If the directory ends with "-p0_ifs", strip the "_ifs" suffix but retain the "-p0" when moving to $SOUND_DIR
                # -p0.ifs are children to full .ifs files and appear to be chart updates (black another/leggendaria/etc?)
                # Need to exist as their own ifs file or directory

                base_name=$(basename "$dir" _ifs)
                inner_base_name="${base_name%-p0}"
                inner_dir="$dir/$inner_base_name"
                move_dir="$SOUND_DIR/$base_name"  # Keeps the -p0 when moving

            else
                # Otherwise, strip just _ifs for standard ifs files
                base_name=$(basename "$dir" _ifs)
                inner_dir="$dir/$base_name"
                move_dir="$SOUND_DIR/$base_name"  # Moves normally without -p0, since these are full song containers
            fi

            # Check if the inner directory exists
            if [ -d "$inner_dir" ]; then
                echo "Moving $inner_dir to $move_dir"

                # Move the inner directory to $SOUND_DIR with the correct naming
                mv "$inner_dir" "$move_dir"

                # Remove the now empty parent directory
                rmdir "$dir"
            else
                echo "Inner directory $inner_dir not found in $dir"
            fi
        fi
    done
}

# Extracts and converts .s3p files
extract_convert_s3p() {
    local s3p_file="$1"
    echo "Processing $s3p_file..."

    # Check if the directory contains .2dx files while ignoring preview files (_pre.2dx)
    if find . -maxdepth 1 -type f -name "*.2dx" ! -name "*_pre.2dx" | grep -q .; then
        echo "Directory contains .2dx (excluding preview files) - skipping."
        return 0
    fi

    # Proceed with extraction
    "S3P_EXTRACT_BIN" "$s3p_file" || { echo "Failed to extract $s3p_file"; return 1; }
    local s3p_out_dir="${s3p_file}.out"
    [ -d "$s3p_out_dir" ] && cd "$s3p_out_dir" || { echo "Directory $s3p_out_dir not found"; return 1; }

    # Convert .wma files to .wav
    for wma_file in *.wma; do
        [ -e "$wma_file" ] || continue
        echo -ne "Converting $wma_file..."\\r
        "FFMPEG_BIN" -loglevel error -i "$wma_file" -acodec adpcm_ms -b:a 256k "${wma_file%.*}.wav" && rm -f "$wma_file"
        # Remove loglevel argument if you want to see what ffmpeg is doing
    done
}

# Repacks .wav files into .2dx
repack_to_2dx() {
    echo "Repacking song/chart data to .2dx..."

    # Get the current working directory name
    dir_name=$(basename "$OLDPWD")

    # Remove the ".out" suffix from the directory name if it exists
    output_name="${dir_name%.out}"

    "BUILD_2DX_BIN" "$output_name".2dx || { echo "Failed to repack files"; return 1; }
    mv "$output_name".2dx ..
    cd ..
    rm -rf *s3p*
}

# Iterates through directories and .ifs in $SOUND_DIR
process_sound_files() {
    if [ ! -d "$SOUND_DIR" ]; then
        echo "Directory $SOUND_DIR does not exist."
        exit 1
    fi

    # Unpack.ifs files into directories, do not repack into .ifs - easier for the sake of managing future updates
    for ifs_file in "$SOUND_DIR"/*.ifs; do
        echo "Processing .ifs file: $ifs_file"
        (extract_ifs "$ifs_file")
        cleanup_ifs
    done

    # Convert all directories
    for dir in "$SOUND_DIR"/*/; do
        echo "Processing directory: $dir"
        (cd "$dir" && extract_convert_s3p *.s3p && repack_to_2dx)
    done
}

# Execute the main function
process_sound_files