#!/bin/bash

# --- Functions ---

# Check the status of the last command executed and exit if an error occurred
check_status() {
  if [ $? -ne 0 ]; then
    echo "Error encountered. Exiting."
    exit 1
  fi
}

# Check if a command exists
command_exists() {
  command -v "$1" >/dev/null 2>&1
}

# --- Main Script ---

# Ensure the script is run as root
if [ "$EUID" -ne 0 ]; then
  echo "Error: Please run this script as root."
  exit 1
fi

# Prompt for and confirm the new IP address
read -p "Enter the Primary Node IP address: " NEW_IP
read -p "Confirm the Primary Node IP address: " CONFIRM_IP

if [ "$NEW_IP" != "$CONFIRM_IP" ]; then
  echo "Error: IP addresses do not match. Exiting."
  exit 1
fi
# Prompt for and confirm the Primary Node Hostname
read -p "Enter the Primary Node Hostname: " NEW_HOSTNAME
read -p "Confirm the Primary Node Hostname: " CONFIRM_HOSTNAME

if [ "$NEW_HOSTNAME" != "$CONFIRM_HOSTNAME" ]; then
  echo "Error: Hostnames do not match. Exiting."
  exit 1
fi
# --- OS Version Check ---
echo -e "\n🖥️ Checking OS version..."
OS_VERSION=$(grep 'VERSION_ID' /etc/os-release | cut -d '"' -f2)

if [[ "$OS_VERSION" != "22.04" ]]; then
  echo "❌ OS version mismatch: Ubuntu 22.04 required. Detected: $OS_VERSION"
  exit 1
else
  echo "✅ OS is Ubuntu 22.04"
fi

# --- IP Address Check ---
echo -e "\n🔍 Checking private IP address..."
PRIVATE_IP=$(hostname -I | awk '{print $1}')
echo "Private IP: $PRIVATE_IP"

# --- Ping Check ---
echo -e "\n🌐 Checking internet connectivity..."
PING_HOSTS=("google.com" "yahoo.com" "8.8.8.8" "8.8.4.4")
PING_SUCCESS=false

for host in "${PING_HOSTS[@]}"; do
  if ping -c 2 "$host" > /dev/null; then
    echo "✅ Ping to $host successful."
    PING_SUCCESS=true
    break
  else
    echo "❌ Ping to $host failed."
  fi
done

if [ "$PING_SUCCESS" = false ]; then
  echo "❌ No internet connectivity. Exiting."
  exit 1
fi

# --- RAM & CPU Check ---
echo -e "\n🧠 Checking RAM and CPU..."
TOTAL_RAM=$(free -g | awk '/^Mem:/{print $2}')
CPU_CORES=$(nproc)

echo "Total RAM: ${TOTAL_RAM}GB"
echo "CPU Cores: $CPU_CORES"

if [[ "$TOTAL_RAM" -lt 16 || "$CPU_CORES" -lt 10 ]]; then
  echo "❌ System requirements not met. At least 16GB RAM and 10 CPU cores required."
  exit 1
else
  echo "✅ RAM and CPU requirements met."
fi

# --- GPU Check ---
echo -e "\n🎮 Checking for NVIDIA GPU..."
if ! lspci | grep -i nvidia > /dev/null; then
  echo "❌ NVIDIA GPU not detected."
  #exit 1
else
  echo "✅ NVIDIA GPU detected."
fi

# --- Python3 Check ---
echo -e "\n🐍 Checking Python version..."
if ! command -v python3 > /dev/null; then
  echo "❌ Python3 is not installed. Exiting."
  exit 1
else
  PYTHON_VERSION=$(python3 --version 2>&1)
  echo "✅ Python version: $PYTHON_VERSION"
fi


# Update /etc/hosts with the new IP and Hostname
echo "Updating /etc/hosts with IP: $NEW_IP and Hostname: $NEW_HOSTNAME"

