#!/bin/bash

# Colors
green="\033[1;32m"
red="\033[1;31m"
yellow="\033[1;33m"
nc="\033[0m"

echo -e "${green}Proxmox Apt-Cacher-NG Proxy Updater${nc}"

# Select targets
echo -e "${yellow}Select target to update:${nc}"
echo "1) LXC Containers"
echo "2) VMs"
echo "3) All (LXC + VMs)"
read -rp "Enter choice [1-3]: " TARGET_CHOICE
case "$TARGET_CHOICE" in
  1) UPDATE_LXC=true; UPDATE_VM=false;;
  2) UPDATE_LXC=false; UPDATE_VM=true;;
  3) UPDATE_LXC=true; UPDATE_VM=true;;
  *) echo -e "${red}Invalid choice. Exiting.${nc}"; exit 1;;
esac

# Optionally skip specific IDs
echo
echo -e "${yellow}If you wish to skip specific IDs, enter them comma-separated (e.g. 100,102), or press Enter to skip skip:${nc}"
read -rp "Skip IDs: " SKIP_IDS_RAW
# Convert to array for checking
IFS=',' read -ra SKIP_IDS <<< "$SKIP_IDS_RAW"

# Ask to update the Proxmox node itself
read -rp "Update Proxmox host node config? [y/N]: " UPDATE_NODE_REPLY
if [[ "$UPDATE_NODE_REPLY" =~ ^[Yy]$ ]]; then
  UPDATE_NODE=true
else
  UPDATE_NODE=false
fi

# Read apt-cacher-ng details
read -rp "Enter apt-cacher-ng IP [192.168.1.10]: " CACHE_IP
CACHE_IP="${CACHE_IP:-192.168.1.10}"
read -rp "Enter apt-cacher-ng port [3142]: " CACHE_PORT
CACHE_PORT="${CACHE_PORT:-3142}"

PROXY_CONF="Acquire::http::Proxy \"http://${CACHE_IP}:${CACHE_PORT}\";"
CONF_FILE="00aptproxy"

echo -e "\n${yellow}Using proxy: http://${CACHE_IP}:${CACHE_PORT}${nc}\n"
read -rp "Proceed? [y/N]: " CONFIRM
[[ ! "$CONFIRM" =~ ^[Yy]$ ]] && echo "Aborted." && exit 1

# Function to update a target (container or VM)
update_target() {
  local exec_cmd="$1"
  local id="$2"
  echo -e "\n${yellow}Processing ID: $id${nc}"
  # Check skip
  for skip in "${SKIP_IDS[@]}"; do
    [[ "$id" == "$skip" ]] && echo -e "${red}Skipping ID $id per user request.${nc}" && return
  done
  # Check distro compatibility
  $exec_cmd "bash -c 'if command -v lsb_release &>/dev/null; then DIST=\$(lsb_release -is); else DIST=\$(. /etc/os-release && echo \$ID); fi; if [[ ! \$DIST =~ (ubuntu|debian) ]]; then exit 2; fi'"
  case $? in
    2) echo -e "${red}ID $id is not Debian/Ubuntu based. Skipping.${nc}"; return;;
  esac
  # Remove old proxy files
  $exec_cmd "rm -f /etc/apt/apt.conf.d/*proxy*"
  # If config exists, prompt before overwrite
  $exec_cmd "bash -c 'if [[ -f /etc/apt/apt.conf.d/$CONF_FILE ]]; then exit 1; fi'"
  if [[ $? -eq 1 ]]; then
    read -rp "Config exists in ID $id. Overwrite? [y/N]: " over
    [[ ! "$over" =~ ^[Yy]$ ]] && echo -e "${red}Skipped ID $id per user choice.${nc}" && return
  fi
  # Write new proxy
  $exec_cmd "bash -c 'mkdir -p /etc/apt/apt.conf.d && echo \"$PROXY_CONF\" > /etc/apt/apt.conf.d/$CONF_FILE'"
  echo -e "${green}ID $id updated successfully.${nc}"
}

# Update LXC Containers
if $UPDATE_LXC; then
  echo -e "\n${green}Updating LXC Containers...${nc}"
  for ct in $(pct list | awk 'NR>1 {print $1}'); do
    update_target "pct exec $ct --" "$ct"
  done
fi

# Update VMs via QEMU agent
if $UPDATE_VM; then
  echo -e "\n${green}Updating VMs...${nc}"
  for vm in $(qm list | awk 'NR>1 {print $1}'); do
    # Only process if guest agent available
    qm guest cmd $vm get-osinfo &>/dev/null || { echo -e "${red}VM $vm has no guest agent. Skipping.${nc}"; continue; }
    update_target "qm guest exec $vm --" "$vm"
  done
fi

# Update host node
if $UPDATE_NODE; then
  echo -e "\n${green}Updating Proxmox host node...${nc}"
  rm -f /etc/apt/apt.conf.d/*proxy*
  echo "$PROXY_CONF" > "/etc/apt/apt.conf.d/$CONF_FILE"
  echo -e "${green}Host node updated successfully.${nc}"
fi

echo -e "\n${green}All selected updates complete.${nc}"