#!/usr/bin/env bash
#
# setup-auto-updates.sh
#
# Universal Debian/Ubuntu unattended-upgrades setup for mixed environments:
#  - Robust repository detection (handles stable/codename/future releases)
#  - Works across 30+ VMs/LXCs with different configurations
#  - Future-proof for Trixie and beyond
#  - Fast and reliable detection algorithms
#  - Validates configuration before completion

set -Eeuo pipefail
trap 'echo "[ERROR] Line $LINENO failed" >&2' ERR

REBOOT_TIME="04:00"
TARGET_LOCALE="en_US.UTF-8"

# Repository detection results (global variables)
declare -g MAIN_ARCHIVE=""
declare -g UPDATES_ARCHIVE=""
declare -g SECURITY_ARCHIVE=""
declare -g DETECTED_ORIGINS=()

require_root() {
  if [[ $EUID -ne 0 ]]; then
    echo "[ERROR] Run as root (sudo)" >&2
    exit 1
  fi
}

detect_os() {
  if [[ ! -f /etc/os-release ]]; then
    echo "[ERROR] /etc/os-release not found" >&2
    exit 1
  fi
  . /etc/os-release
  case "$ID" in
    ubuntu) OS="ubuntu" ;;
    debian) OS="debian" ;;
    *) echo "[ERROR] Unsupported OS: $ID" >&2; exit 1 ;;
  esac
  CODENAME="${VERSION_CODENAME:-$UBUNTU_CODENAME}"
  VERSION_ID="${VERSION_ID:-unknown}"
  echo "[INFO] Detected: $PRETTY_NAME ($OS / $CODENAME)"
}

# Extract repository info from sources
parse_apt_sources() {
  local sources_data=""
  
  # Collect all active deb lines from sources.list and sources.list.d
  if [[ -f /etc/apt/sources.list ]]; then
    sources_data+=$(grep -E "^deb\s+" /etc/apt/sources.list 2>/dev/null || true)
    sources_data+=$'\n'
  fi
  
  if [[ -d /etc/apt/sources.list.d ]]; then
    sources_data+=$(find /etc/apt/sources.list.d -name "*.list" -exec grep -E "^deb\s+" {} \; 2>/dev/null || true)
  fi
  
  echo "$sources_data"
}

# Detect repositories for Ubuntu
detect_ubuntu_repos() {
  echo "[INFO] Detecting Ubuntu repository configuration..."
  
  local sources_data
  sources_data=$(parse_apt_sources)
  
  # Ubuntu typically uses consistent codename-based archives
  MAIN_ARCHIVE="$CODENAME"
  UPDATES_ARCHIVE="${CODENAME}-updates"
  SECURITY_ARCHIVE="${CODENAME}-security"
  
  # Verify these exist in sources
  if echo "$sources_data" | grep -q "ubuntu.*$CODENAME\s"; then
    echo "[INFO] Found Ubuntu $CODENAME repositories"
    
    # Check for ESM (Extended Security Maintenance) repos for LTS versions
    local esm_origins=()
    if echo "$sources_data" | grep -q "esm.ubuntu.com"; then
      esm_origins+=("Ubuntu ESM:${CODENAME}-infra-security")
      esm_origins+=("Ubuntu ESM Apps:${CODENAME}-apps-security")
      echo "[INFO] Found Ubuntu ESM repositories"
    fi
    
    # Build origins array
    DETECTED_ORIGINS=(
      "Ubuntu:${MAIN_ARCHIVE}"
      "Ubuntu:${UPDATES_ARCHIVE}" 
      "Ubuntu:${SECURITY_ARCHIVE}"
    )
    
    # Add ESM if present
    DETECTED_ORIGINS+=("${esm_origins[@]}")
    
  else
    echo "[WARN] Standard Ubuntu repositories not found, using defaults"
  fi
}

