#!/usr/bin/env bash
# Root shell setup for Ubuntu
# Run this script as root.
set -euo pipefail

if [ "$(id -u)" -ne 0 ]; then
  echo "Please run as root." >&2
  exit 1
fi

# =========================
# Settings
# =========================
PREFER_EDITOR="${PREFER_EDITOR:-vim}"
PACKAGES=(bash-completion vim less)

msg()  { printf "\033[1;32m[OK]\033[0m %s\n" "$*"; }
warn() { printf "\033[1;33m[WARN]\033[0m %s\n" "$*"; }
err()  { printf "\033[1;31m[ERR]\033[0m %s\n" "$*" >&2; }

backup_file() {
  local f="$1"; [ -f "$f" ] || return 0
  local ts; ts="$(date +%Y%m%d-%H%M%S)"
  cp -a "$f" "${f}.bak.${ts}"
  msg "Backup: ${f}.bak.${ts}"
}

append_once() {
  # $1=file, $2=unique tag line, $3=payload
  local file="$1" tag="$2" payload="$3"
  grep -Fqx "$tag" "$file" 2>/dev/null && return 0
  printf "\n%s\n%s\n" "$tag" "$payload" >>"$file"
}

# =========================
# 0) Ensure base packages (no sudo needed)
# =========================
if command -v apt-get >/dev/null 2>&1; then
  export DEBIAN_FRONTEND=noninteractive
  apt-get update -y || true
  apt-get install -y "${PACKAGES[@]}" || warn "Some packages failed to install; continuing."
else
  warn "apt-get not found; skipping package install."
fi

# =========================
# 1) Target files
# =========================
ROOT_HOME="/root"
BASHRC="${ROOT_HOME}/.bashrc"
INPUTRC="${ROOT_HOME}/.inputrc"
[ -f "$BASHRC" ] || touch "$BASHRC"
[ -f "$INPUTRC" ] || touch "$INPUTRC"
backup_file "$BASHRC"
backup_file "$INPUTRC"

# =========================
# 2) force_color_prompt=yes
# =========================
if grep -Eq '^[[:space:]]*#?[[:space:]]*force_color_prompt[[:space:]]*=' "$BASHRC"; then
  sed -i '0,/^[[:space:]]*#\{0,1\}[[:space:]]*force_color_prompt[[:space:]]*=.*/s//force_color_prompt=yes/' "$BASHRC"
else
  sed -i '1iforce_color_prompt=yes' "$BASHRC"
fi
msg "Enabled force_color_prompt=yes for root."

# =========================
# 3) bash-completion
# =========================
COMPLETION_TAG="# >>> chatgpt-root: bash-completion enable >>>"
COMPLETION_BLOCK=$(cat <<'EOF'
# >>> chatgpt-root: bash-completion enable >>>
if [ -n "$BASH_VERSION" ] && ! shopt -oq posix; then
  if [ -f /usr/share/bash-completion/bash_completion ]; then
    . /usr/share/bash-completion/bash_completion
  elif [ -f /etc/bash_completion ]; then
    . /etc/bash_completion
  fi
fi
# <<< chatgpt-root: bash-completion enable <<<
EOF
)
append_once "$BASHRC" "$COMPLETION_TAG" "$COMPLETION_BLOCK"
msg "Ensured bash-completion is sourced for root."

# =========================
# 4) History: timestamps + merge + safer defaults
# =========================
HIST_TAG="# >>> chatgpt-root: history settings >>>"
HIST_BLOCK=$(cat <<'EOF'
# >>> chatgpt-root: history settings >>>
# Timestamps (YYYY-MM-DD HH:MM:SS)
export HISTTIMEFORMAT='%F %T '
# Safer history: ignore duplicates & commands starting with a space; remove duplicate older entries
export HISTCONTROL=ignoredups:ignorespace:erasedups
export HISTSIZE=300000
export HISTFILESIZE=600000
# Merge history across concurrent shells
shopt -s histappend
if [[ -n "${PROMPT_COMMAND:-}" ]]; then
  case "$PROMPT_COMMAND" in
    *'history -a; history -c; history -r'*) : ;;
    *) PROMPT_COMMAND="history -a; history -c; history -r; $PROMPT_COMMAND" ;;
  esac
else
  PROMPT_COMMAND='history -a; history -c; history -r'
