#!/bin/sh
# Part of raspi-config http://github.com/asb/raspi-config
#
# See LICENSE file for copyright and license details
# SPDX-FileCopyrightText: 2024 M5Stack Technology CO LTD
#
# SPDX-License-Identifier: MIT

INTERACTIVE=True
ASK_TO_REBOOT=0
CORE_CONFIG_VERSION="v0.1"
CORE_CONFIG_VERSION_ID=1

DESCRIPTION="\
This tool provides a straight-forward way of doing initial
configuration of the LLM. Although it can be run
at any time, some of the options may have difficulties if
you have heavily customised your installation.\
"

calc_wt_size() {
  # NOTE: it's tempting to redirect stderr to /dev/null, so supress error 
  # output from tput. However in this case, tput detects neither stdout or 
  # stderr is a tty and so only gives default 80, 24 values
  WT_HEIGHT=17
  WT_WIDTH=$(tput cols)

  if [ -z "$WT_WIDTH" ] || [ "$WT_WIDTH" -lt 60 ]; then
    WT_WIDTH=80
  fi
  if [ "$WT_WIDTH" -gt 178 ]; then
    WT_WIDTH=120
  fi
  WT_MENU_HEIGHT=$(($WT_HEIGHT-7))
}

do_about() {
  whiptail --yesno "$DESCRIPTION" 20 70 1 --yes-button Exit --no-button check_update 
  RET=$?
  if [ $RET -eq 1 ] ; then
    curl --fail https://raw.githubusercontent.com/m5stack/CoreMP135_buildroot-external-st/st/2023.02.10/tools/overlay_debian12/usr/bin/core-config -o /tmp/core-config
    if [ -f "/tmp/core-config" ] ; then
      chmod +x /tmp/core-config
      NEW_CORE_CONFIG_VERSION_ID=`sh /tmp/core-config --print-version-id`
      NEW_CORE_CONFIG_VERSION=`sh /tmp/core-config --print-version`
      NEW_DESCRIPTION=`sh /tmp/core-config --print-description`
      if [ "$NEW_CORE_CONFIG_VERSION_ID" -gt "$CORE_CONFIG_VERSION_ID" ] ; then
        whiptail --yesno "New version found:$NEW_CORE_CONFIG_VERSION 
$NEW_DESCRIPTION" \
        20 60 2 --yes-button update --no-button Exit
        RET=$?
        if [ $RET -eq 0 ] ; then
          cat /tmp/core-config > $0
          rm /tmp/core-config
          exec $0
        fi
      else
        whiptail --msgbox "The current version is already the latest." 20 60 2
      fi
    else
      whiptail --msgbox "Download core-config failed." 20 60 2
    fi
  fi
}

do_expand_rootfs() {
  if ! [ -f /usr/local/m5stack/resize_mmc.sh ]; then
    whiptail --msgbox "/usr/local/m5stack/resize_mmc.sh does not exist. Don't know how to expand" 20 60 2
    return 0
  fi

  whiptail --yesno "Start adjusting the file system size." 20 60 2 
  RET=$?
  if [ $RET -eq 1 ] ; then
    return 0
  fi

  clear ;

  /usr/local/m5stack/resize_mmc.sh

  whiptail --yesno "Root partition has been resized.\nThe filesystem will be enlarged upon the next reboot" 20 60 2 \
    --yes-button reboot --no-button Exit
  RET=$?
  if [ $RET -eq 0 ] ; then
    sh -c "sleep 1 ; reboot " &
    clear ;
    exit 0 ;
  fi
}

do_change_pass() {
  whiptail --msgbox "You will now be asked to enter a new password for the root user" 20 60 1
  passwd root &&
  whiptail --msgbox "Password changed successfully" 20 60 1
}

do_configure_keyboard() {
  dpkg-reconfigure keyboard-configuration &&
  printf "Reloading keymap. This may take a short while\n" &&
  invoke-rc.d keyboard-setup start
}

do_change_locale() {
  dpkg-reconfigure locales
}

do_change_timezone() {
  dpkg-reconfigure tzdata
}

