#!/usr/bin/env bash

# ---------------------------------------------------------
# Formatting & Colors
# ---------------------------------------------------------
BOLD=$'\e[1m'
RED=$'\e[31m'
GREEN=$'\e[32m'
YELLOW=$'\e[33m'
BLUE=$'\e[34m'
RESET=$'\e[0m'

# ---------------------------------------------------------
# Utilities
# ---------------------------------------------------------
LOG_FILE="/tmp/workstation_setup.log"
> "$LOG_FILE"
FAILED_TASKS=()

run_task() {
  local task_name="$1"
  shift

  local plain_pad
  plain_pad=$(printf "%-50s" "${task_name}...")

  printf "  %s " "$plain_pad"

  local tmp_log
  tmp_log=$(mktemp)

  "$@" > "$tmp_log" 2>&1 < /dev/null &
  local pid=$!
  local spinner='|/-\'
  local i=0

  while kill -0 "$pid" 2>/dev/null; do
    if [ -t 1 ]; then
      printf "\r  %s [%s]" "$plain_pad" "${spinner:i%4:1}"
      i=$((i + 1))
    fi
    sleep 0.1
  done

  local status=0
  wait "$pid" || status=$?

  if [ -t 1 ]; then
    printf "\r  %s [" "$plain_pad"
    [ "$status" -eq 0 ] && printf "${GREEN}${BOLD}  OK  ${RESET}" || printf "${RED}${BOLD} FAIL ${RESET}"
    printf "]\n"
  else
    [ "$status" -eq 0 ] && echo -e "[${GREEN}${BOLD}  OK  ${RESET}]" || echo -e "[${RED}${BOLD} FAIL ${RESET}]"
  fi

  if [ "$status" -eq 0 ]; then
    echo "  $plain_pad [  OK  ]" >> "$LOG_FILE"
  else
    FAILED_TASKS+=("$task_name")
    echo "  $plain_pad [ FAIL ]" >> "$LOG_FILE"
  fi

  [ -s "$tmp_log" ] && sed 's/^/    /' "$tmp_log" >> "$LOG_FILE"
  rm -f "$tmp_log"
  return "$status"
}

show_summary() {
  if [ ${#FAILED_TASKS[@]} -gt 0 ]; then
    echo -e "${YELLOW}${BOLD}=== Tasks with Errors ===${RESET}"
    for task in "${FAILED_TASKS[@]}"; do
      echo -e "  ${RED}✗${RESET} $task"
    done
    echo -e "\nFull log: ${BOLD}$LOG_FILE${RESET}"
  else
    echo -e "${GREEN}${BOLD}All tasks completed successfully.${RESET}"
  fi
}

apt_cmd() {
  sudo env DEBIAN_FRONTEND=noninteractive NEEDRESTART_MODE=a APT_LISTCHANGES_FRONTEND=none apt-get "$@"
}

has_cmd() {
  command -v "$1" >/dev/null 2>&1
}

log_section() {
  local section_name="$1"

  # 1. Print to terminal (with bold formatting)
  echo -e "${BOLD}--- ${section_name} ---${RESET}"

  # 2. Print to log file (conditionally add newline if file is not empty)
  if [ -s "$LOG_FILE" ]; then
    echo -e "\n--- ${section_name} ---" >> "$LOG_FILE"
  else
    echo -e "--- ${section_name} ---" >> "$LOG_FILE"
  fi
}

# ---------------------------------------------------------
# Argument Parsing & State
# ---------------------------------------------------------
INTERACTIVE=1

MODULES=(
  "System Packages & Repositories"
  "Homebrew & CLI Tooling"
  "Dotfiles (Zsh, Stow)"
  "Podman Engine"
  "Docker Engine"
  "DNSMasq Container"
  "SSH Config (COLORTERM)"
  "Arch WSL Bootstrap"
)

SELECTED=(1 1 1 1 1 1 1 1) # Default all ON

while [[ "$#" -gt 0 ]]; do
  case $1 in
    --all) SELECTED=(1 1 1 1 1 1 1 1); INTERACTIVE=0; shift ;;
    -y|--yes) INTERACTIVE=0; shift ;;
    *) echo -e "${RED}Unknown parameter passed: $1${RESET}"; exit 1 ;;
  esac
done

READ_FROM_TTY=0

# If stdin is piped (e.g. curl | bash), still allow interaction via /dev/tty.
if [ "$INTERACTIVE" -eq 1 ] && [ ! -t 0 ]; then
  if [ -r /dev/tty ]; then
    READ_FROM_TTY=1
  else
    echo -e "${YELLOW}Warning: No TTY detected (script likely piped). Running in headless --all mode.${RESET}"
    SELECTED=(1 1 1 1 1 1 1)
    INTERACTIVE=0
  fi
