#!/usr/bin/env bash

#########################################################################################
#### salt-roster-generator.sh — Generate Salt SSH roster files from external sources ####
#### Supports CSV files, Ansible inventory (INI format), and AWS EC2 instances       ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.01                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   ./salt-roster-generator.sh --csv hosts.csv -o roster                          ####
####   ./salt-roster-generator.sh --ansible-inventory /etc/ansible/hosts -o roster   ####
####   ./salt-roster-generator.sh --ec2 --region us-east-1 -o roster                 ####
####   ./salt-roster-generator.sh --ec2 --tag "Environment=production" -o roster     ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────
VERSION="1.0"
SCRIPT_NAME="$(basename "$0")"
readonly SCRIPT_NAME

INPUT_MODE=""
CSV_FILE=""
ANSIBLE_FILE=""
EC2_MODE=false
EC2_REGION=""
EC2_RUNNING_ONLY=true
declare -a EC2_TAGS=()

OUTPUT_FILE=""
DEFAULT_USER="root"
DEFAULT_PORT="22"
ENABLE_SUDO=false
SSH_PRIV=""

# ── Logging ───────────────────────────────────────────────────────────
err()  { echo "[ERROR] $*" >&2; }
die()  { err "$@"; exit 1; }

# ══════════════════════════════════════════════════════════════════════
#                       USAGE
# ══════════════════════════════════════════════════════════════════════

usage() {
    cat <<EOF
${SCRIPT_NAME} — Generate Salt SSH roster files from external inventory sources

USAGE:
  ${SCRIPT_NAME} [OPTIONS]

INPUT SOURCES (pick one):
  --csv FILE                CSV file (columns: name,host,user,port)
  --ansible-inventory FILE  Ansible INI inventory file
  --ec2                     AWS EC2 instances (requires aws cli)

EC2 OPTIONS:
  --region REGION           AWS region (default: from aws cli config)
  --tag KEY=VALUE           Filter by tag (can repeat)
  --running-only            Only include running instances (default: true)

OUTPUT OPTIONS:
  -o, --output FILE         Output roster file (default: stdout)

SSH OPTIONS (applied to all hosts):
  --user USER               SSH user (default: root)
  --sudo                    Enable sudo for all hosts
  --priv KEY                SSH private key path
  --port PORT               SSH port (default: 22)

OTHER:
  --help, -h                Show this help
  --version                 Show version

EXAMPLES:
  ./${SCRIPT_NAME} --csv hosts.csv -o roster
  ./${SCRIPT_NAME} --ansible-inventory /etc/ansible/hosts -o roster
  ./${SCRIPT_NAME} --ec2 --region us-east-1 -o roster
  ./${SCRIPT_NAME} --ec2 --tag "Environment=production" --tag "Role=web" -o roster
EOF
}

# ══════════════════════════════════════════════════════════════════════
#                       ARGUMENT PARSING
# ══════════════════════════════════════════════════════════════════════

parse_args() {
    [[ $# -eq 0 ]] && { usage; exit 1; }

    while [[ $# -gt 0 ]]; do
        case "$1" in
            --csv)
                INPUT_MODE="csv"
                CSV_FILE="$2"
                shift 2
                ;;
            --ansible-inventory)
                INPUT_MODE="ansible"
                ANSIBLE_FILE="$2"
                shift 2
                ;;
            --ec2)
                INPUT_MODE="ec2"
                EC2_MODE=true
                shift
                ;;
            --region)
                EC2_REGION="$2"
                shift 2
                ;;
            --tag)
                EC2_TAGS+=("$2")
                shift 2
                ;;
            --running-only)
                EC2_RUNNING_ONLY=true
                shift
                ;;
            -o|--output)
                OUTPUT_FILE="$2"
                shift 2
                ;;
            --user)
                DEFAULT_USER="$2"
                shift 2
                ;;
            --sudo)
                ENABLE_SUDO=true
                shift
                ;;
            --priv)
                SSH_PRIV="$2"
                shift 2
                ;;
            --port)
                DEFAULT_PORT="$2"
                shift 2
                ;;
            --version)
                echo "${SCRIPT_NAME} v${VERSION}"
                exit 0
                ;;
            --help|-h)
                usage
                exit 0
                ;;
            -*)
                die "Unknown option: $1 — run ${SCRIPT_NAME} --help"
                ;;
            *)
                die "Unexpected argument: $1 — run ${SCRIPT_NAME} --help"
                ;;
        esac
    done
}