do_change_hostname() {
  whiptail --msgbox "\
Please note: RFCs mandate that a hostname's labels \
may contain only the ASCII letters 'a' through 'z' (case-insensitive), 
the digits '0' through '9', and the hyphen.
Hostname labels cannot begin or end with a hyphen. 
No other symbols, punctuation characters, or blank spaces are permitted.\
" 20 70 1

  CURRENT_HOSTNAME=`cat /etc/hostname | tr -d " \t\n\r"`
  NEW_HOSTNAME=$(whiptail --inputbox "Please enter a hostname" 20 60 "$CURRENT_HOSTNAME" 3>&1 1>&2 2>&3)
  if [ $? -eq 0 ]; then
    echo $NEW_HOSTNAME > /etc/hostname
    sed -i "s/127.0.1.1.*$CURRENT_HOSTNAME/127.0.1.1\t$NEW_HOSTNAME/g" /etc/hosts
    ASK_TO_REBOOT=1
  fi
}

do_ssh() {
  whiptail --yesno "Would you like the SSH server enabled or disabled?" 20 60 2 \
    --yes-button Enable --no-button Disable
  RET=$?
  if [ $RET -eq 0 ]; then
    systemctl enable ssh &&
    systemctl start ssh &&
    whiptail --msgbox "SSH server enabled" 20 60 1
  elif [ $RET -eq 1 ]; then
    systemctl stop ssh &&
    systemctl disable ssh &&
    whiptail --msgbox "SSH server disabled" 20 60 1
  else
    return $RET
  fi
}


do_ssh_root() {
  CURRENT_STATUS="no" # assume hdmi output enabled
  if grep -q "^PermitRootLogin.*yes" /etc/ssh/sshd_config; then
    CURRENT_STATUS="yes"
  fi

  whiptail --yesno "Would you like the SSH server enabled or disabled?" 20 60 2 \
    --yes-button Enable --no-button Disable
  RET=$?
  if [ $RET -eq 0 ]; then
    if [ "$CURRENT_STATUS" = "no" ] ; then
      sed -i '$a PermitRootLogin yes' /etc/ssh/sshd_config
      systemctl restart ssh
    fi
    whiptail --msgbox "SSH server enabled" 20 60 1
  elif [ $RET -eq 1 ]; then
    if [ "$CURRENT_STATUS" = "yes" ] ; then
      sed -i /etc/ssh/sshd_config -e "/^PermitRootLogin.*yes/d"
      systemctl restart ssh
    fi
    whiptail --msgbox "SSH server disabled" 20 60 1
  else
    return $RET
  fi
}

do_serial() {
  CURRENT_STATUS="yes" # assume ttyAMA0 output enabled
  if ! grep -q "^T.*:.*:respawn:.*ttyAMA0" /etc/inittab; then
    CURRENT_STATUS="no"
  fi

  whiptail --yesno "Would you like a login shell to be accessible over serial?" 20 60 2
  RET=$?
  if [ $RET -eq 1 ]; then
    sed -i /etc/inittab -e "s|^.*:.*:respawn:.*ttyAMA0|#&|"
    sed -i /boot/cmdline.txt -e "s/console=ttyAMA0,[0-9]\+ //"
    whiptail --msgbox "Serial is now disabled" 20 60 1
  elif [ $RET -eq 0 ]; then
    sed -i /etc/inittab -e "s|^#\(.*:.*:respawn:.*ttyAMA0\)|\1|"
    if ! grep -q "^T.*:.*:respawn:.*ttyAMA0" /etc/inittab; then
      printf "T0:23:respawn:/sbin/getty -L ttyAMA0 115200 vt100\n" >> /etc/inittab
    fi
    if ! grep -q "console=ttyAMA0" /boot/cmdline.txt; then
      sed -i /boot/cmdline.txt -e "s/root=/console=ttyAMA0,115200 root=/"
    fi
    whiptail --msgbox "Serial is now enabled" 20 60 1
  else
    return $RET
  fi
}