fi

# ---------------------------------------------------------
# Native Interactive Module Selection (Single Keystroke)
# ---------------------------------------------------------
if [ "$INTERACTIVE" -eq 1 ]; then
  while true; do
    clear
    echo -e "${BLUE}${BOLD}======================================================${RESET}"
    echo -e "${BOLD}         Workstation Setup - Module Selection         ${RESET}"
    echo -e "${BLUE}${BOLD}======================================================${RESET}\n"

    for i in "${!MODULES[@]}"; do
      num=$((i+1))
      if [ "${SELECTED[$i]}" -eq 1 ]; then
        echo -e "  $num) ${GREEN}[x]${RESET} ${MODULES[$i]}"
      else
        echo -e "  $num) ${RED}[ ]${RESET} ${MODULES[$i]}"
      fi
    done

    echo -e "\nPress 1-8 to toggle, a for all/none, Enter to proceed, q to quit..."

    if [ "$READ_FROM_TTY" -eq 1 ]; then
      read -s -n 1 REPLY < /dev/tty
    else
      read -s -n 1 REPLY
    fi

    case $REPLY in
      1|2|3|4|5|6|7|8)
        idx=$((REPLY-1))
        SELECTED[$idx]=$((1 - SELECTED[$idx]))
        ;;
      a|A)
        all_selected=1
        for i in "${!SELECTED[@]}"; do
          if [ "${SELECTED[$i]}" -eq 0 ]; then
            all_selected=0
            break
          fi
        done

        if [ "$all_selected" -eq 1 ]; then
          SELECTED=(0 0 0 0 0 0 0 0)
        else
          SELECTED=(1 1 1 1 1 1 1 1)
        fi
        ;;
      ''|$'\n')
        break
        ;;
      q|Q)
        echo -e "\nInstallation cancelled."
        exit 0
        ;;
      *)
        ;;
    esac
  done
fi

DO_SYSTEM=${SELECTED[0]}
DO_HOMEBREW=${SELECTED[1]}
DO_DOTFILES=${SELECTED[2]}
DO_PODMAN=${SELECTED[3]}
DO_DOCKER=${SELECTED[4]}
DO_DNSMASQ=${SELECTED[5]}
DO_SSH=${SELECTED[6]}
DO_ARCH_WSL=${SELECTED[7]}

if [ "$DO_SYSTEM" -eq 0 ] && [ "$DO_HOMEBREW" -eq 0 ] && [ "$DO_DOTFILES" -eq 0 ] && [ "$DO_PODMAN" -eq 0 ] && [ "$DO_DOCKER" -eq 0 ] && [ "$DO_DNSMASQ" -eq 0 ] && [ "$DO_SSH" -eq 0 ] && [ "$DO_ARCH_WSL" -eq 0 ]; then
  echo "No modules selected. Exiting."
  exit 0
fi

clear
echo -e "${BLUE}${BOLD}======================================================${RESET}"
echo -e "${BOLD}Starting Installation...${RESET}"
echo -e "Detailed logs saved to: ${BOLD}$LOG_FILE${RESET}"
echo -e "${BLUE}${BOLD}======================================================${RESET}\n"

# Request sudo upfront (redirecting stdin from /dev/tty fixes pipe issues)
sudo -v < /dev/tty
while true; do sudo -n true; sleep 60; kill -0 "$$" || exit; done 2>/dev/null &

USER_DIR=$HOME
USER_ID=$(whoami)

# Detect OS
if [ -f /etc/os-release ]; then
  . /etc/os-release
  OS=$ID
else
  echo -e "${RED}Cannot detect OS.${RESET}"
  exit 1
fi
# ---------------------------------------------------------
# 1. System Packages & Repos
# ---------------------------------------------------------

configure_pacman() {
  if [ ! -f /etc/pacman.conf ]; then
    return 1
  fi

  for opt in Color ILoveCandy DisableDownloadTimeout; do
    if ! grep -q "^$opt" /etc/pacman.conf; then
      sudo sed -i "/^\[options\]/a $opt" /etc/pacman.conf
    fi
  done
}

add_cachyos_repo() {
  grep -iq cachyos /etc/pacman.conf && return 0
  local tmp_dir; tmp_dir=$(mktemp -d)
  cd "$tmp_dir" && curl -sL https://mirror.cachyos.org/cachyos-repo.tar.xz | tar xJ && cd cachyos-repo && yes | sudo ./cachyos-repo.sh
  rm -rf "$tmp_dir"
}

