#!/bin/bash
BASHRC_FILE="/etc/bashrc"
ROOT_BASHRC_FILE="/root/.bashrc"
PROFILE_FILE="/etc/profile"
ROOT_BASH_PROFILE="/root/.bash_profile"

read -r -d '' PROMPT_BLOCK << 'EOF'
if [ -n "$STY" ]; then
  screen_info="\[\e[0;1;33m\](screen $(echo "$STY" | cut -d. -f2)) \[\e[0m\]"
else
  screen_info=""
fi
PS1="${screen_info}\[\e[0;1;38;5;46m\]\u\[\e[0m\]@\[\e[0;1;38;5;214m\]\h\[\e[0m\][\[\e[0;1;97m\]\W\[\e[0m\]]\[\e[0;38;5;46m\]\$ \[\e[0m\]"
EOF

read -r -d '' ROOT_PROMPT_BLOCK << 'EOF'
if [ -n "$STY" ]; then
  screen_info="\[\e[0;1;33m\](screen $(echo "$STY" | cut -d. -f2)) \[\e[0m\]"
else
  screen_info=""
fi
PS1="${screen_info}\[\e[0;1;91m\]\u\[\e[0m\]@\[\e[0;1;38;5;214m\]\h\[\e[0m\][\[\e[0;1;97m\]\W\[\e[0m\]]\[\e[0;1;91m\]# \[\e[0m\]"
EOF

append_prompt() {
  local file="$1"
  local block="$2"
  echo "$block" | sudo tee -a "$file" > /dev/null
}

remove_last_PS1_line() {
  local file="$1"
  if [ -f "$file" ]; then
    if tail -n 1 "$file" | grep -q '^PS1='; then
      sudo sed -i '${/^PS1=/d;}' "$file"
    fi
  fi
}

append_prompt "$BASHRC_FILE" "$PROMPT_BLOCK"
append_prompt "$ROOT_BASHRC_FILE" "$ROOT_PROMPT_BLOCK"
remove_last_PS1_line "$PROFILE_FILE"
remove_last_PS1_line "$ROOT_BASH_PROFILE"