#!/bin/bash
################################################################################
# Script Name: incus-metrics-exporter.sh
# Version: 1.0
# Description: Prometheus exporter for Incus container/VM metrics not covered
#              by the built-in /1.0/metrics endpoint — storage pool usage,
#              snapshot counts and age, instance inventory, and image cache.
#              Designed for node_exporter textfile collector so all Incus data
#              flows through a single Prometheus scrape target.
#
# Author: Phil Connor
# Contact: contact@mylinux.work
# Website: https://mylinux.work
# License: MIT
#
# Prerequisites:
#   - Incus installed and running
#   - User in incus/lxd group (or root)
#   - jq for JSON parsing
#   - netcat (nc) for HTTP mode
#
# Usage:
#   ./incus-metrics-exporter.sh                     # stdout
#   ./incus-metrics-exporter.sh --textfile          # node_exporter textfile
#   ./incus-metrics-exporter.sh --http -p 9199      # HTTP server
#
################################################################################

# ============================================================================
# CONFIGURATION VARIABLES
# ============================================================================

TEXTFILE_DIR="/var/lib/node_exporter"
OUTPUT_FILE=""
HTTP_MODE=false
HTTP_PORT=9199

# ============================================================================
# HELPER FUNCTIONS
# ============================================================================

prom_escape() {
    local s="$1"
    s=${s//\\/\\\\}
    s=${s//\"/\\\"}
    s=${s//$'\n'/\\n}
    printf '%s\n' "$s"
}

show_usage() {
    cat <<EOF
Usage: $0 [OPTIONS]

Export Incus supplemental metrics as Prometheus metrics (v1.0).

MODES:
    --textfile      Write to node_exporter textfile collector
    --http          Run HTTP server on port $HTTP_PORT

OPTIONS:
    -p, --port          HTTP port (default: 9199)
    -o, --output        Output file path

EXAMPLES:
    $0 --textfile                              # Write to textfile collector
    $0 --http --port 9199                      # Run HTTP server
    $0 -o /tmp/incus_metrics.prom              # Write to custom file

EOF
    exit 0
}

parse_args() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help) show_usage ;;
            --textfile) OUTPUT_FILE="$TEXTFILE_DIR/incus_metrics_exporter.prom"; shift ;;
            --http) HTTP_MODE=true; shift ;;
            -p|--port) HTTP_PORT="$2"; shift 2 ;;
            -o|--output) OUTPUT_FILE="$2"; shift 2 ;;
            *) echo "Unknown option: $1" >&2; exit 1 ;;
        esac
    done
}

# ============================================================================
# METRICS GENERATION
# ============================================================================