check_system_deps() {
  for cmd in make ps curl file git zsh stow micro; do
    if ! has_cmd "$cmd"; then
      return 1 # A dependency is missing
    fi
  done
  return 0 # All dependencies exist
}

if [ "$DO_SYSTEM" -eq 1 ]; then
  log_section "System Packages"

  if check_system_deps; then
    run_task "Base dependencies already installed" true
  else
    if [[ "$OS" == "arch" || "$OS" == "cachyos" ]]; then

      run_task "Configuring pacman.conf options" configure_pacman
      run_task "Adding CachyOS repo" add_cachyos_repo

      run_task "Syncing pacman repositories" sudo pacman -Syu --noconfirm
      run_task "Installing base dependencies" sudo pacman -S --noconfirm base-devel procps-ng curl file git zsh stow micro helix paru ripgrep fastfetch yazi lazydocker lazygit bat fd eza zoxide
      run_task "Symlinking helix to /usr/bin/hx" sudo ln -sf /usr/bin/helix /usr/bin/hx

    # Expanded OS check to include Parrot OS and LMDE
    elif [[ "$OS" == "debian" || "$OS" == "ubuntu" || "$OS" == "parrot" || "$OS" == "lmde" ]]; then

      run_task "Updating apt repositories" apt_cmd update
      run_task "Installing base dependencies" apt_cmd install -y ca-certificates curl gnupg zsh stow micro git

    else
      run_task "Unsupported OS for system packages: $OS" false
    fi
  fi

  echo ""
fi
# ---------------------------------------------------------
# 2. Homebrew & Packages
# ---------------------------------------------------------
if [ "$DO_HOMEBREW" -eq 1 ]; then
  log_section "Homebrew"
  HOMEBREW_PREFIX="/home/linuxbrew/.linuxbrew"
  if [ ! -f "$HOMEBREW_PREFIX/bin/brew" ]; then
    run_task "Creating homebrew directory" sudo mkdir -p /home/linuxbrew
    run_task "Setting homebrew ownership" sudo chown "$USER_ID:$USER_ID" /home/linuxbrew
    run_task "Installing Homebrew" bash -c 'NONINTERACTIVE=1 /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)"'
  else
    run_task "Homebrew already installed" true
  fi

  run_task "Adding brew shellenv to .bashrc" bash -c 'grep -q "brew shellenv" ~/.bashrc || echo '"'"'eval "$(/home/linuxbrew/.linuxbrew/bin/brew shellenv)"'"'"' >> ~/.bashrc'

  eval "$("$HOMEBREW_PREFIX"/bin/brew shellenv)"

  run_task "Installing Homebrew CLI tools" brew install ripgrep fastfetch yazi lazydocker lazygit bat fd eza zoxide helix

  run_task "Linking brew binaries to /usr/bin" sudo bash -c '
    for f in /home/linuxbrew/.linuxbrew/bin/*; do
      [ -f "$f" ] && [ ! -e "/usr/bin/$(basename "$f")" ] && ln -s "$f" "/usr/bin/$(basename "$f")"
    done
  '
  echo ""
fi
# ---------------------------------------------------------
# 3. Dotfiles
# ---------------------------------------------------------

clear_stow_conflicts() {
  local dotfiles="$1" target="$2" pkgs="$3"
  local conflicts
  conflicts=$(stow -d "$dotfiles" -t "$target" -n --no-folding $pkgs 2>&1 | sed -n 's/.*existing target is not owned by stow:[[:space:]]*//p')
  for f in $conflicts; do
    rm -f "$target/$f"
  done
}

