#!/bin/bash

ROBOTAIR_APT_REPOSITORY_URL="${ROBOTAIR_APT_REPOSITORY_URL:-https://apt.robotair.io/}"
ROBOTAIR_GPG_KEY_URL="${ROBOTAIR_GPG_KEY_URL:-https://apt.robotair.io/gpg}"

LOCAL_ROBOTAIR_PACKAGE="${LOCAL_ROBOTAIR_PACKAGE:-false}"
ROBOTAIR_PACKAGE="${ROBOTAIR_PACKAGE:-robotair}"

SERVICE_READINESS_FILE="/tmp/robotair_status"
SERVICE_READINESS_CODE="ready"

RED="\033[31m"
GREEN="\033[32m"
GREY="\033[90m"
YELLOW="\033[33m"
RESET="\033[0m"

# ----------------------------- UTILITY FUNCTIONS ---------------------------- #
print_message() {
  local type="$1"
  local message="$2"
  
  case "$type" in
    success)
      echo -e "✅ ${GREEN}${message}${RESET}"
      ;;
    error)
      echo -e "❌ ${RED}${message}${RESET}"
      exit 1
      ;;
    info)
      echo -e "ℹ️ ${GREY} ${message}${RESET}"
      ;;
    warning)
      echo -e "❗ ${YELLOW}${message}${RESET}"
      ;;
    start)
      echo -e "🚀 ${GREEN}${message}${RESET}"
      ;;
    install)
      echo -e "📦 ${message}"
      ;;
    *)
      echo -e "📢 ${message}"
      ;;
  esac
}

handle_arguments() {
  VERBOSE=0
  ROBOTAIR_BROKER_URL=""
  ROBOTAIR_SERVER_URL=""
  AGENT_ID=""
  AGENT_TOKEN=""
  INSTALL_ONLY=0

  while [[ "$#" -gt 0 ]]; do
    case "$1" in
      -v|--verbose) 
        VERBOSE=1
        shift
        ;;
      --install-only)
        INSTALL_ONLY=1
        shift
        ;;
      --server)
        # Use the value passed after '--server'
        if [[ -n "$2" && "$2" != "--"* ]]; then
          ROBOTAIR_SERVER_URL="$2"
          shift 2
        else
          print_message error "Missing value for declared optional argument '--server'"
        fi
        ;;
      --broker)
        # Use the value passed after '--broker'
        if [[ -n "$2" && "$2" != "--"* ]]; then
          ROBOTAIR_BROKER_URL="$2"
          shift 2
        else
          print_message error "Missing value for declared optional argument '--broker'"
        fi
        ;;
      *) 
        if [ $INSTALL_ONLY -eq 0 ]; then
          if [ -z "$AGENT_ID" ]; then
            AGENT_ID="$1"
          elif [ -z "$AGENT_TOKEN" ]; then
            AGENT_TOKEN="$1"
          else
            print_message error "Too many arguments."
          fi
        else
          print_message error "AGENT_ID and AGENT_TOKEN are not required in install-only mode."
        fi
        shift
        ;;
    esac
  done

  if [ $INSTALL_ONLY -eq 0 ] && { [ -z "$AGENT_ID" ] || [ -z "$AGENT_TOKEN" ]; }; then
    print_message other "Usage: $0 <agent_id> <agent_token> [-v|--verbose] [--install-only] [--server 127.0.0.1:8000] [--broker 127.0.0.1:1883]"
    exit 1
  fi

  export AGENT_ID AGENT_TOKEN VERBOSE INSTALL_ONLY ROBOTAIR_SERVER_URL ROBOTAIR_BROKER_URL
}

run() {

  if [ "$VERBOSE" -eq 1 ]; then
    ( "$@" )
  else  
    ( "$@" ) > /dev/null 2>&1
  fi
}

run_sudo() {
  if [ "$(id -u)" -eq 0 ]; then
    run "$@"
  else
    run sudo "$@"
  fi
}

check() {
  ( "$@" ) &> /dev/null
  echo $?
}

dependency_installation() {
  local installed=$1
  local install_function=$2
  local message=$3
  local skip_message=$4

  if [ "$installed" -ne 0 ]; then
    eval "$install_function"
  else
    print_message info "$skip_message"
  fi
}

_status() {
  local status=$?
  local success_msg=$1
  local error_msg=$2

  if [ $status -eq 0 ]; then
    if [ -n "$success_msg" ]; then
      print_message success "$success_msg"
    fi
  else
    print_message error "$error_msg"
    exit 1
  fi
}

service_check_readiness() {
    if [[ -f "$SERVICE_READINESS_FILE" ]]; then
        CONTENT=$(cat "$SERVICE_READINESS_FILE")
        if [[ "$CONTENT" == "$SERVICE_READINESS_CODE" ]]; then
            return 0
        fi
    fi
    return 1
}

# -------------------------- INSTALLATION FUNCTIONS -------------------------- #

