#!/usr/bin/env bash
#==============================================================================
# APT Source Switcher - Ubuntu -> TUNA (Tsinghua University) Mirror
#
# Author:      Smith Wang
# Version:     1.0.0
# License:     MIT
#==============================================================================
# Module Dependencies:
#   - bash (>= 5.0)
#   - coreutils: cp, mv, mkdir, date, id, chmod, chown
#   - util-linux / distro tools: (optional) lsb_release
#   - text tools: sed, awk, grep
#   - apt: apt-get
#
# Notes:
#   - Ubuntu 24.04 typically uses Deb822 sources file: /etc/apt/sources.list.d/ubuntu.sources
#   - Ubuntu 20.04/22.04 often uses traditional /etc/apt/sources.list
#==============================================================================

set -euo pipefail
IFS=$'\n\t'
umask 022

#------------------------------------------------------------------------------
# Global Constants
#------------------------------------------------------------------------------
readonly SCRIPT_NAME="$(basename "$0")"
readonly SCRIPT_VERSION="1.0.0"
readonly TUNA_UBUNTU_URI="https://mirrors.tuna.tsinghua.edu.cn/ubuntu/"
readonly DEFAULT_BACKUP_DIR="/etc/apt/backup"
readonly APT_SOURCES_LIST="/etc/apt/sources.list"
readonly APT_DEB822_SOURCES="/etc/apt/sources.list.d/ubuntu.sources"

# Log levels: DEBUG=0, INFO=1, WARN=2, ERROR=3
readonly LOG_LEVEL_DEBUG=0
readonly LOG_LEVEL_INFO=1
readonly LOG_LEVEL_WARN=2
readonly LOG_LEVEL_ERROR=3

#------------------------------------------------------------------------------
# Runtime Variables (defaults)
#------------------------------------------------------------------------------
log_level="$LOG_LEVEL_INFO"
backup_dir="$DEFAULT_BACKUP_DIR"
do_update="false"
assume_yes="false"

ubuntu_codename=""
ubuntu_version_id=""
sources_mode=""   # "deb822" or "list"

#------------------------------------------------------------------------------
# Function Call Graph (ASCII)
#
#   main
#    |
#    +--> parse_args
#    +--> setup_traps
#    +--> require_root
#    +--> detect_ubuntu
#    |      |
#    |      +--> read_os_release
#    |
#    +--> choose_sources_mode
#    +--> ensure_backup_dir
#    +--> backup_sources
#    +--> confirm_action
#    +--> apply_tuna_sources
#    |      |
#    |      +--> write_sources_list_tuna
#    |      +--> patch_deb822_sources_tuna
#    |
#    +--> apt_update (optional)
#    +--> summary
#------------------------------------------------------------------------------

#------------------------------------------------------------------------------
# Logging
#------------------------------------------------------------------------------

### Log message with level.
# @param level int Numeric log level (0=DEBUG,1=INFO,2=WARN,3=ERROR)
# @param message string Message to print
# @return 0 success
# @require date printf
log() {
  local level="$1"
  local message="$2"

  if [[ "$level" -lt "$log_level" ]]; then
    return 0
  fi

  local level_name="INFO"
  case "$level" in
    0) level_name="DEBUG" ;;
    1) level_name="INFO"  ;;
    2) level_name="WARN"  ;;
    3) level_name="ERROR" ;;
    *) level_name="INFO"  ;;
  esac

  # > Use RFC3339-ish timestamp to help operations & auditing
  printf '%s [%s] %s: %s\n' "$(date '+%Y-%m-%d %H:%M:%S')" "$level_name" "$SCRIPT_NAME" "$message" >&2
}

### Convenience wrappers.
# @return 0
log_debug() { log "$LOG_LEVEL_DEBUG" "$1"; }
log_info()  { log "$LOG_LEVEL_INFO"  "$1"; }
log_warn()  { log "$LOG_LEVEL_WARN"  "$1"; }
log_error() { log "$LOG_LEVEL_ERROR" "$1"; }

#------------------------------------------------------------------------------
# Error handling / traps
#------------------------------------------------------------------------------