# ══════════════════════════════════════════════════════════════════════
#                       VALIDATION
# ══════════════════════════════════════════════════════════════════════

validate() {
    [[ -z "$INPUT_MODE" ]] && die "No input source specified. Use --csv, --ansible-inventory, or --ec2."

    case "$INPUT_MODE" in
        csv)
            [[ -z "$CSV_FILE" ]] && die "--csv requires a file argument."
            [[ ! -f "$CSV_FILE" ]] && die "CSV file not found: $CSV_FILE"
            ;;
        ansible)
            [[ -z "$ANSIBLE_FILE" ]] && die "--ansible-inventory requires a file argument."
            [[ ! -f "$ANSIBLE_FILE" ]] && die "Ansible inventory file not found: $ANSIBLE_FILE"
            ;;
        ec2)
            if ! command -v aws &>/dev/null; then
                die "aws cli not found. Install it to use --ec2 mode."
            fi
            ;;
    esac
}

# ══════════════════════════════════════════════════════════════════════
#                       ROSTER GENERATION
# ══════════════════════════════════════════════════════════════════════

generate_header() {
    local source_label="$1"
    echo "# Salt SSH Roster"
    echo "# Generated by ${SCRIPT_NAME} v${VERSION}"
    echo "# Source: ${source_label}"
    echo "# Date: $(date '+%Y-%m-%d')"
    echo ""
}

generate_entry() {
    local name="$1"
    local host="$2"
    local user="${3:-$DEFAULT_USER}"
    local port="${4:-$DEFAULT_PORT}"
    local sudo="${5:-}"
    local priv="${6:-$SSH_PRIV}"

    [[ -z "$name" || -z "$host" ]] && return

    echo "${name}:"
    echo "  host: ${host}"
    echo "  user: ${user}"
    echo "  port: ${port}"

    if [[ "$ENABLE_SUDO" == true || "$sudo" == "true" || "$sudo" == "True" ]]; then
        echo "  sudo: True"
    fi

    if [[ -n "$priv" ]]; then
        echo "  priv: ${priv}"
    fi

    echo ""
}