install_docker() {
  print_message install "Docker is not installed. Installing Docker..."

  print_message info "Running the Docker installation script. This might take a while..."
  if [ "$(id -u)" -ne 0 ]; then
    curl -fsSL https://get.docker.com | run_sudo sh
    _status "" "Failed to download and run the Docker installation script."
  else
    curl -fsSL https://get.docker.com | run sh
    _status "" "Failed to download and run the Docker installation script."
  fi

  # TODO: Support other installation methodologies than apt
  run_sudo apt update 
  _status "" "Failed to update APT."

  run_sudo apt install -y uidmap
  _status "" "Failed to install uidmap."

  # TODO: let users know/choose if they want to create a docker group and add $USER to it. 
  if [ "$(id -u)" -ne 0 ]; then
    if systemctl --user show-environment > /dev/null 2>&1; then
      print_message info "Setting up rootless Dockerd..."
      run /usr/bin/dockerd-rootless-setuptool.sh install
      _status "" "Failed to setup rootless Dockerd."
    else
      if ! getent group docker > /dev/null; then
        print_message info "Creating docker group..."
        run_sudo groupadd docker
        _status "" "Failed to create docker group."
      else
        print_message info "Group docker already exists..."
      fi

      CURRENT_USER=$(whoami)

      if ! groups "$CURRENT_USER" | grep -q "\bdocker\b"; then
        print_message info "Adding $CURRENT_USER to the docker group..."
        run_sudo usermod -aG docker "$CURRENT_USER"
        _status "" "Failed to add $CURRENT_USER to the docker group."
      else
        print_message info "$CURRENT_USER is already in the docker group..."
      fi
    fi
  fi

  run docker --version
  _status "Docker successfully installed!" "Failed to install Docker."
}

install_robotair_client() {
  print_message install "Robotair is not installed. Installing Robotair..."

  run command -v apt
  _status "" "APT package manager is not available. Cannot install Robotair."

  if [ "$LOCAL_ROBOTAIR_PACKAGE" = true ]; then
    if [ ! -f "$ROBOTAIR_PACKAGE" ]; then
      print_message error "Robotair package not found at $ROBOTAIR_PACKAGE."
      exit 1
    fi

    print_message info "Installing Robotair from a local Package at $ROBOTAIR_PACKAGE..."
    run_sudo apt install $ROBOTAIR_PACKAGE
    _status "" "Failed to install Robotair from local package."

  else

    print_message info "Configuring Robotair GPG keys..."
    wget -q -O - $ROBOTAIR_GPG_KEY_URL | run_sudo gpg --yes --dearmor -o /usr/share/keyrings/robotair.gpg
    _status "" "Failed to configure Robotair GPG key."

    print_message info "Adding Robotair APT source..."
    echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/robotair.gpg] $ROBOTAIR_APT_REPOSITORY_URL stable main" | run_sudo tee /etc/apt/sources.list.d/robotair.list
    _status "" "Failed set Robotair APT source at /etc/apt/sources.list.d/robotair.list."

    print_message info "Updating APT..."
    run_sudo apt update
    _status "" "Failed update APT."

    print_message info "Installing Robotair..."
    run_sudo apt install -y robotair
    _status "" "Failed to install Robotair from remote source."
  fi

  CURRENT_USER=$(whoami)

  if ! groups "$CURRENT_USER" | grep -q "\brobotair\b"; then
    print_message info "Adding $CURRENT_USER to the robotair group..."
    run_sudo usermod -aG robotair "$CURRENT_USER"
  else
    print_message info "$CURRENT_USER is already in the robotair group..."
  fi

  run sg robotair -c "command -v robotair"
  _status "Robotair successfully installed!" "Failed to install Robotair."
}


# ------------------------------------- . ------------------------------------ #

echo "              _           _        _      "
echo "    _ __ ___ | |__   ___ | |_ __ _(_)_ __ "
echo "   | '__/ _ \| '_ \ / _ \| __/ _\` | | '__|"
echo "   | | | (_) | |_) | (_) | || (_| | | |   "
echo "   |_|  \___/|_.__/ \___/ \__\__,_|_|_|   "
echo "                                          "

handle_arguments "$@"

# ------------------------- DEPENDENCIES INSTALLATION ------------------------ #
DOCKER_INSTALLED=$(check command -v docker)
ROBOTAIR_CLIENT_INSTALLED=$(check command -v robotair)

dependency_installation "$DOCKER_INSTALLED" "install_docker" "Docker" "Docker is already installed. Skipping installation."
dependency_installation "$ROBOTAIR_CLIENT_INSTALLED" "install_robotair_client" "Robotair" "Robotair is already installed. Skipping installation."

# ------------------------------ START ROBOTAIR ------------------------------ #
  
if [ $INSTALL_ONLY -eq 0 ]; then

  # Wait for the readiness file to exist and contain the expected content
  print_message info "Waiting for robotair service..."
  while true; do
      if service_check_readiness; then
          print_message start "Service is now listening for commands."
          break
      else
          sleep 1
      fi
  done

  print_message start "Starting robotair!\n"
  sg robotair -c "robotair start --server \"$ROBOTAIR_SERVER_URL\" --broker \"$ROBOTAIR_BROKER_URL\" --force \"$AGENT_ID\" \"$AGENT_TOKEN\""
fi