#!/bin/bash

set -euo pipefail

status_ok='[OK]'
status_warning='[!]'
if [[ -t 1 && ${TERM:-dumb} != dumb && -z ${NO_COLOR+x} ]]; then
    status_ok=$'\033[32m[OK]\033[0m'
    status_warning=$'\033[31m[!]\033[0m'
fi

_usage() {
    cat <<EOF
MySQL backup archive fixer (FM-10055)

Usage:
  ${0##*/} --archive-path BACKUP.tar [OPTIONS]

  To fix an archive without this tool, run following command against the mysql dump file extracted from the archive:

    sed -i 's/KEY \`idx_threat_code\` (\`threat_code\`)/UNIQUE KEY \`idx_threat_code\` (\`threat_code\`)/' mdm-databases.sql

Options:
  --archive-path PATH
      Backup tar archive to inspect or fix (required).
      Plain tar, Zstandard, and bzip2 are detected by content.
      Repairs preserve both outer archive and inner dump compression.
  --fix
      Change KEY idx_threat_code to UNIQUE KEY in database dumps.
      Without this option, check whether repair is needed without changes.
  --no-backup
      Do not retain the original archive as PATH.bak.
      Leave any existing PATH.bak untouched.
  --workdir DIRECTORY
      Use an existing scratch directory for temporary files.
      Default: a new temporary directory beside the archive.
      Warning: the directory and all its contents are removed on exit.
  -h, --help
      Print this help and exit.

Examples:

  Check if MDM backup archive needs repair:

    bash ${0##*/} --archive-path BACKUP.tar

  Repair MDM backup archive (optionally without backing up the original archive):

    bash ${0##*/} --archive-path BACKUP.tar --fix
    bash ${0##*/} --archive-path BACKUP.tar --fix --no-backup
EOF
}

_check_installed_compressors() {
    command -v tar >/dev/null || _fatal 'GNU tar is required'
    command -v zstd >/dev/null || _fatal 'zstd is required for compressed dumps'
    command -v bzip2 >/dev/null || _fatal 'bzip2 is required for compressed dumps'
}

_check_free_space() {
    local target=$1
    local required_space=$2
    local free_space

    required_space=$((required_space + (required_space + 9) / 10))
    free_space=$(LC_ALL=C df --output=avail --block-size=1 -- "$target" | awk 'NR == 2 { print $1 }')
    [[ "$free_space" =~ ^[0-9]+$ ]] || _fatal "Cannot determine available space in $target"
    printf 'Workdir space: required %s bytes (including 10%% margin), available %s bytes\n' \
        "$required_space" "$free_space"
    (( free_space >= required_space )) || _fatal "Insufficient disk space in $target"
}

_stream_archive() {
    "${archive_decoder[@]}" -- "$archive"
}

_stream_dump() {
    local member=$1
    tar "${archive_read_options[@]}" -xOf "$archive" -- "$member" | {
        local signature
        local -a decoder=(cat)
        head -c 4 > "$workdir/header"
        signature=$(od -An -tx1 -N4 "$workdir/header" | tr -d ' \n')
        case "$signature" in
            28b52ffd)
                command -v zstd >/dev/null || _fatal 'zstd is required for this dump'
                decoder=(zstd -q -dc)
                ;;
            425a683[1-9])
                command -v bzip2 >/dev/null || _fatal 'bzip2 is required for this dump'
                decoder=(bzip2 -dc)
                ;;
        esac
        { cat -- "$workdir/header"; cat; } | "${decoder[@]}"
    }
}

_fatal() {
    printf 'Error: %s\n' "$*" >&2
    exit 1
}

archive=''
fix=false
backup=true
workdir=''

while (( $# )); do
    case "$1" in
        --archive-path)
            (( $# >= 2 )) || _fatal '--archive-path requires a file path'
            archive=$2
            shift 2
            ;;
        --archive-path=*)
            archive=${1#*=}
            shift
            ;;
        --fix)
            fix=true
            shift
            ;;
        --no-backup)
            backup=false
            shift
            ;;
        --workdir)
            (( $# >= 2 )) || _fatal '--workdir requires a directory path'
            workdir=$2
            shift 2
            ;;
        --workdir=*)
            workdir=${1#*=}
            shift
            ;;
        --help|-h)
            _usage
            exit 0
            ;;
        *)
            _usage >&2
            echo "" >&2
            echo "" >&2
            _fatal "Unknown option: $1"
            exit 1
            ;;
    esac
done

_check_installed_compressors

[[ -n "$archive" ]] || { _usage >&2; exit 1; }
[[ -f "$archive" ]] || _fatal "Archive not found: $archive"
archive=$(realpath -- "$archive")

archive_compression=none
archive_decoder=(cat)
archive_read_options=()
archive_signature=$(od -An -tx1 -N4 "$archive" | tr -d ' \n')
case "$archive_signature" in
    28b52ffd)
        archive_compression=zstd
        archive_decoder=(zstd -q -dc)
        archive_read_options=(--use-compress-program=zstd)
        ;;
    425a683[1-9])
        archive_compression=bzip2
        archive_decoder=(bzip2 -dc)
        archive_read_options=(--use-compress-program=bzip2)
        ;;
