#!/usr/bin/env bash
#
# Webtop / Selkies-GStreamer Prometheus Metrics Exporter
#
# Prometheus textfile collector exporter for Webtop (LinuxServer.io)
# and Selkies-GStreamer remote desktop containers. Uses Docker inspect,
# Docker stats, and container introspection to collect container state,
# resource usage, session counts, and health metrics.
#
# Usage:
#   ./webtop-selkies-exporter.sh
#   ./webtop-selkies-exporter.sh --textfile
#   WEBTOP_HOST=10.0.0.5 WEBTOP_PORT=3000 ./webtop-selkies-exporter.sh --textfile
#   CONTAINER_PATTERN="selkies" ./webtop-selkies-exporter.sh --install
#
# Parameters:
#   --textfile   Write to textfile collector directory
#   --install    Create cron job for automatic collection
#   --help       Show usage
#
# Environment:
#   WEBTOP_HOST        Container host for HTTP checks (default: localhost)
#   WEBTOP_PORT        Container web port for HTTP checks (default: 3000)
#   CONTAINER_PATTERN  Docker container name pattern to match (default: webtop)
#   TEXTFILE_DIR       Textfile collector directory (default: /var/lib/node_exporter/textfile_collector)
#   CURL_TIMEOUT       HTTP request timeout in seconds (default: 5)
#
# Author:  Phil Connor
# Contact: contact@mylinux.work
# Website: https://mylinux.work
# License: MIT
# Version: 1.0
#
# Metrics Exported:
#   Core:
#     - webtop_up
#     - webtop_exporter_info{version}
#     - webtop_container_state{name,image}
#
#   Container Resources:
#     - webtop_cpu_usage_percent{name}
#     - webtop_memory_usage_bytes{name}
#     - webtop_memory_limit_bytes{name}
#     - webtop_network_rx_bytes{name}
#     - webtop_network_tx_bytes{name}
#
#   Sessions:
#     - webtop_vnc_sessions_active{name}
#     - webtop_http_connections_active{name}
#
#   Health:
#     - webtop_container_restarts{name}
#     - webtop_container_uptime_seconds{name}
#     - webtop_health_check_status{name}
#
#   Exporter:
#     - webtop_exporter_duration_seconds
#     - webtop_exporter_last_run_timestamp

set -euo pipefail

# --- Configuration ---
readonly VERSION="1.0"
readonly SCRIPT_NAME="$(basename "$0")"
WEBTOP_HOST="${WEBTOP_HOST:-localhost}"
WEBTOP_PORT="${WEBTOP_PORT:-3000}"
CONTAINER_PATTERN="${CONTAINER_PATTERN:-webtop}"
TEXTFILE_DIR="${TEXTFILE_DIR:-/var/lib/node_exporter/textfile_collector}"
CURL_TIMEOUT="${CURL_TIMEOUT:-5}"
TEXTFILE_MODE=false
OUTPUT=""
START_TIME=""

# --- Functions ---

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

Webtop / Selkies-GStreamer Prometheus Metrics Exporter

Options:
  --textfile    Write metrics to textfile collector directory
  --install     Create cron job for automatic collection
  --help        Show this help message

Environment Variables:
  WEBTOP_HOST        Container host for HTTP checks (default: localhost)
  WEBTOP_PORT        Container web port for HTTP checks (default: 3000)
  CONTAINER_PATTERN  Docker container name pattern to match (default: webtop)
  TEXTFILE_DIR       Output directory (default: /var/lib/node_exporter/textfile_collector)
  CURL_TIMEOUT       Request timeout in seconds (default: 5)

Examples:
  $SCRIPT_NAME
  $SCRIPT_NAME --textfile
  WEBTOP_HOST=10.0.0.5 WEBTOP_PORT=3000 $SCRIPT_NAME --textfile
  CONTAINER_PATTERN="selkies" $SCRIPT_NAME --install
EOF
    exit 0
}

