#!/bin/bash
# Author: Steven Shiau <steven _at_ clonezilla org>
# License: GPL
# Description: If this program is able to find the arch, it will print. Otherwise, nothing will be shown.

Usage() {
  echo "Usage:"
  echo "$(basename $0) [-c|--drbl-client] KERNEL_VER"
  echo "-c|--drbl-client is to check DRBL client's kernel arch."
  echo "If not specified, it will check DRBL server's kernel arch."
  echo "Ex: $0 2.4.20-31.9smp"
}

# Load DRBL setting and functions
DRBL_SCRIPT_PATH="${DRBL_SCRIPT_PATH:-/usr/share/drbl}"

. $DRBL_SCRIPT_PATH/sbin/drbl-conf-functions

# Parse command-line options
while [ $# -gt 0 ]; do
  case "$1" in
    -c|--drbl-client)
            shift; check_drbl_client_kernel="yes"
	    ;;
    -*)     echo "${0}: ${1}: invalid option" >&2
            Usage >& 2
            exit 2 ;;
    *)      break ;;
  esac
done
KERNEL=$1

[ -z "$KERNEL" ] && exit 1

if [ "$check_drbl_client_kernel" = "yes" ]; then
  CONFIG_FILE=$drbl_common_root/boot/config-$KERNEL
else
  CONFIG_FILE=/boot/config-$KERNEL
fi

if [ ! -f "$CONFIG_FILE" ]; then
  # Determine alternative paths for fallbacks
  if [ "$check_drbl_client_kernel" = "yes" ]; then
    KERNEL_FILE=$drbl_common_root/boot/vmlinuz-$KERNEL
    MODULES_DIR=$drbl_common_root/lib/modules/$KERNEL
    [ ! -d "$MODULES_DIR" ] && MODULES_DIR=$drbl_common_root/usr/lib/modules/$KERNEL
  else
    KERNEL_FILE=/boot/vmlinuz-$KERNEL
    MODULES_DIR=/lib/modules/$KERNEL
    [ ! -d "$MODULES_DIR" ] && MODULES_DIR=/usr/lib/modules/$KERNEL
  fi

  # Fallback 1: If requested kernel is the running kernel, use uname -m
  if [ "$KERNEL" = "$(uname -r)" ]; then
    case "$(uname -m)" in
      x86_64)
        echo "x86_64"
        exit 0
        ;;
      aarch64|arm64)
        echo "arm64"
        exit 0
        ;;
    esac
  fi

  # Fallback 2: Try to find file or directory to query package manager
  target=""
  if [ -f "$KERNEL_FILE" ]; then
    target="$KERNEL_FILE"
  elif [ -d "$MODULES_DIR" ]; then
    target="$MODULES_DIR"
  fi

  if [ -n "$target" ]; then
    if type dpkg &>/dev/null && type dpkg-query &>/dev/null; then
      pkg_name="$(dpkg -S "$target" 2>/dev/null | cut -d: -f1)"
      if [ -n "$pkg_name" ]; then
        pkg_arch="$(dpkg-query -W -f='${Architecture}\n' "$pkg_name" 2>/dev/null)"
        case "$pkg_arch" in
          amd64)
            echo "x86_64"
            exit 0
            ;;
          arm64)
            echo "arm64"
            exit 0
            ;;
        esac
      fi
    elif type rpm &>/dev/null; then
      pkg_arch="$(rpm -q --qf '%{ARCH}\n' -f "$target" 2>/dev/null)"
      case "$pkg_arch" in
        x86_64)
          echo "x86_64"
          exit 0
          ;;
        aarch64)
          echo "arm64"
          exit 0
          ;;
      esac
    fi
  fi

  # Fallback 3: Use file utility on the kernel file if available
  if [ -f "$KERNEL_FILE" ] && type file &>/dev/null; then
    file_info="$(file -b "$KERNEL_FILE" 2>/dev/null)"
    if echo "$file_info" | grep -qi -E "x86_64|x86-64|boot executable bzImage|Intel 80386"; then
      echo "x86_64"
      exit 0
    elif echo "$file_info" | grep -qi -E "arm64|aarch64"; then
      echo "arm64"
      exit 0
    fi
  fi

  # Fallback 4: Try checking the module file formats in the modules directory
  if [ -d "$MODULES_DIR" ] && type file &>/dev/null; then
    # find a kernel module .ko and check its file type
    ko_file="$(find "$MODULES_DIR" -name "*.ko*" -type f -print -quit 2>/dev/null)"
    if [ -n "$ko_file" ]; then
      file_info="$(file -b "$ko_file" 2>/dev/null)"
      if echo "$file_info" | grep -qi -E "x86-64|x86_64"; then
        echo "x86_64"
        exit 0
      elif echo "$file_info" | grep -qi -E "aarch64|arm64"; then
        echo "arm64"
        exit 0
      fi
    fi
  fi

  # If all fallbacks failed, exit with failure
  exit 1
fi

CPULIST="CONFIG_X86_64 CONFIG_ARM64"
for icpu in $CPULIST; do
  if grep -q "$icpu=y" $CONFIG_FILE; then
    arch="$(echo $icpu | sed -e "s/CONFIG_M//g" -e "s/CONFIG_//g" -e "s/TSC$//g" -e "s/MMX$//g" | tr "[A-Z]" "[a-z]")"
    case "$arch" in
     [3456]86) arch=i${arch} ;;
            *) arch=${arch} ;;
    esac
    echo "$arch"
  fi
done
