#!/bin/bash

# A simple script to verify HTTP endpoints and generate a Prometheus scrape configuration file.
#
# Exit codes:
#   0: File content was already correct. No changes made.
#   1: An error occurred (e.g., file not found).
#   2: File content was updated.

#set -x

BASE_DIR="/usr/share/soe-utils"

# Set some default values
ENDPOINTS_FILE="${BASE_DIR}/prometheus-exporters.txt"
TEMPLATE_DIR="${BASE_DIR}/prometheus-generate-config"
OUTPUT_FILE="prometheus-scrape-config.yml"
SITE_CODE="ZZ"
SITE_NAME="Unknown"
QUIET_MODE=false
DEBUG_MODE=false
EXPORTER_COUNT=0

# Check for SOE config file and source it if it exists
[ -r /etc/default/soe-config ] && . /etc/default/soe-config

# Determine bash version
bash_major_version=$(echo "$BASH_VERSION" | cut -d'.' -f1)

# Function to show the help message
show_help() {
  echo "Usage: $0 --file <endpoints_file> [--output <output_file>]"
  echo "Options:"
  echo "  -f, --file    The input file containing the list of endpoints (required)."
  echo "                Default: $ENDPOINTS_FILE"
  echo "  -o, --output  The name of the output Prometheus YAML file (optional)."
  echo "                Default: $OUTPUT_FILE"
  echo "  -q, --quiet   Suppress all progress messages. Only final success/error"
  echo "                messages will be shown. Useful for cron jobs."
  echo "  -d, --debug   Verbose messaging, disables --quiet if both used"
  echo ""
  echo "  --site-code   The site code"
  echo "                Default: $SITE_CODE"
  echo "  --site-name   The site name"
  echo "                Default: $SITE_NAME"
  echo "  --site-host   The site hostname override"
  echo ""
  echo "  --base-dir    The base directory to set for exporter list and templates"
  echo "                useful for testing"
  echo "                Default: $BASE_DIR"
  echo ""
  echo "  -h, --help    Show this help message and exit."
}

# Parse command-line options
while [[ "$#" -gt 0 ]]; do
  case "$1" in
    -f|--file)
      if [ -z "$2" ]; then
        echo "Error: -f|--file requires an argument." >&2
        show_help
        exit 1
      fi
      ENDPOINTS_FILE="$2"
      shift 2
      ;;
    -o|--output)
      if [ -z "$2" ]; then
        echo "Error: -o|--output requires an argument." >&2
        show_help
        exit 1
      fi
      OUTPUT_FILE="$2"
      shift 2
      ;;
    --site-code)
      if [ -z "$2" ]; then
        echo "Error: --site-code requires an argument." >&2
        show_help
        exit 1
      fi
      # Make sure site code is uppercase
      if [ "$bash_major_version" -ge 4 ]; then
        # This is the modern, built-in method for Bash 4.0+.
        SITE_CODE="${2^^}"
      else
        # This is the classic method using the 'tr' (translate) command.
        SITE_CODE=$(echo "$2" | tr '[:lower:]' '[:upper:]')
      fi
      shift 2
      ;;
    --site-name)
      if [ -z "$2" ]; then
        echo "Error: --site-name requires an argument." >&2
        show_help
        exit 1
      fi
      SITE_NAME="$2"
      shift 2
      ;;
    --site-host)
      if [ -z "$2" ]; then
        echo "Error: --site-host requires an argument." >&2
        show_help
        exit 1
      fi
      HOSTNAME_OVERRIDE="$2"
      shift 2
      ;;
    --base-dir)
      if [ -z "$2" ]; then
        echo "Error: --base-dir requires an argument." >&2
        show_help
        exit 1
      fi
      BASE_DIR="$2"

      # TODO: consider if `-f/--file` was used to overwrite endpoint file
      ENDPOINTS_FILE="${BASE_DIR}/prometheus-exporters.txt"
      TEMPLATE_DIR="${BASE_DIR}/prometheus-generate-config"
      shift 2
      ;;
    -q|--quiet)
      QUIET_MODE=true
      shift
      ;;
    -d|--debug)
      DEBUG_MODE=true
      shift
      ;;
    -h|--help)
      show_help
      exit 0
      ;;
    *)
      echo "Error: Unknown option: $1" >&2
      show_help
      exit 1
      ;;
  esac