check_dependencies() {
    local missing=()
    for cmd in docker jq; do
        if ! command -v "$cmd" &>/dev/null; then
            missing+=("$cmd")
        fi
    done
    if [[ ${#missing[@]} -gt 0 ]]; then
        echo "ERROR: Missing required commands: ${missing[*]}" >&2
        echo "Install with: apt install ${missing[*]}  OR  dnf install ${missing[*]}" >&2
        exit 1
    fi
}

validate_config() {
    if ! docker info &>/dev/null; then
        echo "ERROR: Cannot connect to Docker daemon. Is Docker running? Do you have permission?" >&2
        exit 1
    fi
}

api_get() {
    local url="$1"
    curl -sf --max-time "$CURL_TIMEOUT" "$url" 2>/dev/null || echo ""
}

add_metric() {
    local name="$1"
    local type="$2"
    local help="$3"
    local value="$4"
    local labels="${5:-}"

    if [[ -n "$labels" ]]; then
        OUTPUT+="# HELP ${name} ${help}
# TYPE ${name} ${type}
${name}{${labels}} ${value}
"
    else
        OUTPUT+="# HELP ${name} ${help}
# TYPE ${name} ${type}
${name} ${value}
"
    fi
}

add_metric_value() {
    local name="$1"
    local value="$2"
    local labels="${3:-}"

    if [[ -n "$labels" ]]; then
        OUTPUT+="${name}{${labels}} ${value}
"
    else
        OUTPUT+="${name} ${value}
"
    fi
}

# Discover containers matching the pattern.
# Returns newline-separated container IDs.
discover_containers() {
    docker ps -a --filter "name=${CONTAINER_PATTERN}" --format '{{.ID}}' 2>/dev/null
}

# Collect container state and basic info.
# Sets webtop_up to 1 if at least one matching container is running.
collect_container_state() {
    local container_ids="$1"
    local any_running=0

    OUTPUT+="# HELP webtop_container_state Container running state (1=running, 0=stopped)
# TYPE webtop_container_state gauge
"

    while IFS= read -r cid; do
        [[ -z "$cid" ]] && continue

        local inspect_json
        inspect_json=$(docker inspect "$cid" 2>/dev/null) || continue

        local name image state
        name=$(echo "$inspect_json" | jq -r '.[0].Name // ""' | sed 's|^/||')
        image=$(echo "$inspect_json" | jq -r '.[0].Config.Image // ""')
        state=$(echo "$inspect_json" | jq -r '.[0].State.Status // ""')

        local state_val=0
        if [[ "$state" == "running" ]]; then
            state_val=1
            any_running=1
        fi

        add_metric_value "webtop_container_state" "$state_val" "name=\"${name}\",image=\"${image}\""
    done <<< "$container_ids"

    add_metric "webtop_up" "gauge" "Whether any matching container is running (1=up, 0=down)" "$any_running"
}

# Collect CPU and memory from docker stats.
collect_resource_metrics() {
    local container_ids="$1"

    local stats_json
    stats_json=$(docker stats --no-stream --format '{"name":"{{.Name}}","cpu":"{{.CPUPerc}}","mem_usage":"{{.MemUsage}}","net_io":"{{.NetIO}}"}' $(echo "$container_ids" | tr '\n' ' ') 2>/dev/null) || return

    [[ -z "$stats_json" ]] && return

    OUTPUT+="# HELP webtop_cpu_usage_percent Container CPU usage percentage
# TYPE webtop_cpu_usage_percent gauge
"
    while IFS= read -r line; do
        [[ -z "$line" ]] && continue
        local name cpu_str
        name=$(echo "$line" | jq -r '.name // ""')
        cpu_str=$(echo "$line" | jq -r '.cpu // "0"')
        local cpu_val
        cpu_val=$(echo "$cpu_str" | grep -oP '[\d.]+' | head -1 || echo "0")
        add_metric_value "webtop_cpu_usage_percent" "$cpu_val" "name=\"${name}\""
    done <<< "$stats_json"

    OUTPUT+="# HELP webtop_memory_usage_bytes Container memory usage in bytes
# TYPE webtop_memory_usage_bytes gauge
"
    while IFS= read -r line; do
        [[ -z "$line" ]] && continue
        local name mem_str
        name=$(echo "$line" | jq -r '.name // ""')
        mem_str=$(echo "$line" | jq -r '.mem_usage // ""')
        # Format: "123.4MiB / 1.5GiB"
        local usage_part
        usage_part=$(echo "$mem_str" | awk -F' / ' '{print $1}')
        local usage_bytes
        usage_bytes=$(parse_docker_size "$usage_part")
        add_metric_value "webtop_memory_usage_bytes" "$usage_bytes" "name=\"${name}\""
    done <<< "$stats_json"

    OUTPUT+="# HELP webtop_memory_limit_bytes Container memory limit in bytes
# TYPE webtop_memory_limit_bytes gauge
"
    while IFS= read -r line; do
        [[ -z "$line" ]] && continue
        local name mem_str
        name=$(echo "$line" | jq -r '.name // ""')
        mem_str=$(echo "$line" | jq -r '.mem_usage // ""')
        local limit_part
        limit_part=$(echo "$mem_str" | awk -F' / ' '{print $2}')
        local limit_bytes
        limit_bytes=$(parse_docker_size "$limit_part")
        add_metric_value "webtop_memory_limit_bytes" "$limit_bytes" "name=\"${name}\""
    done <<< "$stats_json"

    OUTPUT+="# HELP webtop_network_rx_bytes Network bytes received
# TYPE webtop_network_rx_bytes counter
"
    while IFS= read -r line; do
        [[ -z "$line" ]] && continue
        local name net_str
        name=$(echo "$line" | jq -r '.name // ""')
        net_str=$(echo "$line" | jq -r '.net_io // ""')
        # Format: "1.23kB / 4.56MB"
        local rx_part
        rx_part=$(echo "$net_str" | awk -F' / ' '{print $1}')
        local rx_bytes
        rx_bytes=$(parse_docker_size "$rx_part")
        add_metric_value "webtop_network_rx_bytes" "$rx_bytes" "name=\"${name}\""
    done <<< "$stats_json"

    OUTPUT+="# HELP webtop_network_tx_bytes Network bytes transmitted
# TYPE webtop_network_tx_bytes counter
"
    while IFS= read -r line; do
        [[ -z "$line" ]] && continue
        local name net_str
        name=$(echo "$line" | jq -r '.name // ""')
        net_str=$(echo "$line" | jq -r '.net_io // ""')
        local tx_part
        tx_part=$(echo "$net_str" | awk -F' / ' '{print $2}')
        local tx_bytes
        tx_bytes=$(parse_docker_size "$tx_part")
        add_metric_value "webtop_network_tx_bytes" "$tx_bytes" "name=\"${name}\""
    done <<< "$stats_json"
}

# Parse Docker size strings like "1.23GiB", "456MiB", "789kB" to bytes.
parse_docker_size() {
    local size_str="$1"

    if [[ -z "$size_str" ]]; then
        echo "0"
        return
    fi

    local number unit
    number=$(echo "$size_str" | grep -oP '[\d.]+' | head -1)
    unit=$(echo "$size_str" | grep -oP '[A-Za-z]+' | head -1)

    if [[ -z "$number" ]]; then
        echo "0"
        return
    fi

    case "$unit" in
        B)
            echo "$number" | awk '{printf "%.0f", $1}' ;;
        kB|KB)
            echo "$number" | awk '{printf "%.0f", $1 * 1000}' ;;
        KiB)
            echo "$number" | awk '{printf "%.0f", $1 * 1024}' ;;
        MB)
            echo "$number" | awk '{printf "%.0f", $1 * 1000000}' ;;
        MiB)
            echo "$number" | awk '{printf "%.0f", $1 * 1048576}' ;;
        GB)
            echo "$number" | awk '{printf "%.0f", $1 * 1000000000}' ;;
        GiB)
            echo "$number" | awk '{printf "%.0f", $1 * 1073741824}' ;;
        TB)
            echo "$number" | awk '{printf "%.0f", $1 * 1000000000000}' ;;
        TiB)
            echo "$number" | awk '{printf "%.0f", $1 * 1099511627776}' ;;
        *)
            echo "$number" | awk '{printf "%.0f", $1}' ;;
    esac
}

