#!/bin/bash

##############################################################################
#### Promtail to Grafana Alloy Migration Script                           ####
####                                                                      ####
#### Detects OS, reads existing Promtail config for Loki URL/hostname,    ####
#### generates equivalent Alloy River config, installs Alloy, and         ####
#### handles the cutover from Promtail to Alloy.                          ####
####                                                                      ####
#### Supports: Ubuntu, Debian, RHEL, CentOS, Rocky, Alma, Amazon Linux   ####
####                                                                      ####
#### Author: Phil Connor                                                  ####
#### License: MIT                                                         ####
#### Contact: contact@mylinux.work                                        ####
#### Version: 1.0.0-030326                                                ####
##############################################################################

set -euo pipefail

readonly SCRIPT_NAME=$(basename "$0")
readonly SCRIPT_VERSION="1.0.0-030326"

# Defaults
ALLOY_CONFIG_DIR="/etc/alloy"
ALLOY_CONFIG_FILE="/etc/alloy/config.alloy"
PROMTAIL_CONFIG="/etc/promtail/config.yml"
LOKI_URL=""
CUSTOM_HOSTNAME=""
DRY_RUN=false
GENERATE_ONLY=false
SKIP_INSTALL=false
SKIP_CUTOVER=false
KEEP_PROMTAIL=true
VERBOSE=false
INCLUDE_JOURNAL=true
INCLUDE_NGINX=false
INCLUDE_APACHE=false

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

log()   { echo -e "${GREEN}[INFO]${NC} $1"; }
warn()  { echo -e "${YELLOW}[WARN]${NC} $1"; }
error() { echo -e "${RED}[ERROR]${NC} $1" >&2; }
debug() { [[ "$VERBOSE" == true ]] && echo -e "${BLUE}[DEBUG]${NC} $1"; }

show_help() {
    cat << EOF
Usage: $SCRIPT_NAME [OPTIONS]

Migrate from Promtail to Grafana Alloy. Generates an Alloy config that
maintains Promtail-compatible labels so existing dashboards keep working.

OPTIONS:
    --loki-url URL        Loki push URL (default: extracted from Promtail config)
    --hostname NAME       Override hostname (default: auto-detect or from Promtail)
    --promtail-config F   Path to existing Promtail config (default: /etc/promtail/config.yml)
    --output FILE         Alloy config output path (default: /etc/alloy/config.alloy)
    --generate-only       Only generate the Alloy config, don't install or cutover
    --skip-install        Skip Alloy installation (already installed)
    --skip-cutover        Generate config and install, but don't stop Promtail
    --no-journal          Skip systemd journal collection
    --include-nginx       Include nginx log collection
    --include-apache      Include Apache log collection
    --remove-promtail     Remove Promtail package after cutover (default: keep)
    --dry-run             Show what would be done without making changes
    --verbose             Enable verbose output
    --version             Show version
    --help, -h            Show this help message

EXAMPLES:
    # Auto-detect everything from existing Promtail config
    sudo $SCRIPT_NAME

    # Specify Loki URL and hostname
    sudo $SCRIPT_NAME --loki-url http://loki.example.com:3100 --hostname web-01

    # Generate config only (don't install or cutover)
    $SCRIPT_NAME --generate-only --loki-url http://loki:3100 --output /tmp/config.alloy

    # Full migration with nginx logs
    sudo $SCRIPT_NAME --include-nginx --remove-promtail

    # Dry run to see what would happen
    sudo $SCRIPT_NAME --dry-run
EOF
}