done

# Check if the mandatory endpoints file was provided
if [ -z "$ENDPOINTS_FILE" ]; then
  echo "Error: The endpoints file must be specified with -f or --file." >&2
  show_help
  exit 1
fi

# Check if the directory holding templates exists
if [ ! -d "$TEMPLATE_DIR" ]; then
  echo "Error: The template directory doesn't exist"
  echo "       $TEMPLATE_DIR"
  exit 1
fi

if [ "$DEBUG_MODE" = true ]; then
  QUIET_MODE=false
fi

# Get the fully qualified domain name (FQDN) of the host running the script.
# First the hostname from the /etc/hostname file, on SOE it should be FQDN
HOSTNAME_FILE=$(cat /etc/hostname 2>/dev/null)
if [ -z "$HOSTNAME_FILE" ]; then
  #echo "WARNING: The hostname file isn't populated."
  HOSTNAME_FILE="NO_FILE"
fi

# Get the FQDN which reads it from `/etc/hosts` and DNS as defined by `nsswitch.conf`
# Linux: hostname --fqdn | hostname -f
# MacOS: hostname -f
HOSTNAME_CMD=$(hostname -f)
if [ -z "$HOSTNAME_CMD" ]; then
  #echo "WARNING: Could not determine the full hostname"
  HOSTNAME_CMD="NO_CMD"
fi

# Compare the short name and the FQDN
if [ "$HOSTNAME_FILE" == "$HOSTNAME_CMD" ]; then
  # FQDN is valid, use it for the host ID
  HOST_ID="$HOSTNAME_CMD"
  echo "Using FQDN '$HOST_ID' for instance label."
else
  if [ -n "$HOSTNAME_OVERRIDE" ]; then
    HOST_ID="$HOSTNAME_OVERRIDE"
  else
    # FQDN is not valid (e.g., matches short name). Fall back to hostname-IP combo.
    echo "ERROR: Hostname mis-match, make sure there isn't anything out of place. Exiting"
    exit 1
  fi
fi

# Set a temporary output file
TEMP_OUTPUT_FILE="$OUTPUT_FILE.tmp"

# Get the current date and time in ISO format with the timezone
CURRENT_DATE=$(date +"%Y-%m-%dT%H:%M:%S%z")

# Add header comments to the output file
echo "# This file was automatically generated by $0" > "$TEMP_OUTPUT_FILE"
echo "# Do not modify this file as it will be overwritten" >> "$TEMP_OUTPUT_FILE"
echo "# Generation date: $CURRENT_DATE" >> "$TEMP_OUTPUT_FILE"
echo "---" >> "$TEMP_OUTPUT_FILE"
echo "scrape_configs:" >> "$TEMP_OUTPUT_FILE"