if [ "$DO_DOTFILES" -eq 1 ]; then
  log_section "Dotfiles"
  DOTFILES_DEST="$USER_DIR/dotfiles"

  if [ ! -d "$DOTFILES_DEST" ]; then
    run_task "Cloning dotfiles repository" git clone https://github.com/danne56/dotfiles "$DOTFILES_DEST"
  else
    run_task "Pulling latest dotfiles" git -C "$DOTFILES_DEST" pull origin main
  fi

  BACKUP_FILES=(
    ".zshrc" ".zprofile" ".bashrc" ".tmux.conf" ".p10k.zsh" ".gitconfig"
    ".config/yazi" ".config/micro"
  )

  run_task "Backing up existing configs" bash -c '
    ts=$(date +%s)
    for file in "$@"; do
      p="$HOME/$file"
      if [ ! -L "$p" ] && [ -d "$p" ]; then
        find "$p" -type f -exec sh -c "mv -f \"\$1\" \"\$1.bak.$ts\"" _ {} \;
        rmdir "$p" 2>/dev/null || true
      elif [ ! -L "$p" ] && [ -e "$p" ]; then
        mv -f "$p" "$p.bak.$ts"
      fi

      p="/root/$file"
      if [ ! -L "$p" ] && [ -d "$p" ]; then
        sudo find "$p" -type f -exec sh -c "mv -f \"\$1\" \"\$1.bak.$ts\"" _ {} \;
        sudo rmdir "$p" 2>/dev/null || true
      elif [ ! -L "$p" ] && [ -e "$p" ]; then
        sudo mv -f "$p" "$p.bak.$ts"
      fi
    done
  ' _ "${BACKUP_FILES[@]}"

  PKGS="zsh bash micro yazi git shell helix tmux scripts"
  run_task "Clearing stow conflicts for user" clear_stow_conflicts "$DOTFILES_DEST" "$USER_DIR" "$PKGS"
  run_task "Stowing packages for user" stow -d "$DOTFILES_DEST" -t "$USER_DIR" -v --restow --no-folding $PKGS
  run_task "Clearing stow conflicts for root" sudo bash -c "$(declare -f clear_stow_conflicts); clear_stow_conflicts \$1 \$2 \"\$3\"" _ "$DOTFILES_DEST" /root "$PKGS"
  run_task "Stowing packages for root" sudo stow -d "$DOTFILES_DEST" -t /root -v --restow --no-folding $PKGS

  if [ "$(getent passwd "$USER_ID" | cut -d: -f7)" = "/usr/bin/zsh" ]; then
    run_task "Default shell already Zsh" true
  else
    run_task "Setting default shell to Zsh" sudo chsh -s /usr/bin/zsh "$USER_ID"
  fi

  run_task "Creating /etc/zsh directory" sudo mkdir -p /etc/zsh
  run_task "Setting ZDOTDIR in zshenv" sudo bash -c 'echo "export ZDOTDIR=\"\$HOME/.config/zsh\"" > /etc/zsh/zshenv'
  run_task "Setting zshenv permissions" sudo chmod 0644 /etc/zsh/zshenv

  if [ -f /etc/sudoers.d/10-env_keep ] && grep -q 'MICRO_TRUECOLOR' /etc/sudoers.d/10-env_keep; then
    run_task "Sudo env_keep already configured" true
  else
    run_task "Configuring sudo env_keep" bash -c 'echo '"'"'Defaults env_keep += "EDITOR MICRO_TRUECOLOR"'"'"' | sudo tee /etc/sudoers.d/10-env_keep > /dev/null && sudo chmod 0440 /etc/sudoers.d/10-env_keep'
  fi

  echo ""