parse_arguments() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            --loki-url)       LOKI_URL="$2"; shift 2 ;;
            --hostname)       CUSTOM_HOSTNAME="$2"; shift 2 ;;
            --promtail-config) PROMTAIL_CONFIG="$2"; shift 2 ;;
            --output)         ALLOY_CONFIG_FILE="$2"; shift 2 ;;
            --generate-only)  GENERATE_ONLY=true; shift ;;
            --skip-install)   SKIP_INSTALL=true; shift ;;
            --skip-cutover)   SKIP_CUTOVER=true; shift ;;
            --no-journal)     INCLUDE_JOURNAL=false; shift ;;
            --include-nginx)  INCLUDE_NGINX=true; shift ;;
            --include-apache) INCLUDE_APACHE=true; shift ;;
            --remove-promtail) KEEP_PROMTAIL=false; shift ;;
            --dry-run)        DRY_RUN=true; shift ;;
            --verbose)        VERBOSE=true; shift ;;
            --version)        echo "$SCRIPT_NAME version $SCRIPT_VERSION"; exit 0 ;;
            --help|-h)        show_help; exit 0 ;;
            *)                error "Unknown option: $1"; show_help; exit 1 ;;
        esac
    done
}

detect_os() {
    if [[ -f /etc/os-release ]]; then
        . /etc/os-release
        OS=$ID
        OS_PRETTY="$PRETTY_NAME"
    else
        error "Cannot detect OS"
        exit 1
    fi
    debug "Detected OS: $OS_PRETTY"
}

detect_hostname() {
    if [[ -n "$CUSTOM_HOSTNAME" ]]; then
        DETECTED_HOSTNAME="$CUSTOM_HOSTNAME"
        debug "Using custom hostname: $DETECTED_HOSTNAME"
        return
    fi

    # Try to extract from Promtail config
    if [[ -f "$PROMTAIL_CONFIG" ]]; then
        local pt_host
        pt_host=$(grep -m1 'host:' "$PROMTAIL_CONFIG" 2>/dev/null | awk '{print $2}' | tr -d '"' || true)
        if [[ -n "$pt_host" ]]; then
            DETECTED_HOSTNAME="$pt_host"
            debug "Extracted hostname from Promtail config: $DETECTED_HOSTNAME"
            return
        fi
    fi

    DETECTED_HOSTNAME=$(hostname -f 2>/dev/null || hostname)
    debug "Using system hostname: $DETECTED_HOSTNAME"
}

detect_loki_url() {
    if [[ -n "$LOKI_URL" ]]; then
        debug "Using provided Loki URL: $LOKI_URL"
        return
    fi

    # Extract from Promtail config
    if [[ -f "$PROMTAIL_CONFIG" ]]; then
        LOKI_URL=$(grep -m1 'url:' "$PROMTAIL_CONFIG" 2>/dev/null | awk '{print $2}' | tr -d '"' | sed 's|/loki/api/v1/push||' || true)
        if [[ -n "$LOKI_URL" ]]; then
            debug "Extracted Loki URL from Promtail config: $LOKI_URL"
            return
        fi
    fi

    error "Could not determine Loki URL"
    error "Provide with --loki-url or ensure Promtail config exists at $PROMTAIL_CONFIG"
    exit 1
}

check_promtail_status() {
    if systemctl is-active --quiet promtail 2>/dev/null; then
        PROMTAIL_RUNNING=true
        log "Promtail is currently running"
    else
        PROMTAIL_RUNNING=false
        debug "Promtail is not running"
    fi
}

# Generate an Alloy loki.source.file block if the log file exists
generate_file_source() {
    local label="$1"
    local path="$2"
    local job="$3"
    local extra_labels="$4"

    if [[ "$DRY_RUN" == true ]] || [[ -f "$path" ]] || [[ "$path" == *"*"* ]]; then
        cat << EOF

loki.source.file "$label" {
    targets = [
        {
            "__path__" = "$path",
            "job"      = "$job",
            "host"     = "$DETECTED_HOSTNAME",${extra_labels}
        },
    ]
    forward_to = [loki.write.default.receiver]
}
EOF
    else
        debug "Skipping $path (file does not exist)"
    fi
}