# Loop through each line in the endpoints file
while read -r url interval job_name target_sd; do
  # Skip comments and empty lines
  if [[ "$url" =~ ^#.* ]] || [[ -z "$url" ]]; then
    continue
  fi

  if [ "$DEBUG_MODE" = true ]; then
    echo "[$job_name] Endpoint: '$url every $interval"
  fi

  # Use curl to get the HTTP status code
  status_code=$(curl -s -o /dev/null -w "%{http_code}" "$url")

  # Check if the status code indicates success (2xx)
  if [[ "$status_code" =~ ^2 ]]; then
    ((EXPORTER_COUNT++))
    if [ "$QUIET_MODE" = false ]; then
      echo "SUCCESS: '$job_name' endpoint is up. Generating YAML entry."
    fi

    # Get the host, port, and path
    EXPORTER_HOST=$(echo "$url" | awk -F'//' '{print $2}' | awk -F':' '{print $1}')
    EXPORTER_PORT=$(echo "$url" | sed -r 's/.*:([0-9]+)\/.*/\1/')
    EXPORTER_PATH=$(echo "$url" | cut -d'/' -f4-)

    # Create the instance label by combining the host and the port
    instance_label="${HOST_ID}:${EXPORTER_PORT}"

    # NOTE: Job name is used to track site it belongs to
    # Make sure site code is lowercase
    if [ "$bash_major_version" -ge 4 ]; then
      # This is the modern, built-in method for Bash 4.0+.
      lower_sitecode="${SITE_CODE,,}"
    else
      # This is the classic method using the 'tr' (translate) command.
      lower_sitecode=$(echo "$SITE_CODE" | tr '[:upper:]' '[:lower:]')
    fi
    site_job_name="${job_name}-${lower_sitecode}"

    # Choose a template to use when generating config
    TEMPLATE_FILE=""
    if [ -f "$TEMPLATE_DIR/scrape.$job_name.tpl" ]; then
      TEMPLATE_FILE="$TEMPLATE_DIR/scrape.$job_name.tpl"
    else
      TEMPLATE_FILE="$TEMPLATE_DIR/scrape.tpl"
      if [ ! -f "$TEMPLATE_FILE" ]; then
        echo "ERROR: Default template file not found. Exiting"
        if [ -f "$TEMP_OUTPUT_FILE" ]; then
          rm "$TEMP_OUTPUT_FILE"
        fi
        exit 1
      fi
    fi

    # Generate the YAML entry and append it to the output file
    sed -e "s|__LABEL_INSTANCE__|$instance_label|g" \
        -e "s|__LABEL_HOSTNAME__|$HOST_ID|g" \
        -e "s|__LABEL_SITE_CODE__|$SITE_CODE|g" \
        -e "s|__JOB_NAME__|$site_job_name|g" \
        -e "s|__JOB_INTERVAL__|$interval|g" \
        -e "s|__JOB_TARGETS__|$url|g" \
        -e "s|__EXPORTER_HOST__|$EXPORTER_HOST|g" \
        -e "s|__EXPORTER_PORT__|$EXPORTER_PORT|g" \
        -e "s|__TARGET_SD__|$target_sd|g" \
        "$TEMPLATE_FILE" >> "$TEMP_OUTPUT_FILE"
  else
    if [ "$DEBUG_MODE" = true ]; then
      echo "Check Status: FAILED! Endpoint $url returned status code $status_code. Skipping."
    fi
  fi
done < "$ENDPOINTS_FILE"

if (( EXPORTER_COUNT >= 1 )); then
  if [ "$QUIET_MODE" = false ]; then
    echo "Found $EXPORTER_COUNT"
  fi
else
  echo "No Endpoints"
  if [ -f "$TEMP_OUTPUT_FILE" ]; then
    rm "$TEMP_OUTPUT_FILE"
  fi

  if [ -f "$OUTPUT_FILE" ]; then
    rm "$OUTPUT_FILE"
  fi

  exit 0
fi

# Now, compare the newly generated file with the existing one (if it exists)
if [ -f "$OUTPUT_FILE" ]; then
  # Use diff with process substitution to compare files starting from line 4
  if diff -q <(tail -n +4 "$OUTPUT_FILE") <(tail -n +4 "$TEMP_OUTPUT_FILE") >/dev/null; then
    if [ "$QUIET_MODE" = false ]; then
      echo "Configuration is the same. No changes detected."
    fi
    rm "$TEMP_OUTPUT_FILE"
    exit 0
  else
    if [ "$QUIET_MODE" = false ]; then
      echo "Changes detected. Updating '$OUTPUT_FILE'."
    fi
    mv "$TEMP_OUTPUT_FILE" "$OUTPUT_FILE"
    echo "Prometheus configuration file updated at: $OUTPUT_FILE"
    exit 2
  fi
else
  # If the file doesn't exist, just create it
  mv "$TEMP_OUTPUT_FILE" "$OUTPUT_FILE"
  echo "Prometheus configuration file created at: $OUTPUT_FILE"
  exit 2
fi