fi
# ---------------------------------------------------------
# 4. Podman
# ---------------------------------------------------------
if [ "$DO_PODMAN" -eq 1 ]; then
  log_section "Podman"

  if has_cmd docker; then
    if [[ "$OS" == "arch" || "$OS" == "cachyos" ]]; then
      # Only attempt to remove packages that are actually installed
      DOCKER_PKGS=$(pacman -Qq docker docker-compose docker-buildx 2>/dev/null || true)
      if [ -n "$DOCKER_PKGS" ]; then
        run_task "Removing Docker packages" sudo pacman -Rns --noconfirm $DOCKER_PKGS
      fi
    elif [[ "$OS" == "debian" || "$OS" == "ubuntu" ]]; then
      run_task "Removing Docker packages" apt_cmd remove -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin docker-compose
      run_task "Removing distro docker packages" apt_cmd remove -y docker.io docker-compose-v2 docker-compose
    fi

    if systemctl list-unit-files | grep -q '^docker\.service'; then
      run_task "Stopping Docker daemon" sudo systemctl disable --now docker
    else
      run_task "Docker daemon not installed" true
    fi
  else
    run_task "Docker not installed" true
  fi

  if has_cmd podman; then
    run_task "Podman already installed" true
  else
    if [[ "$OS" == "arch" || "$OS" == "cachyos" ]]; then
      run_task "Installing Podman packages" sudo pacman -S --noconfirm podman podman-compose podman-docker docker-compose
    elif [[ "$OS" == "debian" || "$OS" == "ubuntu" ]]; then
      run_task "Updating apt for Podman" apt_cmd update
      run_task "Installing Podman packages" apt_cmd install -y podman podman-compose podman-docker docker-compose-v2
    else
      run_task "Unsupported OS for Podman: $OS" false
    fi
  fi

  if systemctl --user is-enabled --quiet podman.socket 2>/dev/null && systemctl --user is-active --quiet podman.socket 2>/dev/null; then
    run_task "User Podman socket already enabled" true
  else
    run_task "Enabling User Podman socket" systemctl --user enable --now podman.socket
  fi

  if systemctl is-enabled --quiet podman.socket 2>/dev/null && systemctl is-active --quiet podman.socket 2>/dev/null; then
    run_task "Root Podman socket already enabled" true
  else
    run_task "Enabling Root Podman socket" sudo systemctl enable --now podman.socket
  fi

  USER_ZDOTDIR="${ZDOTDIR:-$USER_DIR/.config/zsh}"
  USER_PODMAN_ENV='export DOCKER_HOST=unix://$XDG_RUNTIME_DIR/podman/podman.sock'

  if [ -f "$USER_ZDOTDIR/.zshrc" ] && grep -q 'podman/podman.sock' "$USER_ZDOTDIR/.zshrc"; then
    run_task "DOCKER_HOST export already in user .zshrc" true
  else
    run_task "Adding DOCKER_HOST export to user .zshrc" bash -c 'mkdir -p "$1" && touch "$2" && printf "\n%s\n" "$3" >> "$2"' _ "$USER_ZDOTDIR" "$USER_ZDOTDIR/.zshrc" "$USER_PODMAN_ENV"
  fi

  if [ "${XDG_RUNTIME_DIR:-}" != "" ]; then
    export DOCKER_HOST="unix://$XDG_RUNTIME_DIR/podman/podman.sock"
  fi

# Prevent interactive short-name registry resolution prompts from blocking tasks
  REGISTRIES_CONF="/etc/containers/registries.conf"
  if [ -f "$REGISTRIES_CONF" ] && ! grep -q "unqualified-search-registries" "$REGISTRIES_CONF"; then
    run_task "Setting unqualified search registries" sudo bash -c "echo 'unqualified-search-registries = [\"docker.io\", \"quay.io\"]' >> $REGISTRIES_CONF"
  fi

  # Disable compose warning logs
  CONTAINERS_CONF_D="/etc/containers/containers.conf.d"
  run_task "Creating containers config directory" sudo mkdir -p "$CONTAINERS_CONF_D"
  run_task "Disabling compose warning logs" sudo bash -c 'echo "compose_warning_logs = false" > /etc/containers/containers.conf.d/99-compose-warnings.conf'

  echo ""
fi
# ---------------------------------------------------------
# 4. Docker
# ---------------------------------------------------------
add_debian_docker_repo() {
  local os_name="$1"

  if [ -f /etc/apt/sources.list.d/docker.list ] && grep -q '^deb .*download.docker.com/linux/' /etc/apt/sources.list.d/docker.list; then
    return 0
  fi

  local arch
  arch=$(dpkg --print-architecture)

  # Source os-release to get VERSION_CODENAME
  . /etc/os-release
  local repo_os="$os_name"
  local codename="$VERSION_CODENAME"

  # Normalize Debian derivatives (Parrot, LMDE)
  if [[ "$os_name" == "parrot" || "$os_name" == "lmde" ]]; then
    repo_os="debian"

    # Map derivative codenames to their Debian upstream equivalents
    case "$codename" in
      ara|elsie) codename="bullseye" ;;     # Parrot 5.x / LMDE 5
      lory|faye) codename="bookworm" ;;     # Parrot 6.x / LMDE 6
      *) codename="bookworm" ;;             # Safe default for newer/unknown versions
    esac
  fi

  # Fallback if VERSION_CODENAME is completely missing from os-release
  if [ -z "$codename" ]; then
    codename="bookworm"
  fi

  echo "deb [arch=${arch} signed-by=/etc/apt/keyrings/docker.asc] https://download.docker.com/linux/${repo_os} ${codename} stable" | sudo tee /etc/apt/sources.list.d/docker.list > /dev/null
}

configure_docker_daemon() {
  if [ -f /etc/docker/daemon.json ] && grep -q 'storage-driver.*overlay2' /etc/docker/daemon.json; then
    return 0
  fi

  sudo mkdir -p /etc/docker
  cat << 'EOF' | sudo tee /etc/docker/daemon.json > /dev/null
{
  "storage-driver": "overlay2",
  "log-opts": {
    "max-size": "10m",
    "max-file": "5"
  }
}
EOF
}

