#!/bin/bash

# Check if running as root
if [ "$(id -u)" != "0" ]; then
    echo "This script requires root privileges"
    exit 1
fi

# Check for required commands
for cmd in curl tar systemctl; do
    command -v $cmd &> /dev/null || { echo "Required command $cmd is not installed"; exit 1; }
done

# Initialize variables
PROMETHEUS_IP=""
N9E_IP=""
REGION=""

# Parse command-line arguments
while [ $# -gt 0 ]; do
    case "$1" in
        --prometheus-ip=*) PROMETHEUS_IP="${1#*=}"; shift ;;
        --n9e-ip=*) N9E_IP="${1#*=}"; shift ;;
        --region=*) REGION="${1#*=}"; shift ;;
        *) echo "Unknown parameter: $1"; echo "Usage: $0 --prometheus-ip=<ip> --n9e-ip=<ip> [--region=<region>]"; echo "Example: $0 --prometheus-ip=10.15.1.11 --n9e-ip=127.0.0.1"; exit 1 ;;
    esac
done

# Validate required parameters
[ -z "$PROMETHEUS_IP" ] || [ -z "$N9E_IP" ] && {
    echo "Usage: $0 --prometheus-ip=<ip> --n9e-ip=<ip> [--region=<region>]"
    echo "Example: $0 --prometheus-ip=10.15.1.11 --n9e-ip=127.0.0.1"
    exit 1
}

# Construct URLs
PROMETHEUS_URL="http://$PROMETHEUS_IP:9091/api/v1/write"
N9E_HEARTBEAT_URL="http://$N9E_IP:17000/v1/n9e/heartbeat"

# Get region (prefer parameter, fallback to hostname)
if [ -z "$REGION" ]; then
    HOSTNAME=$(hostname)
    REGION=$(echo "$HOSTNAME" | cut -d'-' -f3)
    [ -z "$REGION" ] && { echo "Cannot extract region from hostname, please specify with --region"; exit 1; }
fi

INSTALL_DIR="/opt/categraf"
CONFIG_FILE="$INSTALL_DIR/conf/config.toml"
DOWNLOAD_URL="https://github.com/flashcatcloud/categraf/releases/download/v0.4.9/categraf-v0.4.9-linux-amd64.tar.gz"

# Download and extract
echo "Downloading and installing Categraf..."
mkdir -p "$INSTALL_DIR" || { echo "Failed to create $INSTALL_DIR"; exit 1; }
curl -sL "$DOWNLOAD_URL" | tar -xz -C "$INSTALL_DIR" --strip-components=1 || { echo "Download or extraction failed"; exit 1; }

# Update configuration file
echo "Updating configuration file..."
cat > "$CONFIG_FILE" << EOF
[global]
  print_configs = false
  hostname = ""
  region = "$REGION"

[[writers]]
  url = "$PROMETHEUS_URL"
  batch = 1000
  chan_size = 1000000
  timeout = 5000
  dial_timeout = 2500
  max_idle_conns_per_host = 100

[heartbeat]
  enable = true
  url = "$N9E_HEARTBEAT_URL"
  interval = 10
  timeout = 5000
  dial_timeout = 2500
  max_idle_conns_per_host = 100
EOF
[ $? -ne 0 ] && { echo "Failed to update configuration file"; exit 1; }

# Create systemd service
echo "Creating systemd service..."
cat > /etc/systemd/system/categraf.service << EOF
[Unit]
Description=Categraf Telemetry Collector
After=network.target

[Service]
Type=simple
ExecStart=$INSTALL_DIR/categraf --configs $INSTALL_DIR/conf
Restart=always
User=root
WorkingDirectory=$INSTALL_DIR

[Install]
WantedBy=multi-user.target
EOF
[ $? -ne 0 ] && { echo "Failed to create systemd service"; exit 1; }

# Start service
echo "Starting Categraf service..."
systemctl daemon-reload && systemctl enable categraf.service && systemctl start categraf.service || { echo "Failed to start service"; exit 1; }

# Check service status
systemctl is-active --quiet categraf.service && echo "Categraf installed and running successfully!" || { echo "Service not running, please check logs"; exit 1; }