#!/bin/bash
set -euo pipefail

# -------- Defaults --------
DOCKER_COMPOSE_FILE="/opt/appdata/.docker-compose/docker-compose.yml"
CONTAINER_BACKUP_DIR="/data/postgres/backup"
RETENTION_COUNT=5
KUMA_BASE_URL="https://uptime.noxs.me/api/push"
TIMESTAMP=$(date '+%Y%m%d-%H%M')

# -------- CLI Parameters --------
DRY_RUN=false
LIST_ONLY=false
RESTORE=false
KUMA_PUSH_URL=""
DATABASES=()
DB_CONTAINER_NAME=""
DB_USER="user_pg"
DB_PASS=""
BACKUP_FILE=""
BACKUP_FILES=()
RESTORE_DB=""
RESTORE_DBS=()

# -------- Password Map --------
declare -A ROLE_PASSWORDS

# -------- Help --------
print_usage() {
  cat <<EOF

Usage:
  backup_db_container.sh --container=<name> --databases="db1 db2" [options]
  backup_db_container.sh --restore --container=<name> --backup-files="file1 file2" --restore-dbs="db1 db2" [options]

Required:
  --container=NAME          Docker Compose service name (e.g. pg, mysql, sqlite)
  --databases="DB1 DB2"     List of DBs or .db paths (space-separated) (for backup)
  --restore                 Restore mode (see below)
  --backup-files="FILES"    Space-separated list of backup files to restore
  --restore-dbs="NAMES"     (Recommended) Space-separated list of DB names or destination files, matched to --backup-files order

Optional:
  --backup-file=PATH        Path to a backup file to restore (single restore, DEPRECATED: use --backup-files)
  --restore-db=NAME         DB to restore into (single restore)
  --backup-dir=PATH         Path in container to save dumps (default: $CONTAINER_BACKUP_DIR)
  --compose-file=PATH       Path to docker-compose.yml (default: $DOCKER_COMPOSE_FILE)
  --retention=N             Number of backups to keep (default: $RETENTION_COUNT)
  --db-user=USER            DB user to connect as (default: $DB_USER)
  --db-pass=PASS            DB password (⚠️ optional, avoid for security if possible)
  --role-pass=USER:PASS     Set password for created/restored PostgreSQL roles (repeatable)
  --kuma-key=KEY            Uptime Kuma push key (converted to full push URL)
  --dry-run                 Show what would be done
  --list-only               Only list databases
  -h, --help                Show this help

Restore Example with role password:
  ./backup_db_container.sh --restore --container=pg \\
      --backup-files="/data/pg/sonarr-main-20250517-1500.backup /data/pg/sonarr-log-20250517-1500.backup" \\
      --restore-dbs="sonarr-main sonarr-log" \\
      --role-pass=sonarr_user:SuperSecret

EOF
}

# -------- Parse CLI --------
for arg in "$@"; do
  case $arg in
    --container=*) DB_CONTAINER_NAME="${arg#*=}" ;;
    --databases=*) IFS=' ' read -r -a DATABASES <<< "${arg#*=}" ;;
    --backup-dir=*) CONTAINER_BACKUP_DIR="${arg#*=}" ;;
    --compose-file=*) DOCKER_COMPOSE_FILE="${arg#*=}" ;;
    --retention=*) RETENTION_COUNT="${arg#*=}" ;;
    --db-user=*) DB_USER="${arg#*=}" ;;
    --db-pass=*) DB_PASS="${arg#*=}" ;;
    --role-pass=*)
      ROLE_VAL="${arg#*=}"
      ROLE_NAME="${ROLE_VAL%%:*}"
      ROLE_PW="${ROLE_VAL#*:}"
      ROLE_PASSWORDS["$ROLE_NAME"]="$ROLE_PW"
      ;;
    --kuma-key=*) KUMA_PUSH_URL="${KUMA_BASE_URL}/${arg#*=}" ;;
    --dry-run) DRY_RUN=true ;;
    --list-only) LIST_ONLY=true ;;
    --restore) RESTORE=true ;;
    --backup-file=*) BACKUP_FILE="${arg#*=}" ;;  # Legacy, single restore
    --backup-files=*) IFS=' ' read -r -a BACKUP_FILES <<< "${arg#*=}" ;;
    --restore-db=*) RESTORE_DB="${arg#*=}" ;;    # Legacy, single restore
    --restore-dbs=*) IFS=' ' read -r -a RESTORE_DBS <<< "${arg#*=}" ;;
    -h|--help) print_usage; exit 0 ;;
    *) echo "❌ Unknown option: $arg"; print_usage; exit 2 ;;
  esac