enable_docker_daemon() {
  if systemctl is-enabled docker >/dev/null 2>&1 && systemctl is-active docker >/dev/null 2>&1; then
    return 0
  fi
  sudo systemctl enable --now docker
}

add_user_to_docker() {
  local target_user="$1"
  if id -nG "$target_user" | grep -qw docker; then
    return 0
  fi
  sudo usermod -aG docker "$target_user"
}

ensure_overlay2_driver() {
  if sudo docker info 2>/dev/null | grep -q 'Storage Driver: overlay2'; then
    return 0
  fi
  sudo systemctl restart docker
  sleep 2
  sudo docker info 2>/dev/null | grep -q 'Storage Driver: overlay2'
}

if [ "$DO_DOCKER" -eq 1 ]; then
  log_section "Docker"

  if [ "$DO_PODMAN" -eq 1 ]; then
    run_task "Skipping Docker because Podman module is selected" true
    echo ""
  else
    if has_cmd docker; then
      run_task "Docker already installed" true
    else
      if [[ "$OS" == "arch" || "$OS" == "cachyos" ]]; then
        run_task "Installing Docker packages" sudo pacman -S --noconfirm docker docker-compose docker-buildx

      # Expand the OS check here to let Parrot and LMDE through
      elif [[ "$OS" == "debian" || "$OS" == "ubuntu" || "$OS" == "parrot" || "$OS" == "lmde" ]]; then
        run_task "Creating keyrings directory" sudo install -m 0755 -d /etc/apt/keyrings
        run_task "Downloading Docker GPG key" sudo curl -4 --http1.1 --tls-max 1.2 -fsSL "https://download.docker.com/linux/debian/gpg" -o /etc/apt/keyrings/docker.asc
        run_task "Setting key permissions" sudo chmod a+r /etc/apt/keyrings/docker.asc

        # Call the new smart helper function
        run_task "Adding Docker apt repository" add_debian_docker_repo "$OS"

        run_task "Updating apt for Docker" apt_cmd update
        run_task "Installing Docker Engine" apt_cmd install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin

      else
        run_task "Unsupported OS for Docker: $OS" false
      fi
    fi

    run_task "Configuring Docker daemon" configure_docker_daemon
    run_task "Enabling Docker daemon" enable_docker_daemon
    run_task "Adding user to docker group" add_user_to_docker "$USER_ID"
    run_task "Verifying overlay2 storage driver" ensure_overlay2_driver

    echo ""
  fi
fi
# ---------------------------------------------------------
# 5. DNSMasq
# ---------------------------------------------------------

write_dnsmasq_files() {
  local dir="$1"
  mkdir -p "$dir/dnsmasq.d"
  cat > "$dir/docker-compose.yaml" <<'EOF'
services:
  dnsmasq:
    image: danne56/dnsmasq
    container_name: dnsmasq
    environment:
      DNS1: "1.0.0.1"
      DNS2: "1.1.1.1"
    ports:
      - 127.0.0.1:53:53/udp
      - 127.0.0.1:53:53/tcp
    cap_add:
      - NET_ADMIN
    restart: always
    volumes:
      - ./dnsmasq.d/:/etc/dnsmasq.d/
EOF
  cat > "$dir/dnsmasq.d/my.conf" <<'EOF'
listen-address=127.0.0.1
address=/mint.test/127.0.0.1
EOF
}

