#!/bin/bash

# Universal Promtail Installation and Configuration Script
# Supports: Debian, Ubuntu, CentOS, RHEL, Fedora, Arch Linux, Alpine, Raspberry Pi OS
# Usage: bash -c "$(curl -fsSL <your-gist-url>/install-promtail.sh)"

set -e

# Configuration
LOKI_ENDPOINT="loki.pfotenballen.de"
LOKI_PORT="3100"
PROMTAIL_VERSION="2.9.2"
PROMTAIL_USER="promtail"
PROMTAIL_DIR="/opt/promtail"
CONFIG_DIR="/etc/promtail"
LOG_DIR="/var/log/promtail"

# OS Detection Variables
OS=""
DIST=""
PACKAGE_MANAGER=""
SERVICE_MANAGER=""

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Logging functions
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Detect OS and Distribution
detect_os() {
    log_info "Detecting operating system..."
    
    if [[ -f /etc/os-release ]]; then
        . /etc/os-release
        OS=$ID
        DIST=$VERSION_ID
    elif command -v lsb_release >/dev/null 2>&1; then
        OS=$(lsb_release -si | tr '[:upper:]' '[:lower:]')
        DIST=$(lsb_release -sr)
    elif [[ -f /etc/redhat-release ]]; then
        OS="rhel"
        DIST=$(cat /etc/redhat-release | sed 's/.*release \([0-9]\).*/\1/')
    else
        log_error "Cannot detect operating system"
        exit 1
    fi
    
    # Determine package manager
    if command -v apt-get >/dev/null 2>&1; then
        PACKAGE_MANAGER="apt"
    elif command -v yum >/dev/null 2>&1; then
        PACKAGE_MANAGER="yum"
    elif command -v dnf >/dev/null 2>&1; then
        PACKAGE_MANAGER="dnf"
    elif command -v pacman >/dev/null 2>&1; then
        PACKAGE_MANAGER="pacman"
    elif command -v apk >/dev/null 2>&1; then
        PACKAGE_MANAGER="apk"
    else
        log_error "No supported package manager found"
        exit 1
    fi
    
    # Determine service manager
    if command -v systemctl >/dev/null 2>&1 && systemctl --version >/dev/null 2>&1; then
        SERVICE_MANAGER="systemd"
    elif command -v service >/dev/null 2>&1; then
        SERVICE_MANAGER="sysv"
    elif command -v rc-service >/dev/null 2>&1; then
        SERVICE_MANAGER="openrc"
    else
        log_error "No supported service manager found"
        exit 1
    fi
    
    log_info "Detected: OS=$OS, Package Manager=$PACKAGE_MANAGER, Service Manager=$SERVICE_MANAGER"
}

# Check if running as root
check_root() {
    if [[ $EUID -ne 0 ]]; then
        log_error "This script must be run as root"
        exit 1
    fi
}

# Check if promtail is already installed
check_promtail_installed() {
    if command -v promtail &> /dev/null || [[ -f "/usr/local/bin/promtail" ]] || [[ -f "$PROMTAIL_DIR/promtail" ]]; then
        return 0
    else
        return 1
    fi
}

# Test Loki endpoint connectivity
test_loki_connectivity() {
    log_info "Testing connectivity to Loki endpoint: $LOKI_ENDPOINT:$LOKI_PORT"
    
    if command -v nc >/dev/null 2>&1; then
        # Use netcat if available
        if timeout 10 nc -z "$LOKI_ENDPOINT" "$LOKI_PORT" 2>/dev/null; then
            log_success "Successfully connected to $LOKI_ENDPOINT:$LOKI_PORT"
            return 0
        fi
    elif command -v telnet >/dev/null 2>&1; then
        # Use telnet as fallback
        if timeout 10 bash -c "echo 'quit' | telnet $LOKI_ENDPOINT $LOKI_PORT" 2>/dev/null | grep -q "Connected"; then
            log_success "Successfully connected to $LOKI_ENDPOINT:$LOKI_PORT"
            return 0
        fi
    elif command -v curl >/dev/null 2>&1; then
        # Use curl as last resort
        if timeout 10 curl -s "http://$LOKI_ENDPOINT:$LOKI_PORT/ready" >/dev/null 2>&1; then
            log_success "Successfully connected to $LOKI_ENDPOINT:$LOKI_PORT"
            return 0
        fi
    fi
    
    log_error "Cannot reach $LOKI_ENDPOINT:$LOKI_PORT"
    log_error "Please check your network connection and Loki server status"
    return 1
}

