#!/bin/bash

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

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

# Choose operation
echo -e "${yellow}Select operation:${nc}"
echo "1) Apply APT proxy"
echo "2) Remove APT proxy"
read -rp "Enter choice [1-2]: " OP
case "$OP" in
  1) ACTION="apply";;
  2) ACTION="remove";;
  *) echo -e "${red}Invalid choice. Exiting.${nc}"; exit 1;;
esac

# Select targets
echo -e "\n${yellow}Select target to process:${nc}"
echo "1) LXC containers"
echo "2) VMs"
echo "3) Both LXC + VMs"
echo "4) Proxmox host only"
echo "5) All (LXC + VMs + Host)"
read -rp "Enter choice [1-5]: " T
case "$T" in
  1) DO_LXC=true; DO_VM=false; DO_HOST=false;;
  2) DO_LXC=false; DO_VM=true;  DO_HOST=false;;
  3) DO_LXC=true;  DO_VM=true;  DO_HOST=false;;
  4) DO_LXC=false; DO_VM=false; DO_HOST=true;;
  5) DO_LXC=true;  DO_VM=true;  DO_HOST=true;;
  *) echo -e "${red}Invalid choice. Exiting.${nc}"; exit 1;;
esac

# Skip specific IDs (for LXC/VM)
read -rp "\nEnter IDs to skip (comma-separated), or press Enter: " SKIP_RAW
IFS=',' read -ra SKIP_IDS <<< "$SKIP_RAW"

# If applying, ask for proxy details
if [ "$ACTION" = "apply" ]; then
  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="Acquire::http::Proxy \"http://${CACHE_IP}:${CACHE_PORT}\";"
  echo -e "\n${yellow}Configured proxy: ${PROXY}${nc}\n"
fi

confirm() {
  read -rp "$1 [y/N]: " y && [[ "$y" =~ ^[Yy]$ ]]
}

process_lxc() {
  local id=$1
  # skip?
  for s in "${SKIP_IDS[@]}"; do [[ "$id" = "$s" ]] && echo "- Skipping LXC $id" && return; done
  echo -e "\nProcessing LXC $id"
  # ensure running
  if ! pct status "$id" | grep -q "running"; then
    echo "- Not running, skipping"
    return
  fi
  if [ "$ACTION" = "apply" ]; then
    # check distro
    pct exec "$id" -- sh -c 'export ID=$(awk -F= "/^ID=/ {print \$2}" /etc/os-release) && case "$ID" in debian|ubuntu) exit 0;; *) exit 1;; esac'
    if [ $? -ne 0 ]; then echo "- Not Debian/Ubuntu, skipping"; return; fi
    # existing?
    if pct exec "$id" -- test -f /etc/apt/apt.conf.d/00aptproxy; then
      confirm "Proxy exists in LXC $id. Overwrite?" || { echo "- Skipped"; return; }
    fi
    pct exec "$id" -- sh -c 'rm -f /etc/apt/apt.conf.d/*proxy*'
    pct exec "$id" -- sh -c "mkdir -p /etc/apt/apt.conf.d && echo '$PROXY' > /etc/apt/apt.conf.d/00aptproxy"
    echo "- Proxy applied to LXC $id"
  else
    # remove
    if pct exec "$id" -- test -f /etc/apt/apt.conf.d/00aptproxy; then
      pct exec "$id" -- sh -c 'rm -f /etc/apt/apt.conf.d/*proxy*'
      echo "- Proxy removed from LXC $id"
    else
      echo "- No proxy config in LXC $id"
    fi
  fi
}

process_vm() {
  local id=$1
  for s in "${SKIP_IDS[@]}"; do [[ "$id" = "$s" ]] && echo "- Skipping VM $id" && return; done
  echo -e "\nProcessing VM $id"
  if ! qm guest cmd "$id" ping &>/dev/null; then
    echo "- Guest agent not available, skipping"
    return
  fi
  if [ "$ACTION" = "apply" ]; then
    qm guest exec "$id" -- sh -c 'export ID=$(awk -F= "/^ID=/ {print \$2}" /etc/os-release) && case "$ID" in debian|ubuntu) exit 0;; *) exit 1;; esac'
    if [ $? -ne 0 ]; then echo "- Not Debian/Ubuntu, skipping"; return; fi
    if qm guest exec "$id" -- test -f /etc/apt/apt.conf.d/00aptproxy; then
      confirm "Proxy exists in VM $id. Overwrite?" || { echo "- Skipped"; return; }
    fi
    qm guest exec "$id" -- sh -c 'rm -f /etc/apt/apt.conf.d/*proxy*'
    qm guest exec "$id" -- sh -c "mkdir -p /etc/apt/apt.conf.d && echo '$PROXY' > /etc/apt/apt.conf.d/00aptproxy"
    echo "- Proxy applied to VM $id"
  else
    if qm guest exec "$id" -- test -f /etc/apt/apt.conf.d/00aptproxy; then
      qm guest exec "$id" -- sh -c 'rm -f /etc/apt/apt.conf.d/*proxy*'
      echo "- Proxy removed from VM $id"
    else
      echo "- No proxy config in VM $id"
    fi
  fi
}

# Process items
if [ "$DO_LXC" = true ]; then
  echo -e "\n${green}-- LXC Containers --${nc}"
  for id in $(pct list | awk 'NR>1{print $1}'); do process_lxc "$id"; done
fi

if [ "$DO_VM" = true ]; then
  echo -e "\n${green}-- VMs --${nc}"
  for id in $(qm list | awk 'NR>1{print $1}'); do process_vm "$id"; done
fi

if [ "$DO_HOST" = true ]; then
  echo -e "\n${green}-- Proxmox Host --${nc}"
  if [ "$ACTION" = "apply" ]; then
    if [ -f /etc/apt/apt.conf.d/00aptproxy ]; then
      confirm "Proxy exists on host. Overwrite?" || { echo "- Skipped host"; exit; }
    fi
    rm -f /etc/apt/apt.conf.d/*proxy*
    echo "$PROXY" > /etc/apt/apt.conf.d/00aptproxy
    echo "- Proxy applied on host"
  else
    if [ -f /etc/apt/apt.conf.d/00aptproxy ]; then
      rm -f /etc/apt/apt.conf.d/*proxy*
      echo "- Proxy removed from host"
    else
      echo "- No proxy config on host"
    fi
  fi
fi

echo -e "\n${green}Operation complete.${nc}"