if [ "$DO_DNSMASQ" -eq 1 ]; then
  log_section "DNSMasq"

  if [[ "$OS" == "arch" || "$OS" == "cachyos" ]]; then
    if pacman -Q dnsmasq >/dev/null 2>&1; then
      run_task "Removing local DNS conflicts" sudo pacman -Rns --noconfirm dnsmasq
    else
      run_task "No local dnsmasq package to remove" true
    fi
  elif [[ "$OS" == "debian" || "$OS" == "ubuntu" ]]; then
    if dpkg-query -W -f='${Status}' dnsmasq 2>/dev/null | grep -q "install ok installed"; then
      run_task "Removing local DNS conflicts" apt_cmd remove -y dnsmasq
    else
      run_task "No local dnsmasq package to remove" true
    fi
  fi

  DNSMASQ_DIR="$USER_DIR/compose/dnsmasq"

  run_task "Writing DNSMasq files" write_dnsmasq_files "$DNSMASQ_DIR"

  if command -v podman >/dev/null 2>&1; then
    COMPOSE_CMD="sudo podman compose"
    DOCKER_CMD="sudo podman"
  elif command -v docker >/dev/null 2>&1; then
    COMPOSE_CMD="sudo docker compose"
    DOCKER_CMD="sudo docker"
  else
    run_task "No container runtime found" false
  fi

  run_task "Pulling DNSMasq image" bash -c "cd $DNSMASQ_DIR && $COMPOSE_CMD pull"
  run_task "Starting DNSMasq container" bash -c "cd $DNSMASQ_DIR && $COMPOSE_CMD up -d"

  run_task "Disabling DNSStubListener in resolved.conf" sudo sed -i 's/^#*DNSStubListener=.*/DNSStubListener=no/' /etc/systemd/resolved.conf
  run_task "Restarting systemd-resolved" sudo systemctl restart systemd-resolved || true

  if grep -qiE '(microsoft|wsl)' /proc/version 2>/dev/null; then
    run_task "Configuring WSL /etc/wsl.conf" sudo bash -c '
      conf="/etc/wsl.conf"
      if [ -f "$conf" ] && grep -q "^\[network\]" "$conf"; then
        if grep -q "generateResolvConf" "$conf"; then
          sed -i "s/^[[:space:]]*generateResolvConf[[:space:]]*=.*/generateResolvConf = false/" "$conf"
        else
          sed -i "/^\[network\]/a generateResolvConf = false" "$conf"
        fi
      elif [ -f "$conf" ]; then
        printf "\n[network]\ngenerateResolvConf = false\n" >> "$conf"
      else
        printf "[network]\ngenerateResolvConf = false\n" > "$conf"
      fi
    '
  fi

  run_task "Removing symlink /etc/resolv.conf" sudo bash -c '[ -L /etc/resolv.conf ] && rm -f /etc/resolv.conf || true'
  run_task "Writing /etc/resolv.conf" sudo bash -c 'echo "# Managed by script - Pointing to local dnsmasq
nameserver 127.0.0.1" > /etc/resolv.conf'
  echo ""
fi
# ---------------------------------------------------------
# 6. SSH Config
# ---------------------------------------------------------
if [ "$DO_SSH" -eq 1 ]; then
  log_section "SSH"

  run_task "Configuring SSH COLORTERM env" bash -c '
    for f in /etc/ssh/ssh_config /etc/ssh/sshd_config; do
      [ -f "$f" ] || continue
      case "$f" in
        *sshd_config*) p="AcceptEnv"; l="AcceptEnv LANG LC_* COLORTERM" ;;
        *)             p="SendEnv";   l="    SendEnv LANG LC_* COLORTERM" ;;
      esac
      grep -qE "^[[:space:]]*$p.*COLORTERM" "$f" 2>/dev/null && continue
      sudo sed -i "/^#*[[:space:]]*$p/d" "$f"
      sudo sed -i "/^#*[[:space:]]*${p%% *}/a $l" "$f"
      changed=1
    done
    [ "${changed:-0}" -eq 1 ] && { sudo systemctl restart sshd 2>/dev/null || sudo systemctl restart ssh 2>/dev/null || true; }
  '
  echo ""
fi
# ---------------------------------------------------------
# 7. Post Install
# ---------------------------------------------------------

adjust_zinit_arch_for_user() {
  local target_file="$1"
  if [ -f "$target_file" ]; then
    sed -i "/^[[:space:]]*[^#].*OMZP::archlinux/s/^/# /" "$target_file"
  fi
}

adjust_zinit_arch_for_root() {
  local target_file="/root/.config/zsh/.zshrc"
  if [ -f "$target_file" ]; then
    sed -i "/^[[:space:]]*[^#].*OMZP::archlinux/s/^/# /" "$target_file"
  fi
}

configure_wsl() {
  sudo tee /etc/wsl.conf > /dev/null <<'EOF'
[boot]
systemd = true

[user]
default = deffa

[interop]
enabled = true
appendWindowsPath = true

[network]
generateResolvConf = false

[automount]
enabled = true
options = "metadata,umask=22,fmask=011"
EOF
}

setup_shared_root_mount() {
  cat << 'EOF' > /etc/profile.d/02-shared-root.sh
#!/usr/bin/env bash
wsl.exe -u root -e mount --make-rshared /
EOF
  chmod +x /etc/profile.d/02-shared-root.sh
}