# Check if the IP-Hostname entry already exists in /etc/hosts
if grep -q "$NEW_IP" /etc/hosts; then
  echo "Entry for IP $NEW_IP already exists in /etc/hosts. Updating entry."
  # Replace the existing entry
  sed -i "s/^.*$NEW_IP.*/$NEW_IP\t$NEW_HOSTNAME/" /etc/hosts
else
  # Add the new IP-Hostname entry to /etc/hosts
  echo "$NEW_IP    $NEW_HOSTNAME" >> /etc/hosts
fi

check_status

# Verify the update
echo "Updated /etc/hosts successfully!"
cat /etc/hosts | grep "$NEW_IP"  # Show the updated entry


echo "Setting up Qubrid AI Controller with IP address: $NEW_IP"
sudo apt update
# --- Uninstall Existing Qubrid Packages ---
echo "Uninstalling any existing Qubrid packages..."
dpkg -l | grep qubrid
if dpkg -l | grep -q "ii  qubrid"; then
  echo "Removing existing Qubrid packages..."
  sudo dpkg --remove qubriddcai-worker-node
 # sudo dpkg --remove qubridaidcpackageinstall-primary-node 
  sudo dpkg --remove qubridaidccontroller-worker-node
  check_status
fi

# --- Uninstall Node.js and npm ---
#echo "Uninstalling Node.js and npm..."
#sudo apt update
#sudo apt remove nodejs npm -y
#check_status

# --- Uninstall MongoDB ---
#echo "Uninstalling MongoDB..."
# sudo systemctl stop mongod || true
# sudo apt-get purge mongodb-org* -y
# sudo rm -rf /var/log/mongodb
# sudo rm -rf /var/lib/mongodb
#check_status

# --- Update Package Lists and Clean Up ---
echo "Updating package lists and cleaning up..."
sudo apt update
sudo apt autoremove -y
sudo apt clean
check_status

# --- Install Dependencies ---
echo "Installing essential dependencies (ca-certificates, curl)..."
sudo apt-get install -y ca-certificates curl
check_status

# --- Install Docker ---
if ! dpkg-query -W -f='${Status}' docker-ce 2>/dev/null | grep -q "install ok installed"; then
  echo "Docker is not installed. Installing..."
  
  # Add Docker GPG key and repository
  sudo install -m 0755 -d /etc/apt/keyrings
  check_status

  sudo curl -fsSL https://download.docker.com/linux/ubuntu/gpg -o /etc/apt/keyrings/docker.asc
  check_status

  sudo chmod a+r /etc/apt/keyrings/docker.asc
  check_status

  echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.asc] https://download.docker.com/linux/ubuntu $(. /etc/os-release && echo "$VERSION_CODENAME") stable" | sudo tee /etc/apt/sources.list.d/docker.list > /dev/null
  check_status

  sudo apt-get update
  check_status

  sudo apt-get install -y docker-ce=5:26.1.0-1~ubuntu.22.04~jammy docker-ce-cli=5:26.1.0-1~ubuntu.22.04~jammy containerd.io docker-buildx-plugin docker-compose-plugin
  check_status

  echo "Docker installation completed."
else
  echo "Docker is already installed."
fi

# --- Install Nvidia Container Toolkit ---
if ! dpkg-query -W -f='${Status}' nvidia-container-toolkit 2>/dev/null | grep -q "install ok installed"; then
  echo "Nvidia container toolkit is not installed. Installing..."
  
  # Add Nvidia container toolkit GPG key and repository
  curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
  check_status

  curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
  check_status

  sudo apt-get update
  check_status

  sudo apt-get install -y nvidia-container-toolkit
  check_status

  sudo systemctl restart docker
  check_status

  echo "Nvidia container toolkit installation completed."
else
  echo "Nvidia container toolkit is already installed."
fi

# --- Install Python 3 pip ---
echo "Installing Python 3 pip..."
sudo apt-get update
sudo apt-get install -y python3-pip
check_status