generate_alloy_config() {
    log "Generating Alloy config for $OS ($DETECTED_HOSTNAME)..."

    local os_label
    case "$OS" in
        ubuntu|debian) os_label="ubuntu" ;;
        rhel|centos|rocky|almalinux|amzn) os_label="rhel-family" ;;
        *) os_label="$OS" ;;
    esac

    local config=""

    # Header
    config+="// Grafana Alloy Configuration for $DETECTED_HOSTNAME
// Migrated from Promtail on $(date +%Y-%m-%d)
// OS: $OS_PRETTY
// Labels maintained for Promtail dashboard compatibility

logging {
    level = \"info\"
}
"

    # Journal source
    if [[ "$INCLUDE_JOURNAL" == true ]]; then
        config+="
// System logs via systemd journal
loki.source.journal \"systemd_journal\" {
    max_age = \"12h\"
    labels = {
        job  = \"systemd-journal\",
        host = \"$DETECTED_HOSTNAME\",
        os   = \"$os_label\",
    }
    forward_to = [loki.relabel.journal_relabel.receiver]
}

loki.relabel \"journal_relabel\" {
    forward_to = [loki.write.default.receiver]

    rule {
        source_labels = [\"__journal__systemd_unit\"]
        target_label  = \"unit\"
    }

    rule {
        source_labels = [\"__journal_priority\"]
        target_label  = \"priority\"
    }

    rule {
        source_labels = [\"__journal__hostname\"]
        target_label  = \"hostname\"
    }
}
"
    fi

    # OS-specific file sources
    case "$OS" in
        ubuntu|debian)
            config+="
// Ubuntu/Debian system logs"
            config+=$(generate_file_source "syslog" "/var/log/syslog" "messages" "
            \"os\"       = \"ubuntu\",")
            config+=$(generate_file_source "auth" "/var/log/auth.log" "auth" "
            \"log_type\" = \"authentication\",")
            config+=$(generate_file_source "kern" "/var/log/kern.log" "kernel" "")
            config+=$(generate_file_source "cron" "/var/log/cron.log" "cron" "")
            config+=$(generate_file_source "mail" "/var/log/mail.log" "mail" "")
            config+=$(generate_file_source "apt" "/var/log/apt/history.log" "packages" "
            \"package_manager\" = \"apt\",")
            config+=$(generate_file_source "boot" "/var/log/boot.log" "boot" "")
            ;;
        rhel|centos|rocky|almalinux|amzn)
            config+="
// RHEL/CentOS/Rocky/Alma/Amazon Linux system logs"
            config+=$(generate_file_source "messages" "/var/log/messages" "messages" "
            \"os\"       = \"rhel-family\",")
            config+=$(generate_file_source "secure" "/var/log/secure" "auth" "
            \"log_type\" = \"authentication\",")
            config+=$(generate_file_source "cron" "/var/log/cron" "cron" "")
            config+=$(generate_file_source "maillog" "/var/log/maillog" "mail" "")
            config+=$(generate_file_source "yum" "/var/log/yum.log" "packages" "
            \"package_manager\" = \"yum\",")
            config+=$(generate_file_source "boot" "/var/log/boot.log" "boot" "")
            ;;
        *)
            config+="
// Generic system logs"
            config+=$(generate_file_source "syslog" "/var/log/syslog" "messages" "")
            config+=$(generate_file_source "auth" "/var/log/auth.log" "auth" "
            \"log_type\" = \"authentication\",")
            ;;
    esac

    # Application wildcard
    config+=$(generate_file_source "application_logs" "/var/log/*.log" "application" "")

    # Nginx
    if [[ "$INCLUDE_NGINX" == true ]]; then
        config+="
// Nginx logs"
        config+=$(generate_file_source "nginx_access" "/var/log/nginx/access.log" "nginx" "
            \"log_type\" = \"access\",")
        config+=$(generate_file_source "nginx_error" "/var/log/nginx/error.log" "nginx" "
            \"log_type\" = \"error\",")
    fi

    # Apache
    if [[ "$INCLUDE_APACHE" == true ]]; then
        config+="