# Collect VNC session count and HTTP connection count.
collect_session_metrics() {
    local container_ids="$1"

    OUTPUT+="# HELP webtop_vnc_sessions_active Active VNC/KasmVNC sessions
# TYPE webtop_vnc_sessions_active gauge
"

    while IFS= read -r cid; do
        [[ -z "$cid" ]] && continue

        local name
        name=$(docker inspect --format '{{.Name}}' "$cid" 2>/dev/null | sed 's|^/||')
        [[ -z "$name" ]] && continue

        local state
        state=$(docker inspect --format '{{.State.Status}}' "$cid" 2>/dev/null)
        [[ "$state" != "running" ]] && continue

        local vnc_sessions=0

        # Method 1: Try KasmVNC API
        local kasmvnc_response
        kasmvnc_response=$(api_get "http://${WEBTOP_HOST}:${WEBTOP_PORT}/api/get_session")
        if [[ -n "$kasmvnc_response" ]]; then
            local session_count
            session_count=$(echo "$kasmvnc_response" | jq 'if type == "array" then length elif type == "object" then 1 else 0 end' 2>/dev/null || echo "0")
            vnc_sessions="$session_count"
        else
            # Method 2: Parse container logs for VNC connection events
            local recent_connects recent_disconnects
            recent_connects=$(docker logs --since 24h "$cid" 2>&1 | grep -ciE '(connection from|client connected|new client|websocket.*opened)' || true)
            recent_disconnects=$(docker logs --since 24h "$cid" 2>&1 | grep -ciE '(client disconnected|client gone|websocket.*closed|connection closed)' || true)
            vnc_sessions=$(( recent_connects - recent_disconnects ))
            if [[ "$vnc_sessions" -lt 0 ]]; then
                vnc_sessions=0
            fi
        fi

        add_metric_value "webtop_vnc_sessions_active" "$vnc_sessions" "name=\"${name}\""
    done <<< "$container_ids"

    OUTPUT+="# HELP webtop_http_connections_active Active HTTP connections to container port
# TYPE webtop_http_connections_active gauge
"

    while IFS= read -r cid; do
        [[ -z "$cid" ]] && continue

        local name
        name=$(docker inspect --format '{{.Name}}' "$cid" 2>/dev/null | sed 's|^/||')
        [[ -z "$name" ]] && continue

        local state
        state=$(docker inspect --format '{{.State.Status}}' "$cid" 2>/dev/null)
        [[ "$state" != "running" ]] && continue

        local http_conns=0

        # Get the host port mapped to the container port
        local host_port
        host_port=$(docker inspect --format "{{(index (index .NetworkSettings.Ports \"${WEBTOP_PORT}/tcp\") 0).HostPort}}" "$cid" 2>/dev/null || echo "$WEBTOP_PORT")

        if command -v ss &>/dev/null; then
            http_conns=$(ss -tn state established "( dport = :${host_port} or sport = :${host_port} )" 2>/dev/null | tail -n +2 | wc -l || echo "0")
        elif command -v netstat &>/dev/null; then
            http_conns=$(netstat -tn 2>/dev/null | grep -c ":${host_port}.*ESTABLISHED" || echo "0")
        fi

        add_metric_value "webtop_http_connections_active" "$http_conns" "name=\"${name}\""
    done <<< "$container_ids"
}