### Trap handler for unexpected errors.
# @param exit_code int Exit code
# @param line_no int Line number where error occurred
# @param cmd string The command that failed
# @return 0
# @require printf
on_err() {
  local exit_code="$1"
  local line_no="$2"
  local cmd="$3"
  log_error "Script failed (exit=${exit_code}) at line ${line_no}: ${cmd}"
}

### Cleanup handler (reserved for future extension).
# @return 0
# @require true
on_exit() {
  true
}

### Setup traps for ERR and EXIT.
# @return 0
# @require trap
setup_traps() {
  # > Preserve error context with BASH_LINENO and BASH_COMMAND
  trap 'on_err "$?" "${LINENO}" "${BASH_COMMAND}"' ERR
  trap 'on_exit' EXIT
}

#------------------------------------------------------------------------------
# Utility / validation
#------------------------------------------------------------------------------

### Print usage.
# @return 0
# @require cat
usage() {
  cat <<'EOF'
Usage:
  sudo ./apt_tuna_switch.sh [options]

Options:
  -y, --yes           Non-interactive; do not prompt.
  -u, --update        Run "apt-get update" after switching.
  -b, --backup-dir    Backup directory (default: /etc/apt/backup)
  -d, --debug         Enable DEBUG logs.
  -h, --help          Show help.

Examples:
  sudo ./apt_tuna_switch.sh -y -u
  sudo ./apt_tuna_switch.sh --backup-dir /root/apt-bak --update
EOF
}

### Parse CLI arguments.
# @param args string[] CLI args
# @return 0 success
# @require printf
parse_args() {
  while [[ $# -gt 0 ]]; do
    case "$1" in
      -y|--yes)
        assume_yes="true"
        shift
        ;;
      -u|--update)
        do_update="true"
        shift
        ;;
      -b|--backup-dir)
        if [[ $# -lt 2 ]]; then
          log_error "Missing value for --backup-dir"
          usage
          exit 2
        fi
        backup_dir="$2"
        shift 2
        ;;
      -d|--debug)
        log_level="$LOG_LEVEL_DEBUG"
        shift
        ;;
      -h|--help)
        usage
        exit 0
        ;;
      *)
        log_error "Unknown argument: $1"
        usage
        exit 2
        ;;
    esac
  done
}

### Ensure running as root.
# @return 0 if root; exit otherwise
# @require id
require_root() {
  if [[ "$(id -u)" -ne 0 ]]; then
    log_error "This script must be run as root. Try: sudo ./${SCRIPT_NAME}"
    exit 1
  fi
}

### Read /etc/os-release fields.
# @return 0
# @require awk grep
read_os_release() {
  if [[ ! -r /etc/os-release ]]; then
    log_error "Cannot read /etc/os-release"
    exit 1
  fi

  # > Parse key fields safely
  local os_id
  os_id="$(awk -F= '$1=="ID"{gsub(/"/,"",$2); print $2}' /etc/os-release | head -n1 || true)"
  ubuntu_version_id="$(awk -F= '$1=="VERSION_ID"{gsub(/"/,"",$2); print $2}' /etc/os-release | head -n1 || true)"
  ubuntu_codename="$(awk -F= '$1=="VERSION_CODENAME"{gsub(/"/,"",$2); print $2}' /etc/os-release | head -n1 || true)"

  if [[ "$os_id" != "ubuntu" ]]; then
    log_error "Unsupported OS ID: ${os_id:-unknown}. This script supports Ubuntu only."
    exit 1
  fi

  if [[ -z "$ubuntu_version_id" ]]; then
    log_error "Failed to detect Ubuntu VERSION_ID from /etc/os-release"
    exit 1
  fi

  # > For some environments, VERSION_CODENAME may be empty; try UBUNTU_CODENAME
  if [[ -z "$ubuntu_codename" ]]; then
    ubuntu_codename="$(awk -F= '$1=="UBUNTU_CODENAME"{gsub(/"/,"",$2); print $2}' /etc/os-release | head -n1 || true)"
  fi
}