// Apache logs"
        config+=$(generate_file_source "apache_access" "/var/log/apache2/access.log" "apache" "
            \"log_type\" = \"access\",")
        config+=$(generate_file_source "apache_error" "/var/log/apache2/error.log" "apache" "
            \"log_type\" = \"error\",")
        config+=$(generate_file_source "httpd_access" "/var/log/httpd/access_log" "apache" "
            \"log_type\" = \"access\",")
        config+=$(generate_file_source "httpd_error" "/var/log/httpd/error_log" "apache" "
            \"log_type\" = \"error\",")
    fi

    # Loki write endpoint
    config+="

// Write to Loki
loki.write \"default\" {
    endpoint {
        url = \"${LOKI_URL}/loki/api/v1/push\"
    }
}
"

    GENERATED_CONFIG="$config"
}

write_config() {
    local output_file="$1"

    if [[ "$DRY_RUN" == true ]]; then
        log "DRY RUN: Would write config to $output_file"
        echo "--- Generated config.alloy ---"
        echo "$GENERATED_CONFIG"
        echo "--- End config ---"
        return
    fi

    local output_dir
    output_dir=$(dirname "$output_file")
    mkdir -p "$output_dir"

    # Backup existing config
    if [[ -f "$output_file" ]]; then
        local backup="${output_file}.bak.$(date +%Y%m%d%H%M%S)"
        cp "$output_file" "$backup"
        log "Backed up existing config to $backup"
    fi

    echo "$GENERATED_CONFIG" > "$output_file"
    log "Alloy config written to $output_file"
}

install_alloy() {
    if [[ "$DRY_RUN" == true ]]; then
        log "DRY RUN: Would install Grafana Alloy"
        return
    fi

    # Check if already installed
    if command -v alloy >/dev/null 2>&1; then
        log "Alloy is already installed: $(alloy --version 2>&1 | head -1)"
        return
    fi

    log "Installing Grafana Alloy..."

    case "$OS" in
        ubuntu|debian)
            apt-get install -y apt-transport-https software-properties-common
            mkdir -p /etc/apt/keyrings/
            wget -q -O - https://apt.grafana.com/gpg.key | gpg --dearmor | tee /etc/apt/keyrings/grafana.gpg > /dev/null
            echo "deb [signed-by=/etc/apt/keyrings/grafana.gpg] https://apt.grafana.com stable main" | tee /etc/apt/sources.list.d/grafana.list
            apt-get update -qq
            apt-get install -y alloy
            ;;
        rhel|centos|rocky|almalinux|amzn)
            cat > /etc/yum.repos.d/grafana.repo << 'REPO'
[grafana]
name=grafana
baseurl=https://rpm.grafana.com
repo_gpgcheck=1
enabled=1
gpgcheck=1
gpgkey=https://rpm.grafana.com/gpg.key
sslverify=1
sslcacert=/etc/pki/tls/certs/ca-bundle.crt
REPO
            if command -v dnf >/dev/null 2>&1; then
                dnf install -y alloy
            else
                yum install -y alloy
            fi
            ;;
        *)
            error "Unsupported OS for automatic installation: $OS"
            error "Install Alloy manually: https://grafana.com/docs/alloy/latest/set-up/install/"
            exit 1
            ;;
    esac

    log "Alloy installed: $(alloy --version 2>&1 | head -1)"
}

validate_config() {
    if [[ "$DRY_RUN" == true ]]; then
        log "DRY RUN: Would validate config with 'alloy fmt'"
        return
    fi

    if ! command -v alloy >/dev/null 2>&1; then
        warn "Alloy not installed, skipping validation"
        return
    fi

    log "Validating Alloy config..."
    if alloy fmt "$ALLOY_CONFIG_FILE" >/dev/null 2>&1; then
        log "Config validation passed"
    else
        error "Config validation failed. Check $ALLOY_CONFIG_FILE for syntax errors"
        error "Run: alloy fmt $ALLOY_CONFIG_FILE"
        exit 1
    fi
}