# Collect restart count, uptime, and health check status.
collect_health_metrics() {
    local container_ids="$1"

    OUTPUT+="# HELP webtop_container_restarts Container restart count
# TYPE webtop_container_restarts counter
"
    while IFS= read -r cid; do
        [[ -z "$cid" ]] && continue

        local inspect_json
        inspect_json=$(docker inspect "$cid" 2>/dev/null) || continue

        local name restart_count
        name=$(echo "$inspect_json" | jq -r '.[0].Name // ""' | sed 's|^/||')
        restart_count=$(echo "$inspect_json" | jq -r '.[0].RestartCount // 0')

        add_metric_value "webtop_container_restarts" "$restart_count" "name=\"${name}\""
    done <<< "$container_ids"

    OUTPUT+="# HELP webtop_container_uptime_seconds Seconds since container started
# TYPE webtop_container_uptime_seconds gauge
"
    while IFS= read -r cid; do
        [[ -z "$cid" ]] && continue

        local inspect_json
        inspect_json=$(docker inspect "$cid" 2>/dev/null) || continue

        local name started_at state
        name=$(echo "$inspect_json" | jq -r '.[0].Name // ""' | sed 's|^/||')
        state=$(echo "$inspect_json" | jq -r '.[0].State.Status // ""')
        started_at=$(echo "$inspect_json" | jq -r '.[0].State.StartedAt // ""')

        local uptime=0
        if [[ "$state" == "running" && -n "$started_at" && "$started_at" != "0001-01-01T00:00:00Z" ]]; then
            local started_epoch now_epoch
            started_epoch=$(date -d "$started_at" +%s 2>/dev/null || echo "0")
            now_epoch=$(date +%s)
            if [[ "$started_epoch" -gt 0 ]]; then
                uptime=$(( now_epoch - started_epoch ))
            fi
        fi

        add_metric_value "webtop_container_uptime_seconds" "$uptime" "name=\"${name}\""
    done <<< "$container_ids"

    OUTPUT+="# HELP webtop_health_check_status Docker health check status (1=healthy, 0=unhealthy)
# TYPE webtop_health_check_status gauge
"
    while IFS= read -r cid; do
        [[ -z "$cid" ]] && continue

        local inspect_json
        inspect_json=$(docker inspect "$cid" 2>/dev/null) || continue

        local name health_status
        name=$(echo "$inspect_json" | jq -r '.[0].Name // ""' | sed 's|^/||')
        health_status=$(echo "$inspect_json" | jq -r '.[0].State.Health.Status // "none"')

        local health_val=0
        if [[ "$health_status" == "healthy" ]]; then
            health_val=1
        fi

        add_metric_value "webtop_health_check_status" "$health_val" "name=\"${name}\""
    done <<< "$container_ids"
}