# Install dependencies
install_dependencies() {
    log_info "Installing dependencies..."
    
    case $PACKAGE_MANAGER in
        "apt")
            export DEBIAN_FRONTEND=noninteractive
            apt-get update -qq
            apt-get install -y wget curl unzip netcat-openbsd || apt-get install -y wget curl unzip netcat
            ;;
        "yum")
            yum install -y wget curl unzip nc
            ;;
        "dnf")
            dnf install -y wget curl unzip nc
            ;;
        "pacman")
            pacman -Sy --noconfirm wget curl unzip netcat
            ;;
        "apk")
            apk update
            apk add wget curl unzip netcat-openbsd
            ;;
        *)
            log_error "Unsupported package manager: $PACKAGE_MANAGER"
            exit 1
            ;;
    esac
    
    log_success "Dependencies installed"
}

# Create promtail user
create_promtail_user() {
    if ! id "$PROMTAIL_USER" &>/dev/null; then
        log_info "Creating promtail user..."
        
        case $OS in
            "alpine")
                adduser -S -D -H -s /bin/false $PROMTAIL_USER
                ;;
            *)
                if command -v useradd >/dev/null 2>&1; then
                    useradd --system --no-create-home --shell /bin/false $PROMTAIL_USER 2>/dev/null || \
                    useradd -r -M -s /bin/false $PROMTAIL_USER
                else
                    log_error "Cannot create user - useradd not available"
                    exit 1
                fi
                ;;
        esac
        
        log_success "Promtail user created"
    else
        log_info "Promtail user already exists"
    fi
}

# Download and install promtail
install_promtail() {
    log_info "Downloading Promtail v$PROMTAIL_VERSION..."
    
    # Determine architecture
    ARCH=$(uname -m)
    case $ARCH in
        x86_64)
            ARCH_SUFFIX="amd64"
            ;;
        aarch64)
            ARCH_SUFFIX="arm64"
            ;;
        armv7l)
            ARCH_SUFFIX="arm"
            ;;
        arm*)
            ARCH_SUFFIX="arm"
            ;;
        *)
            log_error "Unsupported architecture: $ARCH"
            exit 1
            ;;
    esac
    
    # Create directories
    mkdir -p $PROMTAIL_DIR
    mkdir -p $CONFIG_DIR
    mkdir -p $LOG_DIR
    
    # Download promtail binary
    DOWNLOAD_URL="https://github.com/grafana/loki/releases/download/v$PROMTAIL_VERSION/promtail-linux-$ARCH_SUFFIX.zip"
    
    cd /tmp
    log_info "Downloading from: $DOWNLOAD_URL"
    if ! wget --timeout=30 --tries=3 -q "$DOWNLOAD_URL" -O promtail.zip; then
        log_error "Failed to download Promtail. Check internet connection."
        exit 1
    fi
    
    if ! unzip -q promtail.zip; then
        log_error "Failed to extract Promtail archive"
        exit 1
    fi
    
    # Install binary
    chmod +x promtail-linux-$ARCH_SUFFIX
    mv promtail-linux-$ARCH_SUFFIX /usr/local/bin/promtail
    
    # Set ownership
    chown root:root /usr/local/bin/promtail
    chown -R $PROMTAIL_USER:$PROMTAIL_USER $CONFIG_DIR $LOG_DIR
    
    # Cleanup
    rm -f promtail.zip
    
    log_success "Promtail installed successfully"
}