# Temporary use, a unified dtbo solution will be established later.
temp_hdmi_off_dtbo_base64="0A3+7QAAAWEAAAA4AAABRAAAACgAAAARAAAAEAAAAAAAAAAdAAABDAAAAAAAAAAAAAAAAAAAAAAA
AAABAAAAAAAAAAFmcmFnbWVudEAwAAAAAAADAAAABAAAAAD/////AAAAAV9fb3ZlcmxheV9fAAAA
AAMAAAAJAAAAB2Rpc2FibGVkAAAAAAAAAAIAAAACAAAAAWZyYWdtZW50QDEAAAAAAAMAAAAEAAAA
AP////8AAAABX19vdmVybGF5X18AAAAAAwAAAAkAAAAHZGlzYWJsZWQAAAAAAAAAAgAAAAIAAAAB
X19maXh1cHNfXwAAAAAAAwAAABUAAAAOL2ZyYWdtZW50QDA6dGFyZ2V0OjAAAAAAAAAAAwAAABUA
AAATL2ZyYWdtZW50QDE6dGFyZ2V0OjAAAAAAAAAAAgAAAAIAAAAJdGFyZ2V0AHN0YXR1cwBsdGRj
AHBhbmVsX3JnYgA="

do_hdmi() {
  CURRENT_STATUS="yes" # assume hdmi output enabled
  if grep -q "^.*hdmi-off.*" /boot/extlinux/extlinux.conf; then
    CURRENT_STATUS="no"
  fi

  if [ ! -f "/boot/hdmi-off.dtbo" ] ; then
    printf "$temp_hdmi_off_dtbo_base64" | base64 -d > /boot/hdmi-off.dtbo
  fi

  whiptail --yesno "How do you operate HDMI?" 20 60 2 --yes-button Enable --no-button disabled
  RET=$?
  if [ $RET -eq 0 ]; then
    if [ "$CURRENT_STATUS" = "no" ]; then
        sed -i /boot/extlinux/extlinux.conf -e "/^.*hdmi-off.*/d"
    fi
    whiptail --msgbox "HDMI is now Enable" 20 60 1
  elif [ $RET -eq 1 ]; then
    if [ "$CURRENT_STATUS" = "yes" ]; then
        sed -i '$a \ \ fdtoverlays /boot/hdmi-off.dtbo' /boot/extlinux/extlinux.conf
    fi
    whiptail --msgbox "HDMI is now disabled" 20 60 1
  else
    return $RET
  fi
}

do_mem_size() {
  while true ; do
    ddr_mem_start=1073741824 # 0x40000000
    CMM_VALUE=$(grep 'insmod /soc/ko/ax_cmm.ko cmmpool=anonymous' "/soc/scripts/auto_load_all_drv.sh" | sed -n 's/.*cmmpool=anonymous,0,0x[0-9A-Fa-f]*,\([0-9]*\)M.*/\1/p')
    MEM_VALUE=$(grep -o 'mem=[0-9]*M' /proc/cmdline | sed 's/mem=\([0-9]*\)M/\1/')
    ALL_MEM=4096

    input=$(whiptail --title "Set Kernel MEM size" --inputbox "The size of MEM is between 512 and 2048.\n MEM size:" 10 60 $MEM_VALUE 3>&1 1>&2 2>&3)
    exitstatus=$?
    if [ $exitstatus = 0 ]; then
        if [ "$input" -lt 2049 ]; then
            NEW_CMM_VALUE=$(expr $ALL_MEM - $input)
            NEW_MEM_VALUE=$input
            tmp_val=$(expr $NEW_MEM_VALUE \* 1024 \* 1024)
            CMM_DDR_START=$(expr $ddr_mem_start + $tmp_val)

            ax_cmm_ko_cmd=$(printf 's|cmmpool=anonymous,0,0x[0-9A-Fa-f]\{7,8\},[0-9]\{1,4\}M|cmmpool=anonymous,0,0x%X,%dM|' $CMM_DDR_START $NEW_CMM_VALUE)
            sed -i "$ax_cmm_ko_cmd" /soc/scripts/auto_load_all_drv.sh
            
            cmdline=$(cat /proc/cmdline)
            new_cmdline=$(echo "$cmdline" | sed -E "s/mem=[0-9]+M/mem=${NEW_MEM_VALUE}M/")
            fw_setenv bootargs "$new_cmdline"
            whiptail --yesno "Root partition has been resized.\nThe filesystem will be enlarged upon the next reboot" 20 60 2 \
                --yes-button reboot --no-button Exit
            RET=$?
            if [ $RET -eq 0 ] ; then
                sh -c "sleep 1 ; reboot " &
                clear ;
            fi
            exit 0 ;
        else
            whiptail --msgbox "The CMM size minimum should Less than 2048." 20 60 1;
            continue
        fi
    fi
    exit 0
  done
}