if [ "$DO_DOTFILES" -eq 1 ]; then
  log_section "Post Install"

  USER_ZDOTDIR=${ZDOTDIR:-"$USER_DIR/.config/zsh"}

  if [[ "$OS" != "arch" && "$OS" != "cachyos" ]]; then
    run_task "Adjusting zinit arch plugin for user" adjust_zinit_arch_for_user "$USER_ZDOTDIR/.zshrc"
    run_task "Adjusting zinit arch plugin for root" sudo bash -c "$(declare -f adjust_zinit_arch_for_root); adjust_zinit_arch_for_root"
  else
    run_task "Zinit arch plugin unchanged for user" true
    run_task "Zinit arch plugin unchanged for root" true
  fi

  if has_cmd ya; then
    run_task "Installing yazi packages for user" ya pkg install
    run_task "Installing yazi packages for root" sudo env HOME=/root XDG_CONFIG_HOME=/root/.config ya pkg install
  else
    run_task "Yazi package manager not found (user)" true
    run_task "Yazi package manager not found (root)" true
  fi

  if grep -qiE '(microsoft|wsl)' /proc/version 2>/dev/null; then
    run_task "Configuring wsl.conf" sudo bash -c "$(declare -f configure_wsl); configure_wsl"

    if has_cmd podman || has_cmd docker; then
      run_task "Configuring shared root mount for WSL containers" sudo bash -c "$(declare -f setup_shared_root_mount); setup_shared_root_mount"
    else
      run_task "Skipping shared root mount script (no container runtime)" true
    fi
  else
    run_task "Skipping WSL interop configuration" true
    run_task "Skipping shared root mount script (not WSL)" true
  fi

  echo ""
fi

show_summary

echo -e "${BLUE}${BOLD}======================================================${RESET}"
echo -e "[${GREEN}${BOLD} DONE ${RESET}] Workstation setup is complete."
echo -e "${BLUE}${BOLD}======================================================${RESET}\n"

if [ "$DO_DOTFILES" -eq 1 ] && [ -t 1 ] && [ -z "${ZSH_VERSION:-}" ] && [ "${SHELL:-}" != "/usr/bin/zsh" ] && has_cmd zsh; then
  exec zsh -l
fi
# ---------------------------------------------------------
# 8. Arch WSL Bootstrap
# ---------------------------------------------------------

ensure_sudo_installed() {
  if ! command -v visudo &>/dev/null; then
    if command -v pacman &>/dev/null; then
      sudo pacman -S --noconfirm sudo
    elif command -v apt-get &>/dev/null; then
      sudo apt-get update && sudo apt-get install -y sudo
    else
      return 1
    fi
  fi
}

create_wsl_user() {
  local user="$1"
  if ! id "$user" &>/dev/null; then
    sudo useradd -m -G wheel -s /usr/bin/zsh "$user"
  else
    sudo usermod -aG wheel "$user"
  fi
}

setup_passwordless_sudo() {
  local user="$1"
  local sudoers_file="/etc/sudoers.d/00-$user-nopasswd"
  sudo mkdir -p /etc/sudoers.d
  echo "$user ALL=(ALL:ALL) NOPASSWD: ALL" | sudo tee "$sudoers_file" > /dev/null
  sudo chmod 0440 "$sudoers_file"
  sudo visudo -cf "$sudoers_file"
}

setup_systemd_linger() {
  local user="$1"
  sudo loginctl enable-linger "$user"
}

generate_wsl_locales() {
  sudo sed -i 's/^#en_US.UTF-8 UTF-8/en_US.UTF-8 UTF-8/' /etc/locale.gen
  sudo locale-gen
}

if [ "$DO_ARCH_WSL" -eq 1 ]; then
  log_section "Arch WSL Bootstrap"

  TARGET_USER="${1:-deffa}"

  if ! grep -qiE '(microsoft|wsl)' /proc/version 2>/dev/null; then
    run_task "Not WSL — skipping" true
  elif [[ $EUID -ne 0 ]]; then
    run_task "Script requires root — skipping" false
  else
    user_exists=0
    if id "$TARGET_USER" &>/dev/null; then
      user_exists=1
    fi

    run_task "Ensuring sudo is installed" ensure_sudo_installed
    run_task "Creating user $TARGET_USER" create_wsl_user "$TARGET_USER"

    if [ "$user_exists" -eq 0 ] && id "$TARGET_USER" &>/dev/null; then
      echo "[ACTION REQUIRED] Set password for $TARGET_USER:"
      passwd "$TARGET_USER"
    fi

    run_task "Configuring passwordless sudo for $TARGET_USER" setup_passwordless_sudo "$TARGET_USER"
    run_task "Enabling systemd linger for $TARGET_USER" setup_systemd_linger "$TARGET_USER"
    run_task "Generating locales" generate_wsl_locales
  fi

  echo ""
fi