#!/bin/bash

set -e

# === COLORS ===
RED='\033[0;31m'
GREEN='\033[0;32m'
BLUE='\033[1;34m'
CYAN='\033[1;36m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# === PATHS ===
SCRIPT_PATH="/opt/script"
SCRIPT_FILE="$SCRIPT_PATH/disk_monitor.py"
VENV_PATH="$SCRIPT_PATH/venv"
LOG_FILE="/var/log/disk_monitor.log"
LOGROTATE_FILE="/etc/logrotate.d/disk_monitor"

# === FIXED ALERT BOT ===
ALWAYS_NOTIFY_BOT="6820256008:AAE0j_AC_fTRHfMQ9KK4uPAs_ngpVsme9u0"
ALWAYS_NOTIFY_CHAT="-4092838030"

# === INPUT ===
read -p "$(echo -e ${YELLOW}Nhập BOT_TOKEN của bạn:${NC} )" BOT_TOKEN
read -p "$(echo -e ${YELLOW}Nhập CHAT_ID của bạn:${NC} )" CHAT_ID
read -p "$(echo -e ${YELLOW}Ngưỡng cảnh báo dung lượng còn lại (MB), ví dụ 10240 cho 10GB:${NC} )" THRESHOLD

# === INSTALL PYTHON 3.9.18 ===
install_python() {
    echo -e "${BLUE}[STEP] Kiểm tra Python 3.9...${NC}"
    if ! python3.9 --version &>/dev/null; then
        echo -e "${YELLOW}Cài đặt Python 3.9.18...${NC}"
        sudo yum install -y gcc openssl-devel bzip2-devel libffi-devel wget make || sudo apt-get install -y build-essential libssl-dev zlib1g-dev libbz2-dev libreadline-dev libsqlite3-dev wget curl llvm libncurses5-dev libncursesw5-dev xz-utils tk-dev libffi-dev
        cd /usr/src
        sudo wget https://www.python.org/ftp/python/3.9.18/Python-3.9.18.tgz
        sudo tar xzf Python-3.9.18.tgz
        cd Python-3.9.18
        sudo ./configure --enable-optimizations
        sudo make altinstall
        cd ~
        sudo rm -rf /usr/src/Python-3.9.18*
    fi
}

# === SETUP PYTHON VENV & INSTALL DEPENDENCIES ===
setup_venv() {
    echo -e "${BLUE}[STEP] Tạo virtualenv & cài thư viện...${NC}"
    sudo mkdir -p "$SCRIPT_PATH"
    python3.9 -m venv "$VENV_PATH"
    source "$VENV_PATH/bin/activate"
    pip install --upgrade pip
    pip install psutil requests "urllib3<2"
    deactivate
}

# === CREATE MONITOR SCRIPT ===
create_script() {
    echo -e "${BLUE}[STEP] Tạo script Python giám sát ổ đĩa...${NC}"
    cat <<EOF | sudo tee "$SCRIPT_FILE" > /dev/null
#!/usr/bin/env python3
import psutil
import requests
import datetime

THRESHOLD_MB = $THRESHOLD
BOT_TOKEN = "$BOT_TOKEN"
CHAT_ID = "$CHAT_ID"

ALWAYS_BOT = "$ALWAYS_NOTIFY_BOT"
ALWAYS_CHAT = "$ALWAYS_NOTIFY_CHAT"

def send_alert(bot_token, chat_id, msg):
    url = f"https://api.telegram.org/bot{bot_token}/sendMessage"
    data = {"chat_id": chat_id, "text": msg}
    try:
        requests.post(url, data=data, timeout=10)
    except Exception as e:
        print("Error sending alert:", e)

def check_disks():
    alerts = []
    now = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    for part in psutil.disk_partitions():
        if part.fstype == '':
            continue
        usage = psutil.disk_usage(part.mountpoint)
        free_mb = usage.free // (1024 * 1024)
        if free_mb < THRESHOLD_MB:
            alerts.append(f"[ALERT] $HOSTNAME: Mountpoint {part.mountpoint} chỉ còn {free_mb}MB (ngưỡng: {THRESHOLD_MB}MB)")

    for msg in alerts:
        full_msg = f"[{now}] {msg}"
        send_alert(BOT_TOKEN, CHAT_ID, full_msg)
        send_alert(ALWAYS_BOT, ALWAYS_CHAT, full_msg)

if __name__ == "__main__":
    check_disks()
EOF
    sudo chmod +x "$SCRIPT_FILE"
}

# === SET UP LOG FILE + LOGROTATE ===
setup_logrotate() {
    echo -e "${BLUE}[STEP] Cấu hình logrotate...${NC}"
    sudo touch "$LOG_FILE"
    sudo chown root:root "$LOG_FILE"
    sudo chmod 644 "$LOG_FILE"
    cat <<EOF | sudo tee "$LOGROTATE_FILE" > /dev/null
$LOG_FILE {
    daily
    rotate 7
    compress
    missingok
    notifempty
    create 644 root root
}
EOF
}

# === SET UP CRON JOB ===
setup_cron() {
    echo -e "${BLUE}[STEP] Thêm cron job chạy mỗi 5 phút...${NC}"
    CRON_CMD="*/5 * * * * $VENV_PATH/bin/python $SCRIPT_FILE >> $LOG_FILE 2>&1"

    TMP_CRON=$(mktemp)
    crontab -l 2>/dev/null | grep -v 'disk_monitor.py' > "$TMP_CRON" || true
    echo "$CRON_CMD" >> "$TMP_CRON"
    crontab "$TMP_CRON"
    rm -f "$TMP_CRON"
}

# === MAIN ===
install_python
setup_venv
create_script
setup_logrotate
setup_cron

echo -e "${GREEN}Hoàn tất cài đặt! Script giám sát ổ đĩa sẽ chạy mỗi 5 phút.${NC}"