fi
# <<< chatgpt-root: history settings <<<
EOF
)
append_once "$BASHRC" "$HIST_TAG" "$HIST_BLOCK"
msg "Enabled history timestamps and merging for root."

# =========================
# 5) Prompt, aliases, shopt, editor/pager (root-flavored)
# =========================
QOL_TAG="# >>> chatgpt-root: QoL shell block >>>"
QOL_BLOCK=$(cat <<'EOF'
# >>> chatgpt-root: QoL shell block >>>
# Minimal Git branch helper (fast, no heavy libs)
__chatgpt_git_branch() {
  local b
  b=$(git rev-parse --abbrev-ref HEAD 2>/dev/null) || return 0
  [ "$b" = "HEAD" ] && b=$(git rev-parse --short HEAD 2>/dev/null)
  [ -n "$b" ] && printf ' (%s)' "$b"
}

# Prompt for root: show exit code when non-zero; root@host in red; cwd in magenta; git branch
__chatgpt_set_ps1_root() {
  local exit="$?"
  local red="\[\e[31m\]" green="\[\e[32m\]" yellow="\[\e[33m\]" blue="\[\e[34m\]" magenta="\[\e[35m\]" cyan="\[\e[36m\]" reset="\[\e[0m\]"
  local code=""
  [ "$exit" -ne 0 ] && code="${red}↯${exit}${reset} "
  PS1="${code}${red}\u${reset}@${blue}\h${reset}:${magenta}\w${reset}\$(__chatgpt_git_branch)${cyan}#${reset}"
}
case "${PROMPT_COMMAND:-}" in
  *"__chatgpt_set_ps1_root"*) : ;;
  "") PROMPT_COMMAND="__chatgpt_set_ps1_root" ;;
  *)  PROMPT_COMMAND="__chatgpt_set_ps1_root; $PROMPT_COMMAND" ;;
esac

# Useful shopt toggles
shopt -s checkwinsize cdspell autocd globstar

# Root-friendly aliases
alias ls='ls --color=auto -F --group-directories-first'
alias ll='ls -lah'
alias la='ls -A'
alias grep='grep --color=auto'
alias df='df -h'
alias free='free -h'
alias watchd='watch -n1 -d'
alias cp='cp -i'
alias mv='mv -i'
alias rm='rm -I'   # prompt once if >3 files
alias ..='cd ..'
alias ...='cd ../..'

# Prefer modern tools if present
command -v batcat >/dev/null 2>&1 && alias bat='batcat'
command -v bat >/dev/null 2>&1     && alias cat='bat --paging=never'
command -v lsd  >/dev/null 2>&1    && alias ls='lsd'

# Defaults
export EDITOR="${EDITOR:-vim}"
export VISUAL="$EDITOR"
export PAGER="${PAGER:-less}"
export LESS='-R'
export GPG_TTY="$(tty 2>/dev/null || true)"

# Security hygiene: stricter umask for root
umask 027
# <<< chatgpt-root: QoL shell block <<<
EOF
)
append_once "$BASHRC" "$QOL_TAG" "$QOL_BLOCK"
msg "Added root QoL prompt/aliases/shopt/editor block."

# Ensure the preferred editor is set explicitly
if grep -Fq 'export EDITOR=' "$BASHRC"; then
  sed -i "0,/export EDITOR=.*/s//export EDITOR=\"${PREFER_EDITOR}\"/" "$BASHRC"
else
  echo "export EDITOR=\"${PREFER_EDITOR}\"" >> "$BASHRC"
fi
msg "Default EDITOR for root set to ${PREFER_EDITOR}"

# =========================
# 6) Readline tweaks for root
# =========================
INPUT_TAG="# >>> chatgpt-root: readline tweaks >>>"
INPUT_BLOCK=$(cat <<'EOF'
# >>> chatgpt-root: readline tweaks >>>
set show-all-if-ambiguous on
set menu-complete-display-prefix on
set colored-stats on
set mark-symlinked-directories on
set completion-ignore-case on
set completion-map-case on
"\e[A": history-search-backward
"\e[B": history-search-forward
# <<< chatgpt-root: readline tweaks <<<
EOF
)
append_once "$INPUTRC" "$INPUT_TAG" "$INPUT_BLOCK"
msg "Improved /root/.inputrc completion/behavior."

# =========================
# Done
# =========================
msg "Root shell setup complete. Open a NEW root shell or run: source /root/.bashrc"