# ── CSV Parser ────────────────────────────────────────────────────────
parse_csv() {
    local file="$1"
    local header_line
    header_line=$(head -1 "$file")

    # Map column positions from header
    local -A col_map
    local idx=0
    IFS=',' read -ra cols <<< "$header_line"
    for c in "${cols[@]}"; do
        c=$(echo "$c" | tr -d '[:space:]' | tr '[:upper:]' '[:lower:]')
        col_map["$c"]=$idx
        ((idx++)) || true
    done

    # Require name and host columns
    [[ -z "${col_map[name]:-}" ]] && die "CSV missing required column: name"
    [[ -z "${col_map[host]:-}" ]] && die "CSV missing required column: host"

    local name_idx="${col_map[name]}"
    local host_idx="${col_map[host]}"
    local user_idx="${col_map[user]:-}"
    local port_idx="${col_map[port]:-}"
    local sudo_idx="${col_map[sudo]:-}"
    local priv_idx="${col_map[priv]:-}"

    generate_header "$file"

    local line_num=0
    while IFS= read -r line; do
        ((line_num++)) || true
        [[ $line_num -eq 1 ]] && continue           # skip header
        [[ -z "$line" || "$line" =~ ^[[:space:]]*# ]] && continue

        IFS=',' read -ra fields <<< "$line"

        local name="${fields[$name_idx]:-}"
        local host="${fields[$host_idx]:-}"
        local user="${DEFAULT_USER}"
        local port="${DEFAULT_PORT}"
        local sudo=""
        local priv=""

        [[ -n "$user_idx" && -n "${fields[$user_idx]:-}" ]] && user="${fields[$user_idx]}"
        [[ -n "$port_idx" && -n "${fields[$port_idx]:-}" ]] && port="${fields[$port_idx]}"
        [[ -n "$sudo_idx" && -n "${fields[$sudo_idx]:-}" ]] && sudo="${fields[$sudo_idx]}"
        [[ -n "$priv_idx" && -n "${fields[$priv_idx]:-}" ]] && priv="${fields[$priv_idx]}"

        # Trim whitespace
        name=$(echo "$name" | xargs)
        host=$(echo "$host" | xargs)
        user=$(echo "$user" | xargs)
        port=$(echo "$port" | xargs)

        generate_entry "$name" "$host" "$user" "$port" "$sudo" "$priv"
    done < "$file"
}

# ── Ansible INI Inventory Parser ──────────────────────────────────────
parse_ansible_inventory() {
    local file="$1"

    generate_header "$file"

    while IFS= read -r line; do
        # Strip leading/trailing whitespace
        line=$(echo "$line" | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')

        # Skip empty lines, comments, group headers, group:vars, group:children
        [[ -z "$line" ]] && continue
        [[ "$line" =~ ^# ]] && continue
        [[ "$line" =~ ^\[.*\] ]] && continue

        # Parse host line: hostname key=value key=value ...
        local hostname
        hostname=$(echo "$line" | awk '{print $1}')

        local host="$hostname"
        local user="$DEFAULT_USER"
        local port="$DEFAULT_PORT"
        local priv=""

        # Extract ansible variables
        if echo "$line" | grep -q 'ansible_host='; then
            host=$(echo "$line" | grep -oP 'ansible_host=\K[^\s]+')
        fi
        if echo "$line" | grep -q 'ansible_user='; then
            user=$(echo "$line" | grep -oP 'ansible_user=\K[^\s]+')
        fi
        if echo "$line" | grep -q 'ansible_port='; then
            port=$(echo "$line" | grep -oP 'ansible_port=\K[^\s]+')
        fi
        if echo "$line" | grep -q 'ansible_ssh_private_key_file='; then
            priv=$(echo "$line" | grep -oP 'ansible_ssh_private_key_file=\K[^\s]+')
        fi

        generate_entry "$hostname" "$host" "$user" "$port" "" "$priv"
    done < "$file"
}

# ── EC2 Parser ────────────────────────────────────────────────────────
parse_ec2() {
    local -a aws_args=("ec2" "describe-instances")

    # Region
    if [[ -n "$EC2_REGION" ]]; then
        aws_args+=("--region" "$EC2_REGION")
    fi

    # Build filters
    local -a filters=()
    if [[ "$EC2_RUNNING_ONLY" == true ]]; then
        filters+=("Name=instance-state-name,Values=running")
    fi
    for tag in "${EC2_TAGS[@]}"; do
        local key="${tag%%=*}"
        local val="${tag#*=}"
        filters+=("Name=tag:${key},Values=${val}")
    done
    if [[ ${#filters[@]} -gt 0 ]]; then
        aws_args+=("--filters" "${filters[@]}")
    fi

    aws_args+=(
        "--query" "Reservations[].Instances[].[Tags[?Key==\`Name\`].Value | [0], PrivateIpAddress, InstanceId]"
        "--output" "text"
    )

    local ec2_output
    ec2_output=$(aws "${aws_args[@]}" 2>/dev/null) || die "aws ec2 describe-instances failed. Check credentials and region."

    [[ -z "$ec2_output" ]] && die "No EC2 instances found matching the specified filters."

    local source_label="ec2"
    [[ -n "$EC2_REGION" ]] && source_label="ec2 (${EC2_REGION})"
    generate_header "$source_label"

    while IFS=$'\t' read -r name ip instance_id; do
        # Use instance ID if no Name tag
        if [[ -z "$name" || "$name" == "None" ]]; then
            name="$instance_id"
        fi

        [[ -z "$ip" || "$ip" == "None" ]] && continue

        generate_entry "$name" "$ip" "$DEFAULT_USER" "$DEFAULT_PORT" "" ""
    done <<< "$ec2_output"
}

# ══════════════════════════════════════════════════════════════════════
#                       OUTPUT
# ══════════════════════════════════════════════════════════════════════

write_output() {
    local content="$1"

    if [[ -n "$OUTPUT_FILE" ]]; then
        echo "$content" > "$OUTPUT_FILE"
        echo "Roster written to ${OUTPUT_FILE} ($(echo "$content" | grep -c ':$' || true) hosts)" >&2
    else
        echo "$content"
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                       MAIN
# ══════════════════════════════════════════════════════════════════════

main() {
    parse_args "$@"
    validate

    local roster=""

    case "$INPUT_MODE" in
        csv)     roster=$(parse_csv "$CSV_FILE") ;;
        ansible) roster=$(parse_ansible_inventory "$ANSIBLE_FILE") ;;
        ec2)     roster=$(parse_ec2) ;;
    esac

    write_output "$roster"
}

main "$@"