perform_cutover() {
    if [[ "$DRY_RUN" == true ]]; then
        log "DRY RUN: Would stop Promtail and start Alloy"
        return
    fi

    # Stop Promtail
    if systemctl is-active --quiet promtail 2>/dev/null; then
        log "Stopping Promtail..."
        systemctl stop promtail
        systemctl disable promtail
        log "Promtail stopped and disabled"
    fi

    # Add alloy user to required groups
    if getent group adm >/dev/null 2>&1; then
        usermod -a -G adm alloy 2>/dev/null || true
    fi
    if getent group systemd-journal >/dev/null 2>&1; then
        usermod -a -G systemd-journal alloy 2>/dev/null || true
    fi

    # Start Alloy
    log "Starting Alloy..."
    systemctl enable --now alloy
    sleep 2

    if systemctl is-active --quiet alloy; then
        log "Alloy is running"
    else
        error "Alloy failed to start. Check: journalctl -u alloy --no-pager -n 30"
        error "Rolling back — restarting Promtail"
        systemctl enable --now promtail 2>/dev/null || true
        exit 1
    fi

    # Remove Promtail if requested
    if [[ "$KEEP_PROMTAIL" == false ]]; then
        log "Removing Promtail package..."
        case "$OS" in
            ubuntu|debian) apt-get remove -y promtail 2>/dev/null || true ;;
            *) yum remove -y promtail 2>/dev/null || dnf remove -y promtail 2>/dev/null || true ;;
        esac
        log "Promtail removed"
    else
        log "Promtail package kept (use 'systemctl start promtail' to rollback)"
    fi
}

print_summary() {
    echo ""
    echo "=========================================="
    echo "  Migration Summary"
    echo "=========================================="
    echo "  OS:          $OS_PRETTY"
    echo "  Hostname:    $DETECTED_HOSTNAME"
    echo "  Loki URL:    $LOKI_URL"
    echo "  Alloy config: $ALLOY_CONFIG_FILE"

    if [[ "$DRY_RUN" != true ]] && [[ "$GENERATE_ONLY" != true ]]; then
        echo ""
        echo "  Alloy status: $(systemctl is-active alloy 2>/dev/null || echo 'not checked')"
        echo ""
        echo "  Verify:"
        echo "    systemctl status alloy"
        echo "    journalctl -u alloy -f"
        echo "    curl http://localhost:12345  (Alloy UI)"
        echo ""
        echo "  Rollback:"
        echo "    sudo systemctl stop alloy"
        echo "    sudo systemctl start promtail"
    fi

    if [[ "$GENERATE_ONLY" == true ]]; then
        echo ""
        echo "  Config generated. Review and deploy manually."
    fi

    echo "=========================================="
    echo ""
}

main() {
    parse_arguments "$@"

    log "Promtail → Alloy Migration Script v${SCRIPT_VERSION}"
    echo ""

    # Check root (unless generate-only)
    if [[ "$GENERATE_ONLY" != true ]] && [[ "$DRY_RUN" != true ]] && [[ "$EUID" -ne 0 ]]; then
        error "This script must be run as root (or use --generate-only)"
        exit 1
    fi

    detect_os
    detect_hostname
    detect_loki_url
    check_promtail_status

    # Generate config
    generate_alloy_config
    write_config "$ALLOY_CONFIG_FILE"

    if [[ "$GENERATE_ONLY" == true ]]; then
        print_summary
        exit 0
    fi

    # Install Alloy
    if [[ "$SKIP_INSTALL" != true ]]; then
        install_alloy
    fi

    # Validate
    validate_config

    # Cutover
    if [[ "$SKIP_CUTOVER" != true ]]; then
        perform_cutover
    fi

    print_summary
}

main "$@"