# Create promtail configuration
create_config() {
    log_info "Creating Promtail configuration..."
    
    # Get the actual hostname
    HOSTNAME=$(hostname)
    
    cat > $CONFIG_DIR/promtail.yml << EOF
server:
  http_listen_port: 9080
  grpc_listen_port: 0

positions:
  filename: /var/lib/promtail/positions.yaml

clients:
  - url: http://$LOKI_ENDPOINT:$LOKI_PORT/loki/api/v1/push

scrape_configs:
  # Direct /var/log/ files (system logs)
  - job_name: system-logs
    static_configs:
      - targets:
          - localhost
        labels:
          job: system-logs
          service: system
          host: $HOSTNAME
          __path__: /var/log/*.log

  # Service-specific logs in subdirectories
  - job_name: service-logs
    static_configs:
      - targets:
          - localhost
        labels:
          job: service-logs
          host: $HOSTNAME
          __path__: /var/log/*/*.log
    pipeline_stages:
      - regex:
          expression: '/var/log/(?P<service>[^/]+)/.*'
      - labels:
          service: '{{ .service }}'

  # Recursively capture all nested logs (deeper than one level)
  - job_name: nested-service-logs
    static_configs:
      - targets:
          - localhost
        labels:
          job: nested-service-logs
          host: $HOSTNAME
          __path__: /var/log/**/*.log
    pipeline_stages:
      - regex:
          expression: '/var/log/(?P<service>[^/]+)/.*'
      - labels:
          service: '{{ .service }}'
EOF

    # Create positions directory
    mkdir -p /var/lib/promtail
    chown $PROMTAIL_USER:$PROMTAIL_USER /var/lib/promtail
    
    # Set proper permissions
    chown $PROMTAIL_USER:$PROMTAIL_USER $CONFIG_DIR/promtail.yml
    chmod 640 $CONFIG_DIR/promtail.yml
    
    log_success "Configuration created"
}

# Create service
create_service() {
    case $SERVICE_MANAGER in
        "systemd")
            create_systemd_service
            ;;
        "sysv")
            create_sysv_service
            ;;
        "openrc")
            create_openrc_service
            ;;
        *)
            log_error "Unsupported service manager: $SERVICE_MANAGER"
            exit 1
            ;;
    esac
}

# Create systemd service
create_systemd_service() {
    log_info "Creating systemd service..."
    
    cat > /etc/systemd/system/promtail.service << EOF
[Unit]
Description=Promtail service
Documentation=https://grafana.com/docs/loki/latest/clients/promtail/
After=network.target

[Service]
Type=simple
User=$PROMTAIL_USER
ExecStart=/usr/local/bin/promtail -config.file=$CONFIG_DIR/promtail.yml
Restart=always
RestartSec=10
StandardOutput=journal
StandardError=journal
SyslogIdentifier=promtail

[Install]
WantedBy=multi-user.target
EOF

    systemctl daemon-reload
    log_success "Systemd service created"
}

# Create SysV init script
create_sysv_service() {
    log_info "Creating SysV init script..."
    
    cat > /etc/init.d/promtail << 'EOF'
#!/bin/bash
# promtail        Promtail log collector
# chkconfig: 35 80 20
# description: Promtail log collector for Grafana Loki

. /etc/rc.d/init.d/functions

USER="promtail"
DAEMON="promtail"
ROOT_DIR="/var/lib/promtail"

SERVER="$ROOT_DIR/$DAEMON"
LOCK_FILE="/var/lock/subsys/promtail"

start() {
    if [ -f $LOCK_FILE ]; then
        echo "promtail is locked."
        return 1
    fi
    
    echo -n $"Shutting down $DAEMON: "
    pid=`ps -aefw | grep "$DAEMON" | grep -v " grep " | awk '{print $2}'`
    kill -9 $pid > /dev/null 2>&1
    [ $? -eq 0 ] && echo "OK" || echo "FAILED"
}

stop() {
    echo -n $"Shutting down $DAEMON: "
    pid=`ps -aefw | grep "$DAEMON" | grep -v " grep " | awk '{print $2}'`
    kill -9 $pid > /dev/null 2>&1
    [ $? -eq 0 ] && echo "OK" || echo "FAILED"
    rm -f $LOCK_FILE
}

case "$1" in
    start)
        start
        ;;
    stop)
        stop
        ;;
    status)
        status $DAEMON
        ;;
    restart)
        stop
        start
        ;;
    *)
        echo "Usage: {start|stop|status|restart}"
        exit 1
        ;;
esac