write_output() {
    if [[ "$TEXTFILE_MODE" == true ]]; then
        local output_file="${TEXTFILE_DIR}/webtop.prom"
        local temp_file="${output_file}.$$"

        mkdir -p "$TEXTFILE_DIR"
        echo "$OUTPUT" > "$temp_file"
        mv "$temp_file" "$output_file"
    else
        echo "$OUTPUT"
    fi
}

install_cron() {
    if [[ $EUID -ne 0 ]]; then
        echo "ERROR: --install requires root" >&2
        exit 1
    fi

    local script_path
    script_path=$(readlink -f "$0")

    cat > /etc/cron.d/webtop-selkies-exporter <<EOF
# Webtop/Selkies Prometheus Exporter — runs every 2 minutes
WEBTOP_HOST=${WEBTOP_HOST}
WEBTOP_PORT=${WEBTOP_PORT}
CONTAINER_PATTERN=${CONTAINER_PATTERN}
TEXTFILE_DIR=${TEXTFILE_DIR}
*/2 * * * * root ${script_path} --textfile 2>/dev/null
EOF

    chmod 644 /etc/cron.d/webtop-selkies-exporter
    echo "Installed cron job: /etc/cron.d/webtop-selkies-exporter"
    echo "Metrics will be written to: ${TEXTFILE_DIR}/webtop.prom"
}

# --- Main ---

main() {
    # Parse arguments
    for arg in "$@"; do
        case "$arg" in
            --textfile) TEXTFILE_MODE=true ;;
            --install)
                check_dependencies
                validate_config
                install_cron
                exit 0
                ;;
            --help|-h) usage ;;
            *) echo "Unknown option: $arg" >&2; usage ;;
        esac
    done

    check_dependencies
    validate_config

    START_TIME=$(date +%s%N)

    # Exporter info
    add_metric "webtop_exporter_info" "gauge" "Exporter version information" "1" "version=\"${VERSION}\""

    # Discover matching containers
    local container_ids
    container_ids=$(discover_containers)

    if [[ -z "$container_ids" ]]; then
        add_metric "webtop_up" "gauge" "Whether any matching container is running (1=up, 0=down)" "0"
    else
        collect_container_state "$container_ids"

        # Only collect detailed metrics for running containers
        local running_ids
        running_ids=$(docker ps --filter "name=${CONTAINER_PATTERN}" --filter "status=running" --format '{{.ID}}' 2>/dev/null)

        if [[ -n "$running_ids" ]]; then
            collect_resource_metrics "$running_ids"
            collect_session_metrics "$running_ids"
        fi

        collect_health_metrics "$container_ids"
    fi

    # Exporter performance
    local end_time duration
    end_time=$(date +%s%N)
    duration=$(echo "scale=2; ($end_time - $START_TIME) / 1000000000" | bc 2>/dev/null || echo "0")
    add_metric "webtop_exporter_duration_seconds" "gauge" "Time to generate all metrics" "$duration"
    add_metric "webtop_exporter_last_run_timestamp" "gauge" "Unix timestamp of last successful run" "$(date +%s)"

    write_output
}

main "$@"