do_esp32c6_wifi() {
  CURRENT_STATUS="yes" # assume hdmi output enabled
  OS_NAME=`grep "^NAME=" /etc/os-release | awk -F "=" '{print $2}' | tr -d '"'`
  if [ "$OS_NAME" = "Ubuntu" ] ; then
    if ! grep -q "^/usr/local/m5stack/bin/S03espc6wifi" /etc/rc.local; then
      CURRENT_STATUS="no"
    fi
    if [ ! -L '/etc/rc.local.d/S03espc6wifi' ]; then
      CURRENT_STATUS="no"
    fi
  fi
  if [ "$OS_NAME" = "Buildroot" ] ; then
    if [ ! -L '/etc/init.d/S03espc6wifi' ] ; then
      CURRENT_STATUS="no"
    fi
  fi

  whiptail --yesno "Do you want to enable WiFi?" 20 60 2 --yes-button Enable --no-button Disable
  RET=$?
  if [ $RET -eq 0 ]; then
    if [ "$OS_NAME" = "Ubuntu" ] ; then
      if [ "$CURRENT_STATUS" = "no" ]; then
          if [ -d '/etc/rc.local.d' ]; then
            ln -s /usr/local/m5stack/bin/S03espc6wifi /etc/rc.local.d/S03espc6wifi
          else
            sed -i '2a /usr/local/m5stack/bin/S03espc6wifi start' /etc/rc.local
          fi
          sync
          /usr/local/m5stack/bin/S03espc6wifi start
      fi
    fi
    if [ "$OS_NAME" = "Buildroot" ] ; then
      if [ "$CURRENT_STATUS" = "no" ]; then
          ln -s /usr/local/m5stack/bin/S03espc6wifi /etc/init.d/S03espc6wifi
          sync
          /usr/local/m5stack/bin/S03espc6wifi start
      fi
    fi
    whiptail --msgbox "WIFI is now Enable" 20 60 1
  elif [ $RET -eq 1 ]; then
    if [ "$OS_NAME" = "Ubuntu" ] ; then
      if [ "$CURRENT_STATUS" = "yes" ]; then
          if [ -d '/etc/rc.local.d' ]; then
            rm /etc/rc.local.d/S03espc6wifi
          else
            sed -i /etc/rc.local -e "/^.*S03espc6wifi.*/d"
          fi
          sync
          rmmod esp32_sdio
          /sbin/devmem 0x480100c 32 0x00
      fi
    fi
    if [ "$OS_NAME" = "Buildroot" ] ; then
      if [ "$CURRENT_STATUS" = "yes" ]; then
          rm -f /etc/init.d/S03espc6wifi
          sync
          rmmod esp32_sdio
          /sbin/devmem 0x480100c 32 0x00
      fi
    fi
    whiptail --msgbox "WIFI is now Disable" 20 60 1
  else
    return $RET
  fi
}

do_finish() {
  if [ $ASK_TO_REBOOT -eq 1 ]; then
    whiptail --yesno "Would you like to reboot now?" 20 60 2
    if [ $? -eq 0 ]; then # yes
      sync
      reboot
    fi
  fi
  exit 0
}

#
# Command line options for non-interactive use
#
for i in "$@"
do
  case $i in
  --apply-os-config)
    printf "Not currently supported\n"
    exit $?
    ;;
  --print-version-id)
    echo $CORE_CONFIG_VERSION_ID
    exit 0
    ;;
  --print-version)
    echo $CORE_CONFIG_VERSION
    exit 0
    ;;
  --print-description)
    echo $DESCRIPTION
    exit 0
    ;;
  -h|--help)
    echo "Usage: script.sh [OPTIONS]"
    echo ""
    echo "Options:"
    echo "  --set-lcd-backlight VALUE          Set the LCD backlight brightness to VALUE."
    echo "  --print-lcd-backlight              Print the current LCD backlight brightness."
    echo "  --expand-rootfs                    Expand the root filesystem. Please reboot after running this command."
    echo "  --apply-os-config                  Apply the OS configuration (not currently supported)."
    echo "  --print-version-id                 Print the CORE_CONFIG_VERSION_ID."
    echo "  --print-version                    Print the CORE_CONFIG_VERSION."
    echo "  --print-description                Print the DESCRIPTION."
    echo "  -h, --help                         Show this help message and exit."
    exit 0
    ;;
  *)
    # unknown option
    ;;
  esac
  shift