# --- Install MongoDB ---
if ! dpkg-query -W -f='${Status}' mongodb-org 2>/dev/null | grep -q "install ok installed"; then
  echo "MongoDB is not installed. Installing..."
  
  # Install gnupg and curl
  sudo apt-get install -y gnupg curl
  check_status
  
  # Add MongoDB GPG key and repository
  curl -fsSL https://www.mongodb.org/static/pgp/server-7.0.asc | sudo gpg -o /usr/share/keyrings/mongodb-server-7.0.gpg --dearmor
  check_status

  echo "Adding MongoDB repository..."
  echo "deb [ arch=amd64,arm64 signed-by=/usr/share/keyrings/mongodb-server-7.0.gpg ] https://repo.mongodb.org/apt/ubuntu jammy/mongodb-org/7.0 multiverse" | sudo tee /etc/apt/sources.list.d/mongodb-org-7.0.list
  check_status

  sudo apt-get update
  check_status

  sudo apt-get install -y mongodb-org
  check_status

  echo "MongoDB installation completed."
else
  echo "MongoDB is already installed."
fi

# --- Install Python Packages ---
echo "Installing Python Flask and related packages..."
pip3 install Flask==3.0.3 flask_cors==4.0.1 botocore==1.34.120 pymongo==4.7.2 ping3==4.0.8 apscheduler==3.10.4 docker==7.1.0 Flask-Compress==1.17 huggingface-hub==0.29.3 tornado==6.5 --no-cache-dir

####ubuntu24.04 install ###################
#pip3 install Flask==3.0.3 flask_cors==4.0.1 botocore==1.34.120 pymongo==4.7.2 ping3==4.0.8 apscheduler==3.10.4 docker==7.1.0 Flask-Compress==1.17 huggingface-hub==0.29.3 tornado --no-cache-dir  --break-system-packages
check_status

# --- Install Node.js and npm using NVM ---
#if ! command_exists nvm; then
 # echo "Installing nvm (Node Version Manager)..."
 # curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.2/install.sh | bash
 # export NVM_DIR="$([ -s "$HOME/.nvm/nvm.sh" ] && echo "$HOME/.nvm" || /opt/nvm)"
 # [ -s "$NVM_DIR/nvm.sh" ] && \. "$NVM_DIR/nvm.sh"  # This loads nvm in the current shell
 # [ -s "$NVM_DIR/bash_completion" ] && \. "$NVM_DIR/bash_completion"  # This loads nvm bash_completion
#else
 # echo "nvm is already installed."
#fi

#if command_exists nvm; then
 # if ! command_exists node || [ "$(node -v)" != "v18.20.7" ]; then
  #  echo "Installing Node.js v18.20.7..."
   # nvm install 18.20.7
   # nvm use 18.20.7
 # else
  #  echo "Node.js v18.20.7 is already installed."
 # fi
  
 # echo "Verifying Node.js version:"
 # node -v

 # echo "Current nvm version:"
 # nvm current

#  echo "Installing npm serve globally..."
 # npm install --global serve
 # check_status
#else
 # echo "Error: nvm installation failed or was not detected. Node.js and npm may not be available."
#fi
#export NVM_DIR="$HOME/.nvm"
#[ -s "$NVM_DIR/nvm.sh" ] && \. "$NVM_DIR/nvm.sh"  # This loads nvm
#[ -s "$NVM_DIR/bash_completion" ] && \. "$NVM_DIR/bash_completion"  # This loads nvm bash_completion
# --- Check npm, node, and serve installations ---
#echo "Checking installation locations:"
#which npm
#which node
#which serve

# --- Configure MongoDB for Remote Access and Replication ---
#echo "Configuring MongoDB for remote access and replication..."

# Modify /etc/mongod.conf file without -e flag
#echo "Modifying /etc/mongod.conf..."
#sudo cp /etc/mongod.conf /etc/mongod.conf.bak
#sudo sed -i 's/^  bindIp: 127.0.0.1/  bindIp: 0.0.0.0/' /etc/mongod.conf
#sudo sed -i '/^replication:/,/^$/d' /etc/mongod.conf
#echo -e"\nreplication:\n  replSetName: \"rs0\"" | sudo tee -a /etc/mongod.conf
#sleep 5
#echo -e "\nreplication:\n  replSetName: \"rs0\"" | sudo tee -a /etc/mongod.conf > /dev/null