generate_metrics() {
    local script_start
    script_start=$(date +%s)

    # ========================================================================
    # Prerequisite Check
    # ========================================================================

    if ! command -v jq >/dev/null 2>&1; then
        echo "# ERROR: jq is required but not found" >&2
        cat <<EOF
# HELP incus_exporter_up Exporter status (1=up, 0=down)
# TYPE incus_exporter_up gauge
incus_exporter_up 0

# HELP incus_exporter_info Exporter version information
# TYPE incus_exporter_info gauge
incus_exporter_info{version="1.0"} 1

EOF
        return
    fi

    if ! command -v incus >/dev/null 2>&1; then
        echo "# ERROR: incus command not found" >&2
        cat <<EOF
# HELP incus_exporter_up Exporter status (1=up, 0=down)
# TYPE incus_exporter_up gauge
incus_exporter_up 0

# HELP incus_exporter_info Exporter version information
# TYPE incus_exporter_info gauge
incus_exporter_info{version="1.0"} 1

EOF
        return
    fi

    cat <<EOF
# HELP incus_exporter_up Exporter status (1=up, 0=down)
# TYPE incus_exporter_up gauge
incus_exporter_up 1

# HELP incus_exporter_info Exporter version information
# TYPE incus_exporter_info gauge
incus_exporter_info{version="1.0"} 1

EOF

    # ========================================================================
    # Storage Pool Metrics
    # ========================================================================
    local pool_total_lines=""
    local pool_used_lines=""
    local pool_ratio_lines=""

    local pools_json
    pools_json=$(incus query /1.0/storage-pools 2>/dev/null)

    if [ -n "$pools_json" ] && [ "$pools_json" != "null" ]; then
        local pool_urls
        pool_urls=$(echo "$pools_json" | jq -r '.[]' 2>/dev/null)

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

            local pool_name
            pool_name=$(basename "$pool_url")
            local esc_pool
            esc_pool=$(prom_escape "$pool_name")

            local pool_config
            pool_config=$(incus query "/1.0/storage-pools/$pool_name" 2>/dev/null)
            local driver
            driver=$(echo "$pool_config" | jq -r '.driver // "unknown"' 2>/dev/null)
            local esc_driver
            esc_driver=$(prom_escape "$driver")

            local resources_json
            resources_json=$(incus query "/1.0/storage-pools/$pool_name/resources" 2>/dev/null)

            local total_space=""
            local used_space=""

            if [ -n "$resources_json" ] && [ "$resources_json" != "null" ]; then
                total_space=$(echo "$resources_json" | jq -r '.space.total // empty' 2>/dev/null)
                used_space=$(echo "$resources_json" | jq -r '.space.used // empty' 2>/dev/null)
            fi

            # Fallback: parse incus storage info output
            if [ -z "$total_space" ] || [ -z "$used_space" ]; then
                local storage_info
                storage_info=$(incus storage info "$pool_name" 2>/dev/null)

                if [ -n "$storage_info" ]; then
                    total_space=$(echo "$storage_info" | awk '/Total space:/ {
                        val = $3; unit = $4
                        if (unit == "GiB") val = val * 1073741824
                        else if (unit == "TiB") val = val * 1099511627776
                        else if (unit == "MiB") val = val * 1048576
                        else if (unit == "KiB") val = val * 1024
                        printf "%.0f", val
                    }')
                    used_space=$(echo "$storage_info" | awk '/Space used:/ {
                        val = $3; unit = $4
                        if (unit == "GiB") val = val * 1073741824
                        else if (unit == "TiB") val = val * 1099511627776
                        else if (unit == "MiB") val = val * 1048576
                        else if (unit == "KiB") val = val * 1024
                        printf "%.0f", val
                    }')
                fi
            fi

            if [ -n "$total_space" ] && [ -n "$used_space" ]; then
                local labels="pool=\"$esc_pool\",driver=\"$esc_driver\""
                pool_total_lines="${pool_total_lines}incus_storage_pool_total_bytes{${labels}} $total_space
"
                pool_used_lines="${pool_used_lines}incus_storage_pool_used_bytes{${labels}} $used_space
"
                local ratio
                ratio=$(awk "BEGIN { if ($total_space > 0) printf \"%.6f\", $used_space / $total_space; else print \"0\" }")
                pool_ratio_lines="${pool_ratio_lines}incus_storage_pool_usage_ratio{${labels}} $ratio
"
            else
                echo "# WARNING: could not read storage pool '$pool_name' resources" >&2
            fi
        done <<< "$pool_urls"
    fi

    if [ -n "$pool_total_lines" ]; then
        echo "# HELP incus_storage_pool_total_bytes Total storage pool capacity in bytes"
        echo "# TYPE incus_storage_pool_total_bytes gauge"
        printf '%s' "$pool_total_lines"
        echo ""
    fi

    if [ -n "$pool_used_lines" ]; then
        echo "# HELP incus_storage_pool_used_bytes Used storage pool space in bytes"
        echo "# TYPE incus_storage_pool_used_bytes gauge"
        printf '%s' "$pool_used_lines"
        echo ""
    fi

    if [ -n "$pool_ratio_lines" ]; then
        echo "# HELP incus_storage_pool_usage_ratio Storage pool used/total ratio"
        echo "# TYPE incus_storage_pool_usage_ratio gauge"
        printf '%s' "$pool_ratio_lines"
        echo ""
    fi

    # ========================================================================
    # Instance Inventory Metrics
    # ========================================================================
    local instances_json
    instances_json=$(incus list --format json 2>/dev/null)

    local instance_total_lines=""
    local instance_info_lines=""
    local snapshot_total_lines=""
    local snapshot_oldest_lines=""
    local snapshot_newest_lines=""

    if [ -n "$instances_json" ] && [ "$instances_json" != "null" ] && [ "$instances_json" != "[]" ]; then
        # Count instances by type and status
        local counts
        counts=$(echo "$instances_json" | jq -r '
            group_by(.type, .status) |
            .[] |
            {type: .[0].type, status: .[0].status, count: length} |
            "\(.type) \(.status) \(.count)"
        ' 2>/dev/null)

        while IFS= read -r count_line; do
            [ -z "$count_line" ] && continue
            local inst_type inst_status inst_count
            inst_type=$(echo "$count_line" | awk '{print $1}')
            inst_status=$(echo "$count_line" | awk '{print $2}')
            inst_count=$(echo "$count_line" | awk '{print $3}')
            local esc_type esc_status
            esc_type=$(prom_escape "$inst_type")
            esc_status=$(prom_escape "$inst_status")
            instance_total_lines="${instance_total_lines}incus_instances_total{type=\"$esc_type\",status=\"$esc_status\"} $inst_count
"
        done <<< "$counts"

        # Per-instance info gauge and snapshot metrics
        local instance_data
        instance_data=$(echo "$instances_json" | jq -r '
            .[] |
            "\(.name)\t\(.type)\t\(.status)\t\(.project // "default")"
        ' 2>/dev/null)

        local now
        now=$(date +%s)

        while IFS=$'\t' read -r inst_name inst_type inst_status inst_project; do
            [ -z "$inst_name" ] && continue

            local esc_name esc_type esc_status esc_project
            esc_name=$(prom_escape "$inst_name")
            esc_type=$(prom_escape "$inst_type")
            esc_status=$(prom_escape "$inst_status")
            esc_project=$(prom_escape "$inst_project")

            instance_info_lines="${instance_info_lines}incus_instance_info{name=\"$esc_name\",type=\"$esc_type\",status=\"$esc_status\",project=\"$esc_project\"} 1
"

            # Snapshot metrics
            local snap_json
            snap_json=$(incus query "/1.0/instances/$inst_name/snapshots" --project "$inst_project" 2>/dev/null)

            if [ -n "$snap_json" ] && [ "$snap_json" != "null" ] && [ "$snap_json" != "[]" ]; then
                local snap_count
                snap_count=$(echo "$snap_json" | jq 'length' 2>/dev/null)
                snap_count=${snap_count:-0}

                local snap_labels="name=\"$esc_name\",project=\"$esc_project\""
                snapshot_total_lines="${snapshot_total_lines}incus_instance_snapshots_total{${snap_labels}} $snap_count
"

                if [ "$snap_count" -gt 0 ] 2>/dev/null; then
                    # Fetch details for each snapshot to get created_at timestamps
                    local snap_timestamps=""
                    local snap_urls
                    snap_urls=$(echo "$snap_json" | jq -r '.[]' 2>/dev/null)

                    while IFS= read -r snap_url; do
                        [ -z "$snap_url" ] && continue
                        local snap_detail
                        snap_detail=$(incus query "$snap_url" --project "$inst_project" 2>/dev/null)
                        if [ -n "$snap_detail" ]; then
                            local created_at
                            created_at=$(echo "$snap_detail" | jq -r '.created_at // empty' 2>/dev/null)
                            if [ -n "$created_at" ]; then
                                local snap_epoch
                                snap_epoch=$(date -d "$created_at" +%s 2>/dev/null)
                                if [ -n "$snap_epoch" ]; then
                                    snap_timestamps="${snap_timestamps}${snap_epoch}
"
                                fi
                            fi
                        fi
                    done <<< "$snap_urls"

                    if [ -n "$snap_timestamps" ]; then
                        local oldest newest
                        oldest=$(echo "$snap_timestamps" | sort -n | head -1)
                        newest=$(echo "$snap_timestamps" | sort -rn | head -1)

                        if [ -n "$oldest" ]; then
                            local oldest_age
                            oldest_age=$((now - oldest))
                            snapshot_oldest_lines="${snapshot_oldest_lines}incus_instance_snapshot_oldest_age_seconds{${snap_labels}} $oldest_age
"
                        fi

                        if [ -n "$newest" ]; then
                            local newest_age
                            newest_age=$((now - newest))
                            snapshot_newest_lines="${snapshot_newest_lines}incus_instance_snapshot_newest_age_seconds{${snap_labels}} $newest_age
"
                        fi
                    fi
                fi
            else
                local snap_labels="name=\"$esc_name\",project=\"$esc_project\""
                snapshot_total_lines="${snapshot_total_lines}incus_instance_snapshots_total{${snap_labels}} 0
"
            fi
        done <<< "$instance_data"
    fi

    if [ -n "$instance_total_lines" ]; then
        echo "# HELP incus_instances_total Number of instances by type and status"
        echo "# TYPE incus_instances_total gauge"
        printf '%s' "$instance_total_lines"
        echo ""
    fi

    if [ -n "$instance_info_lines" ]; then
        echo "# HELP incus_instance_info Instance information (always 1)"
        echo "# TYPE incus_instance_info gauge"
        printf '%s' "$instance_info_lines"
        echo ""
    fi

    if [ -n "$snapshot_total_lines" ]; then
        echo "# HELP incus_instance_snapshots_total Number of snapshots per instance"
        echo "# TYPE incus_instance_snapshots_total gauge"
        printf '%s' "$snapshot_total_lines"
        echo ""
    fi

    if [ -n "$snapshot_oldest_lines" ]; then
        echo "# HELP incus_instance_snapshot_oldest_age_seconds Age of oldest snapshot in seconds"
        echo "# TYPE incus_instance_snapshot_oldest_age_seconds gauge"
        printf '%s' "$snapshot_oldest_lines"
        echo ""
    fi

    if [ -n "$snapshot_newest_lines" ]; then
        echo "# HELP incus_instance_snapshot_newest_age_seconds Age of newest snapshot in seconds"
        echo "# TYPE incus_instance_snapshot_newest_age_seconds gauge"
        printf '%s' "$snapshot_newest_lines"
        echo ""
    fi

    # ========================================================================
    # Image Cache Metrics
    # ========================================================================
    local images_json
    images_json=$(incus image list --format json 2>/dev/null)

    local images_total=0
    local image_size_lines=""

    if [ -n "$images_json" ] && [ "$images_json" != "null" ] && [ "$images_json" != "[]" ]; then
        images_total=$(echo "$images_json" | jq 'length' 2>/dev/null)
        images_total=${images_total:-0}

        local image_data
        image_data=$(echo "$images_json" | jq -r '
            .[] |
            "\(.fingerprint[0:12])\t\(((.aliases // []) | if length > 0 then .[0].name else "" end))\t\(.architecture)\t\(.size)"
        ' 2>/dev/null)

        while IFS=$'\t' read -r img_fp img_alias img_arch img_size; do
            [ -z "$img_fp" ] && continue
            local esc_fp esc_alias esc_arch
            esc_fp=$(prom_escape "$img_fp")
            esc_alias=$(prom_escape "$img_alias")
            esc_arch=$(prom_escape "$img_arch")
            image_size_lines="${image_size_lines}incus_image_size_bytes{fingerprint=\"$esc_fp\",alias=\"$esc_alias\",arch=\"$esc_arch\"} $img_size
"
        done <<< "$image_data"
    fi

    echo "# HELP incus_images_total Total number of cached images"
    echo "# TYPE incus_images_total gauge"
    echo "incus_images_total $images_total"
    echo ""

    if [ -n "$image_size_lines" ]; then
        echo "# HELP incus_image_size_bytes Image size in bytes"
        echo "# TYPE incus_image_size_bytes gauge"
        printf '%s' "$image_size_lines"
        echo ""
    fi

    # ========================================================================
    # Built-in Incus Metrics (/1.0/metrics)
    # ========================================================================
    # Pull the native metrics (CPU, memory, network, disk I/O per instance)
    # so everything flows through node_exporter — no second scrape target.
    local builtin_metrics
    builtin_metrics=$(incus query /1.0/metrics 2>/dev/null)

    if [ -n "$builtin_metrics" ]; then
        echo "# ── Built-in Incus instance metrics (/1.0/metrics) ──"
        echo "$builtin_metrics"
        echo ""
    else
        echo "# WARNING: could not read /1.0/metrics (check incus daemon status)" >&2
    fi

    # ========================================================================
    # Exporter Runtime
    # ========================================================================
    local script_end script_duration
    script_end=$(date +%s)
    script_duration=$((script_end - script_start))

    cat <<EOF
# HELP incus_exporter_duration_seconds Time to generate all metrics
# TYPE incus_exporter_duration_seconds gauge
incus_exporter_duration_seconds $script_duration

# HELP incus_exporter_last_run_timestamp Unix timestamp of last successful run
# TYPE incus_exporter_last_run_timestamp gauge
incus_exporter_last_run_timestamp $script_end
EOF

    echo ""
}

# ============================================================================
# HTTP SERVER MODE
# ============================================================================

run_http_server() {
    echo "Starting Incus metrics exporter on port $HTTP_PORT..." >&2

    if ! command -v nc >/dev/null 2>&1; then
        echo "ERROR: netcat (nc) required for HTTP mode" >&2
        exit 1
    fi

    trap 'echo "Shutting down Incus metrics exporter..." >&2; exit 0' INT TERM

    while true; do
        {
            read -r request
            local body
            if [[ "$request" =~ ^GET\ /metrics ]]; then
                body=$(generate_metrics)
                printf "HTTP/1.1 200 OK\r\nContent-Type: text/plain; version=0.0.4\r\nContent-Length: %d\r\nConnection: close\r\n\r\n%s" "${#body}" "$body"
            else
                body=$(cat <<'HTMLEOF'
<!DOCTYPE html>
<html>
<head><title>Incus Metrics Exporter v1.0</title></head>
<body>
<h1>Incus Metrics Exporter v1.0</h1>
<p><a href="/metrics">Metrics</a></p>
<h2>Sections</h2>
<ul>
<li>Storage pool usage (total, used, ratio)</li>
<li>Instance inventory (count by type/status, info per instance)</li>
<li>Snapshot counts and age per instance</li>
<li>Image cache (count and size per image)</li>
</ul>
</body>
</html>
HTMLEOF
)
                printf "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\nContent-Length: %d\r\nConnection: close\r\n\r\n%s" "${#body}" "$body"
            fi
        } | if nc -h 2>&1 | grep -q 'GNU\|traditional'; then
            nc -l -p "$HTTP_PORT" -q 1 2>/dev/null
        else
            nc -l "$HTTP_PORT" 2>/dev/null
        fi
    done
}

# ============================================================================
# MAIN EXECUTION
# ============================================================================

main() {
    parse_args "$@"

    if [ "$HTTP_MODE" = true ]; then
        run_http_server
    elif [ -n "$OUTPUT_FILE" ]; then
        local output_dir
        output_dir="$(dirname "$OUTPUT_FILE")"
        mkdir -p "$output_dir"

        local temp_file
        temp_file=$(mktemp "${output_dir}/.incus_metrics.XXXXXX")

        if ! generate_metrics > "$temp_file" 2>/dev/null; then
            rm -f "$temp_file"
            echo "ERROR: Failed to generate metrics" >&2
            exit 1
        fi

        local file_lines
        file_lines=$(wc -l < "$temp_file" 2>/dev/null || echo 0)

        if [ "$file_lines" -lt 3 ]; then
            rm -f "$temp_file"
            echo "ERROR: Metrics file too small ($file_lines lines), keeping previous" >&2
            exit 1
        fi

        chmod 644 "$temp_file"
        mv -f "$temp_file" "$OUTPUT_FILE"

        echo "Metrics written to $OUTPUT_FILE ($file_lines lines)" >&2
    else
        generate_metrics
    fi
}

main "$@"