done

#if [ "GET" = "${OPT_MEMORY_SPLIT:-}" ]; then
#  set -u # Fail on unset variables
#  get_current_memory_split
#  echo $CURRENT_MEMSPLIT
#  exit 0
#fi

# Everything else needs to be run as root
if [ $(id -u) -ne 0 ]; then
  printf "Script must be run as root. Try 'sudo core-config'\n"
  exit 1
fi


do_internationalisation_menu() {
  FUN=$(whiptail --title "LLM Software Configuration Tool (core-config)" --menu "Internationalisation Options" $WT_HEIGHT $WT_WIDTH $WT_MENU_HEIGHT --cancel-button Back --ok-button Select \
    "I1 Change Locale" "Set up language and regional settings to match your location" \
    "I2 Change Timezone" "Set up timezone to match your location" \
    3>&1 1>&2 2>&3)
  RET=$?
  if [ $RET -eq 1 ]; then
    return 0
  elif [ $RET -eq 0 ]; then
    case "$FUN" in
      I1\ *) do_change_locale ;;
      I2\ *) do_change_timezone ;;
      *) whiptail --msgbox "Programmer error: unrecognized option" 20 60 1 ;;
    esac || whiptail --msgbox "There was an error running option $FUN" 20 60 1
  fi
}

do_advanced_menu() {
  FUN=$(whiptail --title "LLM Software Configuration Tool (core-config)" --menu "Advanced Options" $WT_HEIGHT $WT_WIDTH $WT_MENU_HEIGHT --cancel-button Back --ok-button Select \
    "A1 Hostname" "Set the visible name for this Pi on a network" \
    "A2 SSH" "Enable/Disable remote command line access to your device using SSH" \
    "A3 SSH for root" "Enable/Disable remote command line access to your root using SSH" \
    "A4 MEM" "Set Kernel MEM size, reboot to apply the configuration." \
    "A5 esp32c6-wifi" "Enable/Disable wifi." \
    3>&1 1>&2 2>&3)
  RET=$?
  if [ $RET -eq 1 ]; then
    return 0
  elif [ $RET -eq 0 ]; then
    case "$FUN" in
      A1\ *) do_change_hostname ;;
      A2\ *) do_ssh ;;
      A3\ *) do_ssh_root ;;
      A4\ *) do_mem_size ;;
      A5\ *) do_esp32c6_wifi ;;
      *) whiptail --msgbox "Programmer error: unrecognized option" 20 60 1 ;;
    esac || whiptail --msgbox "There was an error running option $FUN" 20 60 1
  fi
}


#
# Interactive use loop
#
calc_wt_size
while true; do
  FUN=$(whiptail --title "LLM Software Configuration Tool (core-config) $CORE_CONFIG_VERSION" --menu "Setup Options" $WT_HEIGHT $WT_WIDTH $WT_MENU_HEIGHT --cancel-button Finish --ok-button Select \
    "1 Change User Password" "Change password for the default user (root)" \
    "2 Internationalisation Options" "Set up language and regional settings to match your location" \
    "3 Advanced Options" "Configure advanced settings" \
    3>&1 1>&2 2>&3)
  RET=$?
  if [ $RET -eq 1 ]; then
    do_finish
  elif [ $RET -eq 0 ]; then
    case "$FUN" in
      1\ *) do_change_pass ;;
      2\ *) do_internationalisation_menu ;;
      3\ *) do_advanced_menu ;;
      *) whiptail --msgbox "Programmer error: unrecognized option" 20 60 1 ;;
    esac || whiptail --msgbox "There was an error running option $FUN" 20 60 1
  else
    exit 1
  fi
done