#check_status

# Restart MongoDB service
echo "Restarting MongoDB service..."
sudo systemctl restart mongod
check_status
#echo "Waiting for MongoDB to restart (20 seconds)..."
sleep 10

# Initialize MongoDB Replica Set
#echo "Initializing MongoDB Replica Set..."
#if command_exists mongosh; then
 # mongosh --host "$NEW_IP" --port 27017 <<EOF
#rs.initiate()
#rs.status()
#exit
#EOF
 # check_status
#else
 # echo "Error: mongosh is not installed. Please install it to initialize the replica set manually."
#fi

# Create /etc/ai_control.conf
echo "Creating /etc/ai_control.conf..."
echo -e "[DEFAULT]\nMONGODB_IP=$NEW_IP" | sudo tee /etc/ai_control.conf > /dev/null
check_status

# --- Enable and Start MongoDB Service ---
echo "Enabling and starting MongoDB service..."
sudo systemctl enable mongod
sudo systemctl daemon-reload
sudo systemctl start mongod
check_status


# --- Install Qubrid AI Controller ---
echo "Installing qubridai_dc_controller..."
wget https://qubridaic-packages.s3.us-east-1.amazonaws.com/v1.2.0-trail/ubuntu22.04-trial/qubridai_dc_controller-worker-node-v1.2.0.deb -O /tmp/qubridai_dc_controller.deb
sudo dpkg -i /tmp/qubridai_dc_controller.deb
check_status
rm -f /tmp/qubridai_dc_controller.deb
#python3 /usr/local/bin/qubriddcb/logos/delete_values.py
#########after install front and backend############################# 
#cd /usr/local/bin/qubriddcb/logos
#python3 /usr/local/bin/qubriddcb/logos/NIM_id.py
#python3 /usr/local/bin/qubriddcb/logos/images_send.py
#python3 /usr/local/bin/qubriddcb/logos/model_id.py
python3 /usr/local/bin/qubriddcb/logos/library_send.py
# Update API base URL
#FILE_PATH="/var/www/html/qubriddcf/build/config.json"
#echo "Updating API base URL in '$FILE_PATH' to '$NEW_IP' using sed..."
#sudo sed -i "s/\"API_BASE_URL\": \".*\"/\"API_BASE_URL\": \"$NEW_IP\"/" "$FILE_PATH"
#check_status
sleep 2
##############CHECKS ##################

# --- Firewall Check ---
echo -e "\n🧱 Checking firewall status..."
if command -v ufw > /dev/null; then
  FIREWALL_STATUS=$(ufw status)
  if echo "$FIREWALL_STATUS" | grep -q "Status: active"; then
    echo "❌ Firewall is active. Please disable it before proceeding."
    exit 1
  else
    echo "✅ Firewall is disabled."
  fi
else
  echo "⚠️ ufw not found. Skipping firewall check."
fi

# --- Port Checks ---
echo -e "\n🔌 Checking required ports..."

IS_PORT_OPEN=0
for PORT in {4000..5000}; do
  if ss -ltnH | awk '{print $4}' | grep -q ":$PORT$"; then
    echo "✅ Port $PORT is open."
    IS_PORT_OPEN=1
    break
  fi
done

if [[ $IS_PORT_OPEN -eq 0 ]]; then
  echo "❌ No ports open in the 4000-5000 range. Exiting."
  exit 1
fi

if ! ss -ltn | grep -q ":27017"; then
  echo "❌ MongoDB port 27017 is not open. Exiting."
  exit 1
else
  echo "✅ Port 27017 is open."
fi


rm -rf qubridAIcontroller-primary-node.deb
# Final steps
echo "✅ Qubrid AI Controller setup complete!"
echo "🎉 All tasks completed successfully."
rm -rf /usr/local/bin/qubridAIcontroller-install