done

# -------- Validate Required --------
[[ -n "$DB_CONTAINER_NAME" ]] || { echo "❌ Missing --container"; print_usage; exit 1; }

if $RESTORE; then
  if [[ ${#BACKUP_FILES[@]} -eq 0 && -z "$BACKUP_FILE" ]]; then
    echo "❌ --backup-files is required with --restore"; print_usage; exit 1;
  fi
  # Legacy single-restore fallback:
  if [[ -n "$BACKUP_FILE" ]]; then BACKUP_FILES=("$BACKUP_FILE"); fi
else
  [[ ${#DATABASES[@]} -gt 0 ]] || { echo "❌ Missing --databases"; print_usage; exit 1; }
fi

# -------- Utilities --------
urlencode() { jq -sRr @uri <<<"$1"; }
notify_kuma() {
  [[ -n "$KUMA_PUSH_URL" ]] && \
    curl -fsS "${KUMA_PUSH_URL}?status=$1&msg=$(urlencode "$2")" --max-time 10 --retry 3 || true
}

# -------- Detect DB Type --------
output=$(docker compose -f "$DOCKER_COMPOSE_FILE" exec -T "$DB_CONTAINER_NAME" sh -c "which pg_restore || which pg_dump || which mysqldump || which sqlite3 || true")
if [[ "$output" == *pg_restore* || "$output" == *pg_dump* ]]; then
  DB_TYPE="postgresql"
  BACKUP_COMMAND="pg_dump"
  RESTORE_COMMAND="pg_restore"
  BACKUP_OPTIONS="-F c"
  FILE_EXTENSION=".backup"
elif [[ "$output" == *mysqldump* ]]; then
  DB_TYPE="mysql"
  BACKUP_COMMAND="mysqldump"
  RESTORE_COMMAND="mysql"
  BACKUP_OPTIONS="--single-transaction --quick"
  FILE_EXTENSION=".sql"
elif [[ "$output" == *sqlite3* ]]; then
  DB_TYPE="sqlite"
  BACKUP_COMMAND="sqlite3"
  FILE_EXTENSION=".db"
else
  echo "❌ No supported database tool found in container"
  notify_kuma down "Unsupported DB"
  exit 2
fi

# -------- Rotation --------
ROTATION_PATTERN="*.${FILE_EXTENSION#'.'}"
purge_old_backups() {
  echo "🧹 Rotating old backups (keep $RETENTION_COUNT)"
  if $DRY_RUN; then
    docker compose -f "$DOCKER_COMPOSE_FILE" exec -T "$DB_CONTAINER_NAME" \
      bash -c "cd $CONTAINER_BACKUP_DIR && ls -1t $ROTATION_PATTERN | tail -n +$((RETENTION_COUNT + 1))"
  else
    docker compose -f "$DOCKER_COMPOSE_FILE" exec -T "$DB_CONTAINER_NAME" \
      bash -c "cd $CONTAINER_BACKUP_DIR && ls -1t $ROTATION_PATTERN | tail -n +$((RETENTION_COUNT + 1)) | xargs -r rm -v"
  fi
}

# -------- Backup Logic --------
backup_database() {
  local db="$1"
  local base_name=$(basename "$db" .db)
  local file="$CONTAINER_BACKUP_DIR/${base_name//[^a-zA-Z0-9_-]/_}-${TIMESTAMP}${FILE_EXTENSION}"

  echo "📦 Backing up $db → $file"
  if $DRY_RUN || $LIST_ONLY; then
    echo "ℹ️ [Dry Run] Would backup $db to $file"
    return
  fi

  case "$DB_TYPE" in
    postgresql)
      if [[ -n "$DB_PASS" ]]; then
        docker compose -f "$DOCKER_COMPOSE_FILE" exec -T "$DB_CONTAINER_NAME" \
          env PGPASSWORD="$DB_PASS" \
          $BACKUP_COMMAND -U "$DB_USER" -h localhost -d "$db" $BACKUP_OPTIONS -f "$file"
      else
        docker compose -f "$DOCKER_COMPOSE_FILE" exec -T "$DB_CONTAINER_NAME" \
          $BACKUP_COMMAND -U "$DB_USER" -h localhost -d "$db" $BACKUP_OPTIONS -f "$file"
      fi
      ;;
    mysql)
      if [[ -n "$DB_PASS" ]]; then
        docker compose -f "$DOCKER_COMPOSE_FILE" exec -T "$DB_CONTAINER_NAME" \
          sh -c "MYSQL_PWD='$DB_PASS' $BACKUP_COMMAND -u$DB_USER $BACKUP_OPTIONS $db > $file"
      else
        docker compose -f "$DOCKER_COMPOSE_FILE" exec -T "$DB_CONTAINER_NAME" \
          sh -c "$BACKUP_COMMAND -u$DB_USER $BACKUP_OPTIONS $db > $file"
      fi
      ;;
    sqlite)
      docker compose -f "$DOCKER_COMPOSE_FILE" exec -T "$DB_CONTAINER_NAME" \
        sqlite3 "$db" ".backup '$file'"
      ;;
  esac || {
    echo "❌ Backup failed for $db"
    notify_kuma down "Backup failed: $db"
    exit 1
  }

  echo "✅ $db backup completed"
}

# -------- Restore Logic --------
restore_database() {
  local file="$1"
  local target_db="${2:-}"

  echo "⚠️  About to restore: $file"
  if $DRY_RUN; then
    echo "ℹ️ [Dry Run] Would restore $file (type: $DB_TYPE) to $target_db"
    return
  fi

  case "$DB_TYPE" in
    postgresql)
      [[ -n "$target_db" ]] || { echo "❌ --restore-db or --restore-dbs required for PostgreSQL"; exit 1; }
      echo "🧩 Restoring to Postgres DB: $target_db"

      # Try to create the DB if it doesn't exist
      if [[ -n "$DB_PASS" ]]; then
        docker compose -f "$DOCKER_COMPOSE_FILE" exec -T "$DB_CONTAINER_NAME" \
          env PGPASSWORD="$DB_PASS" psql -U "$DB_USER" -h localhost -tc "SELECT 1 FROM pg_database WHERE datname = '$target_db';" | grep -q 1 || \
          docker compose -f "$DOCKER_COMPOSE_FILE" exec -T "$DB_CONTAINER_NAME" \
            env PGPASSWORD="$DB_PASS" createdb -U "$DB_USER" -h localhost "$target_db"
      else
        docker compose -f "$DOCKER_COMPOSE_FILE" exec -T "$DB_CONTAINER_NAME" \
          psql -U "$DB_USER" -h localhost -tc "SELECT 1 FROM pg_database WHERE datname = '$target_db';" | grep -q 1 || \
          docker compose -f "$DOCKER_COMPOSE_FILE" exec -T "$DB_CONTAINER_NAME" \
            createdb -U "$DB_USER" -h localhost "$target_db"
      fi

      # ---- Auto-create missing roles from backup file (before restore) ----
      echo "🔍 Scanning backup for required roles (owners)..."
      # Extract last column of each TOC line, filter out numbers/public/semicolons/blanks, sort unique.
      ROLES=$(docker compose -f "$DOCKER_COMPOSE_FILE" exec -T "$DB_CONTAINER_NAME" \
        bash -c "pg_restore -l '$file' 2>/dev/null | awk '{if (NF>=1) print \$NF}' | grep -Ev '^;|^$' | grep -v '^[0-9]*$' | grep -v '^public$' | sort -u | tr '\n' ' '")
      echo "Roles found: '$ROLES'"

      for ROLE in $ROLES; do
        if [[ -n "$ROLE" ]]; then
          echo "🆕 Ensuring role exists: $ROLE"
          PW="${ROLE_PASSWORDS[$ROLE]:-}"
          if [[ -n "$PW" ]]; then
            CREATE_ROLE_SQL="DO \$\$ BEGIN IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '$ROLE') THEN CREATE ROLE \"$ROLE\" LOGIN PASSWORD '$PW'; ELSE ALTER ROLE \"$ROLE\" WITH PASSWORD '$PW'; END IF; END \$\$;"
          else
            CREATE_ROLE_SQL="DO \$\$ BEGIN IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '$ROLE') THEN CREATE ROLE \"$ROLE\" LOGIN; END IF; END \$\$;"
          fi

          if [[ -n "$DB_PASS" ]]; then
            docker compose -f "$DOCKER_COMPOSE_FILE" exec -T "$DB_CONTAINER_NAME" \
              env PGPASSWORD="$DB_PASS" psql -U "$DB_USER" -h localhost -c "$CREATE_ROLE_SQL"
          else
            docker compose -f "$DOCKER_COMPOSE_FILE" exec -T "$DB_CONTAINER_NAME" \
              psql -U "$DB_USER" -h localhost -c "$CREATE_ROLE_SQL"
          fi
        fi
      done
      # ---- End auto-create roles ----

      # Now restore
      if [[ -n "$DB_PASS" ]]; then
        docker compose -f "$DOCKER_COMPOSE_FILE" exec -T "$DB_CONTAINER_NAME" \
          env PGPASSWORD="$DB_PASS" \
          pg_restore -U "$DB_USER" -d "$target_db" -h localhost --clean --if-exists "$file"
      else
        docker compose -f "$DOCKER_COMPOSE_FILE" exec -T "$DB_CONTAINER_NAME" \
          pg_restore -U "$DB_USER" -d "$target_db" -h localhost --clean --if-exists "$file"
      fi
      ;;
    mysql)
      [[ -n "$target_db" ]] || { echo "❌ --restore-db or --restore-dbs required for MySQL"; exit 1; }
      echo "🧩 Restoring to MySQL DB: $target_db"
      if [[ -n "$DB_PASS" ]]; then
        docker compose -f "$DOCKER_COMPOSE_FILE" exec -T "$DB_CONTAINER_NAME" \
          sh -c "MYSQL_PWD='$DB_PASS' mysql -u$DB_USER $target_db < $file"
      else
        docker compose -f "$DOCKER_COMPOSE_FILE" exec -T "$DB_CONTAINER_NAME" \
          sh -c "mysql -u$DB_USER $target_db < $file"
      fi
      ;;
    sqlite)
      local sqlite_target="${target_db:-$file}"
      echo "🧩 Restoring SQLite DB: $sqlite_target"
      docker compose -f "$DOCKER_COMPOSE_FILE" exec -T "$DB_CONTAINER_NAME" \
        sh -c "cp '$file' '$sqlite_target'"
      ;;
  esac || {
    echo "❌ Restore failed for $file"
    notify_kuma down "Restore failed for $file"
    exit 1
  }
  echo "✅ Restore completed: $file"
  notify_kuma up "DB restore OK: $file"
}

# -------- Main --------
if $LIST_ONLY; then
  echo "📋 Databases to backup:"
  printf ' - %s\n' "${DATABASES[@]}"
  exit 0
fi

if $RESTORE; then
  count=${#BACKUP_FILES[@]}
  if [[ $count -eq 1 && -n "$RESTORE_DB" ]]; then
    restore_database "${BACKUP_FILES[0]}" "$RESTORE_DB"
    exit $?
  fi
  for i in "${!BACKUP_FILES[@]}"; do
    file="${BACKUP_FILES[$i]}"
    target_db=""
    if [[ ${#RESTORE_DBS[@]} -ge $((i+1)) ]]; then
      target_db="${RESTORE_DBS[$i]}"
    fi
    restore_database "$file" "$target_db"
  done
  exit $?
fi

purge_old_backups

for db in "${DATABASES[@]}"; do
  backup_database "$db"
done

echo "✅ All backups completed successfully"
notify_kuma up "All database backups OK"