# Detect repositories for Debian with comprehensive fallback logic
detect_debian_repos() {
  echo "[INFO] Detecting Debian repository configuration..."
  
  # Update package lists for accurate detection
  apt-get update -qq >/dev/null 2>&1 || true
  
  local sources_data
  sources_data=$(parse_apt_sources)
  
  # Strategy 1: Parse sources.list for archive names
  local candidate_archives=()
  
  # Extract archives from main Debian repos (not security)
  while IFS= read -r line; do
    if [[ "$line" =~ deb[[:space:]]+[^[:space:]]+[[:space:]]+([^[:space:]]+) ]] && 
       [[ "$line" == *"debian"* ]] && 
       [[ "$line" != *"security"* ]]; then
      local archive="${BASH_REMATCH[1]}"
      if [[ "$archive" =~ ^(stable|testing|unstable|[a-z]+)$ ]]; then
        candidate_archives+=("$archive")
      fi
    fi
  done <<< "$sources_data"
  
  # Strategy 2: Check apt-cache policy for active origins
  local policy_archives=()
  if command -v apt-cache >/dev/null; then
    while IFS= read -r line; do
      if [[ "$line" =~ o=Debian,a=([^,]+) ]]; then
        local archive="${BASH_REMATCH[1]}"
        policy_archives+=("$archive")
      fi
    done < <(apt-cache policy 2>/dev/null | grep -E "^\s*[0-9]+\s+.*o=Debian,a=" || true)
  fi
  
  # Strategy 3: Test which archives actually have packages
  local working_archives=()
  local test_archives=("${candidate_archives[@]}" "${policy_archives[@]}" "stable" "$CODENAME")
  
  # Remove duplicates and test each
  local unique_archives=($(printf '%s\n' "${test_archives[@]}" | sort -u))
  
  for archive in "${unique_archives[@]}"; do
    if [[ -z "$archive" ]]; then continue; fi
    
    # Test if this archive has packages available
    if apt-cache policy 2>/dev/null | grep -q "o=Debian,a=$archive" || 
       echo "$sources_data" | grep -q "debian.*$archive"; then
      working_archives+=("$archive")
      echo "[DEBUG] Found working archive: $archive"
    fi
  done
  
  # Strategy 4: Choose the best archive
  if [[ ${#working_archives[@]} -gt 0 ]]; then
    # Prefer stable > codename > testing > others
    for preferred in "stable" "$CODENAME" "testing"; do
      for archive in "${working_archives[@]}"; do
        if [[ "$archive" == "$preferred" ]]; then
          MAIN_ARCHIVE="$archive"
          break 2
        fi
      done
    done
    
    # If no preferred found, use first working archive
    if [[ -z "$MAIN_ARCHIVE" ]]; then
      MAIN_ARCHIVE="${working_archives[0]}"
    fi
  else
    # Fallback to codename
    echo "[WARN] No working archives detected, using codename: $CODENAME"
    MAIN_ARCHIVE="$CODENAME"
  fi
  
  # Set related archives
  UPDATES_ARCHIVE="${MAIN_ARCHIVE}-updates"
  SECURITY_ARCHIVE="${MAIN_ARCHIVE}-security"
  
  # Special handling for suite names
  if [[ "$MAIN_ARCHIVE" == "testing" ]]; then
    SECURITY_ARCHIVE="testing-security"
  elif [[ "$MAIN_ARCHIVE" == "unstable" ]]; then
    UPDATES_ARCHIVE=""  # Unstable doesn't have updates
    SECURITY_ARCHIVE=""  # Unstable doesn't have security
  fi
  
  # Verify security archive exists
  if [[ -n "$SECURITY_ARCHIVE" ]]; then
    if ! (apt-cache policy 2>/dev/null | grep -q "o=Debian,a=$SECURITY_ARCHIVE" || 
          echo "$sources_data" | grep -q "security.*$SECURITY_ARCHIVE"); then
      echo "[WARN] Security archive $SECURITY_ARCHIVE not found"
      # Try alternative security archive naming
      if [[ "$MAIN_ARCHIVE" != "stable" ]] && 
         (apt-cache policy 2>/dev/null | grep -q "o=Debian,a=stable-security" || 
          echo "$sources_data" | grep -q "security.*stable-security"); then
        SECURITY_ARCHIVE="stable-security"
        echo "[INFO] Using stable-security instead"
      fi
    fi
  fi
  
  # Build origins array
  DETECTED_ORIGINS=("Debian:${MAIN_ARCHIVE}")
  
  if [[ -n "$UPDATES_ARCHIVE" ]]; then
    DETECTED_ORIGINS+=("Debian:${UPDATES_ARCHIVE}")
  fi
  
  if [[ -n "$SECURITY_ARCHIVE" ]]; then
    DETECTED_ORIGINS+=("Debian-Security:${SECURITY_ARCHIVE}")
  fi
  
  echo "[INFO] Debian archives selected:"
  echo "  Main: $MAIN_ARCHIVE"
  [[ -n "$UPDATES_ARCHIVE" ]] && echo "  Updates: $UPDATES_ARCHIVE"
  [[ -n "$SECURITY_ARCHIVE" ]] && echo "  Security: $SECURITY_ARCHIVE"
}

# Main repository detection function
detect_repositories() {
  echo "[INFO] Detecting repository configuration for $OS..."
  
  case "$OS" in
    ubuntu)
      detect_ubuntu_repos
      ;;
    debian)
      detect_debian_repos
      ;;
    *)
      echo "[ERROR] Unsupported OS for repository detection: $OS"
      exit 1
      ;;
  esac
  
  if [[ ${#DETECTED_ORIGINS[@]} -eq 0 ]]; then
    echo "[ERROR] No valid repository origins detected!"
    exit 1
  fi
  
  echo "[INFO] Final detected origins:"
  printf '  %s\n' "${DETECTED_ORIGINS[@]}"
}

fix_locale() {
  echo "[INFO] Ensuring locale ${TARGET_LOCALE} is generated…"
  apt-get update -y >/dev/null 2>&1 || true
  DEBIAN_FRONTEND=noninteractive apt-get install -y locales >/dev/null 2>&1 || true

  if [[ -f /etc/locale.gen ]]; then
    if grep -qE "^\s*#\s*${TARGET_LOCALE}\s+UTF-8" /etc/locale.gen; then
      sed -ri "s/^\s*#\s*(${TARGET_LOCALE}\s+UTF-8)/\1/" /etc/locale.gen
    elif ! grep -qE "^\s*${TARGET_LOCALE}\s+UTF-8" /etc/locale.gen; then
      echo "${TARGET_LOCALE} UTF-8" >> /etc/locale.gen
    fi
  else
    echo "${TARGET_LOCALE} UTF-8" > /etc/locale.gen
  fi

  locale-gen >/dev/null 2>&1
  update-locale LANG=${TARGET_LOCALE} LC_ALL=${TARGET_LOCALE}
  export LANG=${TARGET_LOCALE}
  export LC_ALL=${TARGET_LOCALE}
}

install_unattended() {
  echo "[INFO] Updating apt cache…"
  apt update
  echo "[INFO] Installing unattended-upgrades…"
  DEBIAN_FRONTEND=noninteractive apt install -y unattended-upgrades
}

enable_unattended() {
  echo "[INFO] Enabling unattended-upgrades via debconf…"
  echo 'unattended-upgrades unattended-upgrades/enable_auto_updates boolean true' | debconf-set-selections
  DEBIAN_FRONTEND=noninteractive dpkg-reconfigure -f noninteractive unattended-upgrades
}

# Generate origins configuration from detected origins
write_50_conf_universal() {
  local origins_config=""
  for origin in "${DETECTED_ORIGINS[@]}"; do
    origins_config+="    \"${origin}\";"$'\n'
  done
  
  cat > /etc/apt/apt.conf.d/50unattended-upgrades <<EOF
// ${OS^}: auto-install regular + security updates (kernels included)
// Detected origins: ${DETECTED_ORIGINS[*]}

Unattended-Upgrade::Allowed-Origins {
${origins_config}};

Unattended-Upgrade::Automatic-Reboot "true";
Unattended-Upgrade::Automatic-Reboot-Time "${REBOOT_TIME}";
Unattended-Upgrade::Automatic-Reboot-WithUsers "false";

Unattended-Upgrade::Remove-Unused-Dependencies "true";
Unattended-Upgrade::Remove-New-Unused-Dependencies "true";

Unattended-Upgrade::MinimalSteps "true";
Unattended-Upgrade::SyslogEnable "true";
Unattended-Upgrade::Verbose "0";
EOF
}

write_20_conf() {
  cat > /etc/apt/apt.conf.d/20auto-upgrades <<'EOF'
APT::Periodic::Update-Package-Lists "1";
APT::Periodic::Download-Upgradeable-Packages "1";
APT::Periodic::AutocleanInterval "7";
APT::Periodic::Unattended-Upgrade "1";
EOF
}

# Comprehensive validation with detailed error diagnosis
validate_config() {
  echo "[INFO] Validating unattended-upgrades configuration..."
  
  local temp_log="/tmp/unattended-upgrades-validation-$.log"
  
  # Test the configuration
  if ! timeout 60 unattended-upgrades --dry-run --debug > "$temp_log" 2>&1; then
    echo "[ERROR] Configuration validation failed (timeout or error)!"
    echo "[ERROR] Last 20 lines of debug output:"
    tail -n 20 "$temp_log"
    rm -f "$temp_log"
    return 1
  fi
  
  # Check for critical errors
  local critical_errors=(
    "not enough values to unpack"
    "Unable to parse.*Allowed-Origins"
    "ValueError:"
    "AttributeError:"
    "ImportError:"
  )
  
  for error_pattern in "${critical_errors[@]}"; do
    if grep -qi "$error_pattern" "$temp_log"; then
      echo "[ERROR] Critical configuration error detected:"
      grep -i "$error_pattern" "$temp_log" || true
      rm -f "$temp_log"
      return 1
    fi
  done
  
  # Analyze the results
  if grep -q "Allowed origins are:" "$temp_log"; then
    local allowed_origins
    allowed_origins=$(grep "Allowed origins are:" "$temp_log" | head -n1)
    echo "[INFO] $allowed_origins"
    
    # Convert our expected origins to the debug output format for comparison
    # "Debian:stable" becomes "o=Debian,a=stable"
    # "Debian-Security:stable-security" becomes "o=Debian-Security,a=stable-security"
    local configs_working=0
    for origin in "${DETECTED_ORIGINS[@]}"; do
      if [[ "$origin" =~ ^([^:]+):(.+)$ ]]; then
        local expected_format="o=${BASH_REMATCH[1]},a=${BASH_REMATCH[2]}"
        if [[ "$allowed_origins" == *"$expected_format"* ]]; then
          ((configs_working++))
          echo "[DEBUG] Found expected origin: $expected_format"
        fi
      fi
    done
    
    if [[ $configs_working -eq 0 ]]; then
      echo "[ERROR] None of the configured origins are being recognized!"
      echo "[ERROR] Expected origins (in config format): ${DETECTED_ORIGINS[*]}"
      echo "[ERROR] Debug log may show format differences - check /tmp/unattended-upgrades-test.log"
      rm -f "$temp_log"
      return 1
    elif [[ $configs_working -lt ${#DETECTED_ORIGINS[@]} ]]; then
      echo "[WARN] Only $configs_working/${#DETECTED_ORIGINS[@]} origins recognized, but continuing..."
    else
      echo "[OK] All $configs_working configured origins are being recognized"
    fi
  fi
  
  # Check for blocking issues
  if grep -q "No packages found that can be upgraded unattended" "$temp_log"; then
    if grep -q "Marking not allowed.*-32768" "$temp_log"; then
      # Count how many packages are being blocked
      local blocked_count
      blocked_count=$(grep -c "Marking not allowed.*-32768" "$temp_log" || echo "0")
      echo "[INFO] $blocked_count non-official packages are correctly being blocked"
      echo "[OK] Configuration is working - no official packages need updates"
    else
      echo "[OK] No packages need updating - configuration appears correct"
    fi
  else
    echo "[OK] Configuration validation passed"
  fi
  
  # Save full log for debugging
  cp "$temp_log" "/tmp/unattended-upgrades-test.log"
  rm -f "$temp_log"
  return 0
}

dry_run_test() {
  echo "[INFO] Testing unattended-upgrades configuration…"
  
  if validate_config; then
    echo "[OK] Configuration test passed"
    return 0
  else
    echo "[ERROR] Configuration test failed - see details above"
    echo "[INFO] Full debug log saved to /tmp/unattended-upgrades-test.log"
    echo "[INFO] You can retry with: sudo unattended-upgrades --dry-run --debug"
    return 1
  fi
}

show_status() {
  echo
  echo "[INFO] Service and timer status:"
  systemctl list-timers --all | grep -E 'apt-(daily|daily-upgrade)\.timer' || true
  echo
  echo "[INFO] Configuration files:"
  ls -la /etc/apt/apt.conf.d/{20auto-upgrades,50unattended-upgrades}
  echo
  echo "[INFO] Detected configuration summary:"
  echo "  OS: $OS $CODENAME"
  echo "  Origins: ${DETECTED_ORIGINS[*]}"
  if [[ "$OS" == "debian" ]]; then
    echo "  Main archive: $MAIN_ARCHIVE"
    [[ -n "$UPDATES_ARCHIVE" ]] && echo "  Updates: $UPDATES_ARCHIVE"
    [[ -n "$SECURITY_ARCHIVE" ]] && echo "  Security: $SECURITY_ARCHIVE"
  fi
}

prompt_self_delete() {
  echo
  read -r -p "Remove this script? [Y/n] " confirm
  case "$confirm" in
    [nN]*) echo "[INFO] Script not removed." ;;
    *) echo "[INFO] Removing script: $0"; rm -- "$0" ;;
  esac
}

main() {
  echo "[INFO] Universal Unattended-Upgrades Setup Script"
  echo "[INFO] Supports mixed Debian/Ubuntu environments with auto-detection"
  echo
  
  require_root
  detect_os
  detect_repositories
  fix_locale
  install_unattended
  enable_unattended

  echo "[INFO] Writing configuration files…"
  write_50_conf_universal
  write_20_conf

  if ! dry_run_test; then
    echo "[ERROR] Setup completed but configuration validation failed"
    echo "[ERROR] Manual review of /etc/apt/apt.conf.d/50unattended-upgrades may be needed"
    show_status
    exit 1
  fi
  
  show_status

  echo
  echo "[OK] ✅ Unattended updates configured successfully!"
  echo "[OK] Security + regular updates enabled, reboot @ ${REBOOT_TIME} if needed"
  echo "[OK] Compatible with current and future Debian/Ubuntu releases"
  
  prompt_self_delete
}

main "$@"