esac

if [[ -z "$workdir" ]]; then
    workdir=$(mktemp -d -- "${archive%/*}/.backup-fixer.XXXXXX")
fi

trap 'rm -rf -- "$workdir"' EXIT
trap 'exit 130' INT
trap 'exit 143' TERM

tar "${archive_read_options[@]}" -tf "$archive" > "$workdir/members"

# LC_ALL=C to ensure consistent sorting and pattern matching regardless of the user's locale.
mapfile -t members < <(
    LC_ALL=C grep -E '^tmp/backup/mysql/databases[^/]*\.sql(\.(zst|bz2))?$' "$workdir/members" |
        LC_ALL=C sort
)
(( ${#members[@]} )) || _fatal 'No database SQL dump members found'

previous=''
for member in "${members[@]}"; do
    [[ "$member" != "$previous" ]] || _fatal "Duplicate dump member: $member"
    previous=$member
    printf 'Dump: %s\n' "$member"
done

echo "Script run time may vary depending on the size of the backup archive, please wait..."

key_pattern="^[[:space:]]*KEY \`idx_threat_code\` \(\`threat_code\`\)"
if [[ "$fix" == false ]]; then
    needs_fix=0
    for member in "${members[@]}"; do
        matching_keys=$(_stream_dump "$member" | sed -nE "/${key_pattern}/p" | wc -l) ||
            _fatal "Could not inspect dump: $member"
        if (( matching_keys > 0 )); then
            printf '%s Fix needed (%s matching key definitions): %s\n' "$status_warning" "$matching_keys" "$member"
            needs_fix=$((needs_fix + 1))
        else
            printf 'No fix needed: %s\n' "$member"
        fi
    done
    if (( needs_fix > 0 )); then
        echo ""
        printf '%s MDM backup elements need repair. Run the script again with --fix to repair them.\n' "$needs_fix"
        echo "Optionally, you can run the script with --fix --no-backup to repair without creating a backup of the original archive."
    else
        printf '%s No repair needed; archive left untouched.\n' "$status_ok"
    fi
    exit 0
fi

uncompressed_size=0
for member in "${members[@]}"; do
    printf 'Estimating unpacked size: %s\n' "$member"
    member_size=$(_stream_dump "$member" | wc -c)
    uncompressed_size=$((uncompressed_size + member_size))
done
if [[ "$archive_compression" == none ]]; then
    archive_size=$(stat --format=%s -- "$archive")
else
    printf 'Estimating unpacked outer archive size (%s).\n' "$archive_compression"
    archive_size=$(_stream_archive | wc -c)
fi
if [[ "$backup" == false ]]; then
    printf 'Backup disabled; peak space estimate is unchanged (backups use hardlinks).\n'
fi
printf 'Space estimate reserves two unpacked archive sizes and four unpacked dump sizes, plus 1 MiB.\n'
_check_free_space "$workdir" "$((2 * archive_size + 4 * uncompressed_size + 1048576))"

changed=0
mkdir -p -- "$workdir/tmp/backup/mysql"

for member in "${members[@]}"; do
    tar "${archive_read_options[@]}" -xOf "$archive" -- "$member" > "$workdir/stored"
    signature=$(od -An -tx1 -N4 "$workdir/stored" | tr -d ' \n')
    compression=none
    case "$signature" in
        28b52ffd)
            command -v zstd >/dev/null || _fatal 'zstd is required for this dump'
            compression=zstd
            zstd -q -dc -- "$workdir/stored" > "$workdir/dump.sql"
            ;;
        425a683[1-9])
            command -v bzip2 >/dev/null || _fatal 'bzip2 is required for this dump'
            compression=bzip2
            bzip2 -dc -- "$workdir/stored" > "$workdir/dump.sql"
            ;;
        *)
            mv -- "$workdir/stored" "$workdir/dump.sql"
            ;;
    esac

    if grep -qE "$key_pattern" "$workdir/dump.sql"; then
        sed -E "s/^([[:space:]]*)KEY \`idx_threat_code\` \(\`threat_code\`\)/\1UNIQUE KEY \`idx_threat_code\` (\`threat_code\`)/" \
            "$workdir/dump.sql" > "$workdir/fixed.sql"
    else
        status=$?
        # If grep did not find a match, it exits with status 1. Any other status indicates an error.
        (( status == 1 )) || _fatal "Could not inspect dump: $member"
        printf 'Unchanged (no matching non-unique key): %s\n' "$member"
        continue
    fi

    case "$compression" in
        zstd)
            zstd -q -c -- "$workdir/fixed.sql" > "$workdir/$member"
            ;;
        bzip2)
            bzip2 -c -- "$workdir/fixed.sql" > "$workdir/$member"
            ;;
        none)
            mv -- "$workdir/fixed.sql" "$workdir/$member"
            ;;
    esac

    if (( changed == 0 )); then
        if [[ "$backup" == true ]]; then
            [[ ! -e "$archive.bak" && ! -L "$archive.bak" ]] || _fatal "Backup already exists: $archive.bak"
        fi
        if [[ "$archive_compression" == none ]]; then
            cp --reflink=auto --preserve=mode,timestamps -- "$archive" "$workdir/archive.tar"
        else
            _stream_archive > "$workdir/archive.tar"
            chmod --reference="$archive" -- "$workdir/archive.tar"
        fi
    fi
    tar --delete -f "$workdir/archive.tar" -- "$member"
    tar -rf "$workdir/archive.tar" -C "$workdir" -- "$member"
    rm -f -- "$workdir/$member" "$workdir/dump.sql" "$workdir/fixed.sql" "$workdir/stored"
    changed=$((changed + 1))
    printf 'Fixed: %s\n' "$member"
done

if (( changed == 0 )); then
    printf 'No changes needed; archive left untouched.\n'
    exit 0
fi

replacement="$workdir/archive.tar"
case "$archive_compression" in
    zstd)
        zstd -q -c -- "$replacement" > "$workdir/archive.repacked"
        replacement="$workdir/archive.repacked"
        ;;
    bzip2)
        bzip2 -c -- "$replacement" > "$workdir/archive.repacked"
        replacement="$workdir/archive.repacked"
        ;;
esac
chmod --reference="$archive" -- "$replacement"

if [[ "$backup" == true ]]; then
    ln -- "$archive" "$archive.bak"
fi
mv -f -- "$replacement" "$archive"
printf 'Updated %s dump(s) in %s\n' "$changed" "$archive"
if [[ "$backup" == true ]]; then
    printf 'Backup: %s.bak\n' "$archive"
else
    printf 'No backup retained.\n'
fi