exit $?
EOF

    chmod +x /etc/init.d/promtail
    chkconfig --add promtail 2>/dev/null || update-rc.d promtail defaults
    log_success "SysV service created"
}

# Create OpenRC service
create_openrc_service() {
    log_info "Creating OpenRC service..."
    
    cat > /etc/init.d/promtail << EOF
#!/sbin/openrc-run

name="promtail"
description="Promtail log collector"

command="/usr/local/bin/promtail"
command_args="-config.file=$CONFIG_DIR/promtail.yml"
command_user="$PROMTAIL_USER"
command_background="yes"
pidfile="/run/\${RC_SVCNAME}.pid"

depend() {
    need net
    after firewall
}
EOF

    chmod +x /etc/init.d/promtail
    log_success "OpenRC service created"
}

# Add promtail user to log access groups
configure_log_access() {
    log_info "Configuring log file access..."
    
    # Try to add to common log access groups
    for group in adm systemd-journal wheel; do
        if getent group $group >/dev/null 2>&1; then
            usermod -a -G $group $PROMTAIL_USER 2>/dev/null || true
            log_info "Added to group: $group"
        fi
    done
    
    log_success "Log access configured"
}

# Start and enable service
start_service() {
    log_info "Starting Promtail service..."
    
    case $SERVICE_MANAGER in
        "systemd")
            systemctl enable promtail
            systemctl start promtail
            sleep 2
            if systemctl is-active --quiet promtail; then
                log_success "Promtail service is running"
                systemctl status promtail --no-pager -l
            else
                log_error "Failed to start Promtail service"
                log_error "Check logs with: journalctl -u promtail -f"
                exit 1
            fi
            ;;
        "sysv")
            service promtail start
            chkconfig promtail on 2>/dev/null || update-rc.d promtail enable
            if service promtail status >/dev/null 2>&1; then
                log_success "Promtail service is running"
            else
                log_error "Failed to start Promtail service"
                exit 1
            fi
            ;;
        "openrc")
            rc-update add promtail default
            rc-service promtail start
            if rc-service promtail status >/dev/null 2>&1; then
                log_success "Promtail service is running"
            else
                log_error "Failed to start Promtail service"
                exit 1
            fi
            ;;
    esac
}

# Main installation process
main() {
    echo "=================================="
    echo "   Promtail Installation Script   "
    echo "=================================="
    echo
    
    check_root
    detect_os
    
    if check_promtail_installed; then
        log_warning "Promtail appears to be already installed"
        if [[ -t 0 ]]; then
            echo "Existing installation found. Do you want to continue and reconfigure? (y/N)"
            read -r response
            if [[ ! "$response" =~ ^[Yy]$ ]]; then
                log_info "Installation cancelled"
                exit 0
            fi
        else
            log_info "Non-interactive mode: Reconfiguring existing installation"
        fi
    fi
    
    # Test Loki connectivity first
    if ! test_loki_connectivity; then
        if [[ -t 0 ]]; then
            echo "Do you want to continue anyway? (y/N)"
            read -r response
            if [[ ! "$response" =~ ^[Yy]$ ]]; then
                log_info "Installation cancelled"
                exit 1
            fi
        else
            log_warning "Non-interactive mode: Continuing despite connectivity issues"
        fi
    fi
    
    install_dependencies
    create_promtail_user
    
    if ! check_promtail_installed; then
        install_promtail
    else
        log_info "Promtail binary already exists, skipping download"
    fi
    
    create_config
    create_service
    configure_log_access
    start_service
    
    echo
    echo "=================================="
    log_success "Promtail installation completed!"
    echo "=================================="
    echo
    echo "Configuration file: $CONFIG_DIR/promtail.yml"
    echo "Service status: systemctl status promtail"
    echo "Service logs: journalctl -u promtail -f"
    echo "Loki endpoint: http://$LOKI_ENDPOINT:$LOKI_PORT"
    echo
    echo "To check if logs are being sent to Loki:"
    echo "curl -G -s \"http://$LOKI_ENDPOINT:$LOKI_PORT/loki/api/v1/query\" --data-urlencode 'query={job=\"system-logs\"}'"
}

# Execute main function
main "$@"