### Detect supported Ubuntu version and codename.
# @return 0 success; exit otherwise
# @require awk
detect_ubuntu() {
  read_os_release

  case "$ubuntu_version_id" in
    "20.04") ubuntu_codename="${ubuntu_codename:-focal}" ;;
    "22.04") ubuntu_codename="${ubuntu_codename:-jammy}" ;;
    "24.04") ubuntu_codename="${ubuntu_codename:-noble}" ;;
    *)
      log_error "Unsupported Ubuntu version: ${ubuntu_version_id}. Supported: 20.04/22.04/24.04"
      exit 1
      ;;
  esac

  if [[ -z "$ubuntu_codename" ]]; then
    log_error "Failed to determine Ubuntu codename."
    exit 1
  fi

  log_info "Detected Ubuntu ${ubuntu_version_id} (${ubuntu_codename})"
}

### Decide which sources format to manage.
# @return 0
# @require test
choose_sources_mode() {
  if [[ -f "$APT_DEB822_SOURCES" ]]; then
    sources_mode="deb822"
  elif [[ -f "$APT_SOURCES_LIST" ]]; then
    sources_mode="list"
  else
    # > Defensive: if neither exists, still proceed by creating sources.list
    sources_mode="list"
  fi
  log_info "Sources mode: ${sources_mode}"
}

### Ensure backup directory exists.
# @param backup_dir string Directory path
# @return 0
# @require mkdir
ensure_backup_dir() {
  local dir="$1"
  if [[ -z "$dir" ]]; then
    log_error "Backup directory is empty."
    exit 1
  fi
  mkdir -p "$dir"
  log_debug "Backup directory ensured: $dir"
}

### Backup an APT sources file if it exists.
# @param src_path string File path to backup
# @param backup_dir string Backup directory
# @return 0
# @require cp date
backup_file_if_exists() {
  local src_path="$1"
  local dir="$2"

  if [[ ! -e "$src_path" ]]; then
    log_warn "Skip backup (not found): $src_path"
    return 0
  fi

  local ts
  ts="$(date '+%Y%m%d-%H%M%S')"
  local base
  base="$(basename "$src_path")"
  local dst="${dir}/${base}.${ts}.bak"

  cp -a "$src_path" "$dst"
  log_info "Backed up: $src_path -> $dst"
}

### Backup relevant source files.
# @return 0
# @require cp
backup_sources() {
  backup_file_if_exists "$APT_SOURCES_LIST" "$backup_dir"
  backup_file_if_exists "$APT_DEB822_SOURCES" "$backup_dir"
}

### Ask for confirmation unless --yes is given.
# @return 0 if confirmed; exit otherwise
# @require read
confirm_action() {
  if [[ "$assume_yes" == "true" ]]; then
    log_info "Non-interactive mode: --yes"
    return 0
  fi

  log_warn "About to replace APT sources with TUNA mirror:"
  log_warn "  ${TUNA_UBUNTU_URI}"
  log_warn "This will modify system APT source configuration."
  printf "Continue? [y/N]: " >&2
  local ans=""
  read -r ans
  case "$ans" in
    y|Y|yes|YES) return 0 ;;
    *) log_info "Cancelled by user."; exit 0 ;;
  esac
}

#------------------------------------------------------------------------------
# Core actions
#------------------------------------------------------------------------------

### Write traditional /etc/apt/sources.list using TUNA mirror.
# @param codename string Ubuntu codename (focal/jammy/noble)
# @return 0
# @require cat chmod chown mv
write_sources_list_tuna() {
  local codename="$1"
  local tmp_file
  tmp_file="$(mktemp)"

  # > Provide standard suites: release, updates, backports, security
  cat >"$tmp_file" <<EOF
#------------------------------------------------------------------------------#
# Ubuntu ${codename} - TUNA Mirror
# Generated by: ${SCRIPT_NAME} v${SCRIPT_VERSION}
# Mirror: ${TUNA_UBUNTU_URI}
#------------------------------------------------------------------------------#

deb ${TUNA_UBUNTU_URI} ${codename} main restricted universe multiverse
deb ${TUNA_UBUNTU_URI} ${codename}-updates main restricted universe multiverse
deb ${TUNA_UBUNTU_URI} ${codename}-backports main restricted universe multiverse
deb ${TUNA_UBUNTU_URI} ${codename}-security main restricted universe multiverse

# If you want source packages, uncomment the following lines:
# deb-src ${TUNA_UBUNTU_URI} ${codename} main restricted universe multiverse
# deb-src ${TUNA_UBUNTU_URI} ${codename}-updates main restricted universe multiverse
# deb-src ${TUNA_UBUNTU_URI} ${codename}-backports main restricted universe multiverse
# deb-src ${TUNA_UBUNTU_URI} ${codename}-security main restricted universe multiverse
EOF

  chmod 0644 "$tmp_file"
  chown root:root "$tmp_file"

  # > Atomic replace
  mkdir -p "$(dirname "$APT_SOURCES_LIST")"
  mv -f "$tmp_file" "$APT_SOURCES_LIST"

  log_info "Updated: $APT_SOURCES_LIST"
}

### Patch Deb822 ubuntu.sources to use TUNA mirror.
# @param deb822_file string Path to ubuntu.sources
# @param tuna_uri string The TUNA mirror base URI
# @return 0
# @require sed cp mktemp chmod chown mv grep
patch_deb822_sources_tuna() {
  local deb822_file="$1"
  local tuna_uri="$2"

  if [[ ! -f "$deb822_file" ]]; then
    log_warn "Deb822 sources file not found: $deb822_file"
    return 0
  fi

  local tmp_file
  tmp_file="$(mktemp)"

  cp -a "$deb822_file" "$tmp_file"

  # > Replace any "URIs:" line to TUNA; keep other Deb822 fields unchanged.
  # > Some systems may have multiple stanzas; this applies globally.
  sed -i -E "s|^URIs:[[:space:]]+.*$|URIs: ${tuna_uri}|g" "$tmp_file"

  # > Defensive check: ensure we still have at least one URIs line
  if ! grep -qE '^URIs:[[:space:]]+' "$tmp_file"; then
    log_error "Deb822 patch failed: no 'URIs:' line found after edit."
    exit 1
  fi

  chmod 0644 "$tmp_file"
  chown root:root "$tmp_file"
  mv -f "$tmp_file" "$deb822_file"

  log_info "Patched Deb822 sources: $deb822_file"
}

### Apply TUNA sources according to detected mode.
# @return 0
# @require true
apply_tuna_sources() {
  case "$sources_mode" in
    deb822)
      patch_deb822_sources_tuna "$APT_DEB822_SOURCES" "$TUNA_UBUNTU_URI"
      ;;
    list)
      write_sources_list_tuna "$ubuntu_codename"
      ;;
    *)
      log_error "Unknown sources mode: $sources_mode"
      exit 1
      ;;
  esac
}

### Run apt-get update if requested.
# @return 0
# @require apt-get
apt_update() {
  if [[ "$do_update" != "true" ]]; then
    log_info "Skip apt-get update (use --update to enable)."
    return 0
  fi

  log_info "Running: apt-get update"
  # > Use noninteractive frontend to reduce prompts in some envs
  DEBIAN_FRONTEND=noninteractive apt-get update
  log_info "apt-get update completed."
}

### Print summary.
# @return 0
summary() {
  log_info "Done."
  log_info "Backup directory: ${backup_dir}"
  log_info "Mirror applied: ${TUNA_UBUNTU_URI}"
  log_info "Ubuntu: ${ubuntu_version_id} (${ubuntu_codename}), mode: ${sources_mode}"
}

#------------------------------------------------------------------------------
# Main
#------------------------------------------------------------------------------

### Main entry.
# @param args string[] CLI args
# @return 0 success; non-zero otherwise
# @require bash
main() {
  parse_args "$@"
  setup_traps
  require_root
  detect_ubuntu
  choose_sources_mode
  ensure_backup_dir "$backup_dir"
  backup_sources
  confirm_action
  apply_tuna_sources
  apt_update
  summary
}

main "$@"
