#!/bin/bash
################################################################################
# Script Name: mailcow-exporter.sh
# Version: 1.0
# Description: Prometheus exporter for Mailcow — exposes container health,
#              mailbox statistics, domain info, rspamd action counts,
#              quarantine status, and version information as Prometheus metrics.
#              Supports stdout, node_exporter textfile collector, and HTTP
#              server modes.
#
# Author: Phil Connor
# Contact: contact@mylinux.work
# Website: https://mylinux.work
# License: MIT
#
# Prerequisites:
#   - curl and jq
#   - Mailcow API key (read-only recommended)
#   - netcat (nc) for HTTP mode
#
# Usage:
#   ./mailcow-exporter.sh                                    # stdout
#   ./mailcow-exporter.sh --textfile                         # node_exporter textfile
#   ./mailcow-exporter.sh --http -p 9201                     # HTTP server
#   ./mailcow-exporter.sh --url https://mail.example.com --api-key KEY
#
################################################################################

# ── Configuration ─────────────────────────────────────────────────────────────

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

# Source environment file if present (for cron deployments)
[ -f /etc/default/mailcow-exporter ] && . /etc/default/mailcow-exporter

MAILCOW_URL="${MAILCOW_URL:-https://localhost}"
MAILCOW_API_KEY="${MAILCOW_API_KEY:-}"

# ── 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]

Prometheus exporter for Mailcow (v1.0).

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

Options:
  -p, --port PORT     HTTP port (default: 9201)
  -o, --output FILE   Output file path
  --url URL           Mailcow base URL (default: $MAILCOW_URL)
  --api-key KEY       Mailcow API key

Environment: MAILCOW_URL, MAILCOW_API_KEY
EOF
    exit 0
}

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

# ── Metrics generation ────────────────────────────────────────────────────────

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

    # Prerequisite checks
    local prereq_fail=""
    command -v curl >/dev/null 2>&1 || prereq_fail="curl is required but not found"
    [ -z "$prereq_fail" ] && ! command -v jq >/dev/null 2>&1 && prereq_fail="jq is required but not found"
    [ -z "$prereq_fail" ] && [ -z "$MAILCOW_API_KEY" ] && prereq_fail="MAILCOW_API_KEY is required but not set"

    if [ -n "$prereq_fail" ]; then
        echo "# ERROR: $prereq_fail" >&2
        cat <<EOF
# HELP mailcow_exporter_up Exporter status (1=up, 0=down)
# TYPE mailcow_exporter_up gauge
mailcow_exporter_up 0
# HELP mailcow_exporter_info Exporter version information
# TYPE mailcow_exporter_info gauge
mailcow_exporter_info{version="1.0"} 1
EOF
        return
    fi

    local exporter_up=1
    local api_curl="curl -sf --max-time 10 -k -H X-API-Key:${MAILCOW_API_KEY}"

    cat <<EOF
# HELP mailcow_exporter_info Exporter version information
# TYPE mailcow_exporter_info gauge
mailcow_exporter_info{version="1.0"} 1

EOF

    # ── Version / Health ─────────────────────────────────────────────
    local version_json
    version_json=$($api_curl "${MAILCOW_URL}/api/v1/get/status/version" 2>/dev/null)
    if [ -n "$version_json" ] && [ "$version_json" != "null" ]; then
        local mc_version
        mc_version=$(echo "$version_json" | jq -r '.version // "unknown"' 2>/dev/null)
        mc_version=$(prom_escape "$mc_version")
        cat <<EOF
# HELP mailcow_version_info Mailcow version information
# TYPE mailcow_version_info gauge
mailcow_version_info{version="${mc_version}"} 1

EOF
    else
        echo "# WARNING: could not read /api/v1/get/status/version" >&2
        exporter_up=0
    fi

    # ── Container Health ─────────────────────────────────────────────
    local containers_json container_lines=""
    containers_json=$($api_curl "${MAILCOW_URL}/api/v1/get/status/containers" 2>/dev/null)
    if [ -n "$containers_json" ] && [ "$containers_json" != "null" ]; then
        local container_data
        container_data=$(echo "$containers_json" | jq -r '
            to_entries[] | "\(.key)\t\(.value.state // "unknown")"
        ' 2>/dev/null)
        if [ -n "$container_data" ]; then
            while IFS=$'\t' read -r c_name c_state; do
                [ -z "$c_name" ] && continue
                local esc_name up_val
                esc_name=$(prom_escape "$c_name")
                [ "$c_state" = "running" ] && up_val=1 || up_val=0
                container_lines="${container_lines}mailcow_container_up{name=\"${esc_name}\"} ${up_val}
"
            done <<< "$container_data"
        fi
    else
        echo "# WARNING: could not read /api/v1/get/status/containers" >&2
        exporter_up=0
    fi

    if [ -n "$container_lines" ]; then
        echo "# HELP mailcow_container_up Container running status (1=running, 0=down)"
        echo "# TYPE mailcow_container_up gauge"
        printf '%s' "$container_lines"
        echo ""
    fi

    # ── Domain Metrics ───────────────────────────────────────────────
    local domains_json domain_count=0 domain_info_lines=""
    domains_json=$($api_curl "${MAILCOW_URL}/api/v1/get/domain/all" 2>/dev/null)
    if [ -n "$domains_json" ] && [ "$domains_json" != "null" ]; then
        domain_count=$(echo "$domains_json" | jq 'length' 2>/dev/null)
        domain_count=${domain_count:-0}
        local domain_data
        domain_data=$(echo "$domains_json" | jq -r '
            .[] | "\(.domain_name)\t\(.active // 0)"
        ' 2>/dev/null)
        if [ -n "$domain_data" ]; then
            while IFS=$'\t' read -r d_name d_active; do
                [ -z "$d_name" ] && continue
                local esc_domain esc_active
                esc_domain=$(prom_escape "$d_name")
                esc_active=$(prom_escape "$d_active")
                domain_info_lines="${domain_info_lines}mailcow_domain_info{domain=\"${esc_domain}\",active=\"${esc_active}\"} 1
"
            done <<< "$domain_data"
        fi
    else
        echo "# WARNING: could not read /api/v1/get/domain/all" >&2
    fi

    cat <<EOF
# HELP mailcow_domain_count Total number of domains
# TYPE mailcow_domain_count gauge
mailcow_domain_count $domain_count

EOF

    if [ -n "$domain_info_lines" ]; then
        echo "# HELP mailcow_domain_info Domain inventory (always 1)"
        echo "# TYPE mailcow_domain_info gauge"
        printf '%s' "$domain_info_lines"
        echo ""
    fi

    # ── Mailbox Metrics ──────────────────────────────────────────────
    local mailboxes_json mailbox_count=0
    local quota_used_lines="" quota_total_lines="" messages_lines="" active_lines=""
    mailboxes_json=$($api_curl "${MAILCOW_URL}/api/v1/get/mailbox/all" 2>/dev/null)
    if [ -n "$mailboxes_json" ] && [ "$mailboxes_json" != "null" ]; then
        mailbox_count=$(echo "$mailboxes_json" | jq 'length' 2>/dev/null)
        mailbox_count=${mailbox_count:-0}
        local mailbox_data
        mailbox_data=$(echo "$mailboxes_json" | jq -r '
            .[] | "\(.username)\t\(.domain)\t\(.quota_used // 0)\t\(.quota // 0)\t\(.messages // 0)\t\(.active // 0)"
        ' 2>/dev/null)
        if [ -n "$mailbox_data" ]; then
            while IFS=$'\t' read -r mb_user mb_domain mb_quota_used mb_quota_total mb_messages mb_active; do
                [ -z "$mb_user" ] && continue
                local esc_mb esc_mbd
                esc_mb=$(prom_escape "$mb_user")
                esc_mbd=$(prom_escape "$mb_domain")
                local labels="mailbox=\"${esc_mb}\",domain=\"${esc_mbd}\""
                quota_used_lines="${quota_used_lines}mailcow_mailbox_quota_used_bytes{${labels}} ${mb_quota_used}
"
                quota_total_lines="${quota_total_lines}mailcow_mailbox_quota_total_bytes{${labels}} ${mb_quota_total}
"
                messages_lines="${messages_lines}mailcow_mailbox_messages{${labels}} ${mb_messages}
"
                active_lines="${active_lines}mailcow_mailbox_active{${labels}} ${mb_active}
"
            done <<< "$mailbox_data"
        fi
    else
        echo "# WARNING: could not read /api/v1/get/mailbox/all" >&2
    fi

    cat <<EOF
# HELP mailcow_mailbox_count Total number of mailboxes
# TYPE mailcow_mailbox_count gauge
mailcow_mailbox_count $mailbox_count

EOF

    if [ -n "$quota_used_lines" ]; then
        echo "# HELP mailcow_mailbox_quota_used_bytes Quota used per mailbox in bytes"
        echo "# TYPE mailcow_mailbox_quota_used_bytes gauge"
        printf '%s' "$quota_used_lines"
        echo ""
    fi
    if [ -n "$quota_total_lines" ]; then
        echo "# HELP mailcow_mailbox_quota_total_bytes Total quota per mailbox in bytes"
        echo "# TYPE mailcow_mailbox_quota_total_bytes gauge"
        printf '%s' "$quota_total_lines"
        echo ""
    fi
    if [ -n "$messages_lines" ]; then
        echo "# HELP mailcow_mailbox_messages Message count per mailbox"
        echo "# TYPE mailcow_mailbox_messages gauge"
        printf '%s' "$messages_lines"
        echo ""
    fi
    if [ -n "$active_lines" ]; then
        echo "# HELP mailcow_mailbox_active Mailbox active status (1=active, 0=inactive)"
        echo "# TYPE mailcow_mailbox_active gauge"
        printf '%s' "$active_lines"
        echo ""
    fi

    # ── Rspamd Metrics ───────────────────────────────────────────────
    local rspamd_json rspamd_scanned=0 action_lines=""
    rspamd_json=$($api_curl "${MAILCOW_URL}/api/v1/get/logs/rspamd-history/100" 2>/dev/null)
    if [ -n "$rspamd_json" ] && [ "$rspamd_json" != "null" ]; then
        rspamd_scanned=$(echo "$rspamd_json" | jq 'length' 2>/dev/null)
        rspamd_scanned=${rspamd_scanned:-0}
        local action_data
        action_data=$(echo "$rspamd_json" | jq -r '
            group_by(.action) | .[] | "\(.[0].action)\t\(length)"
        ' 2>/dev/null)
        if [ -n "$action_data" ]; then
            while IFS=$'\t' read -r act_name act_count; do
                [ -z "$act_name" ] && continue
                local esc_action
                esc_action=$(prom_escape "$act_name")
                action_lines="${action_lines}mailcow_rspamd_action_total{action=\"${esc_action}\"} ${act_count}
"
            done <<< "$action_data"
        fi
    else
        echo "# WARNING: could not read /api/v1/get/logs/rspamd-history/100" >&2
    fi

    cat <<EOF
# HELP mailcow_rspamd_scanned_total Total messages scanned (from recent history)
# TYPE mailcow_rspamd_scanned_total gauge
mailcow_rspamd_scanned_total $rspamd_scanned

EOF

    if [ -n "$action_lines" ]; then
        echo "# HELP mailcow_rspamd_action_total Count of rspamd actions by type"
        echo "# TYPE mailcow_rspamd_action_total gauge"
        printf '%s' "$action_lines"
        echo ""
    fi

    # ── Quarantine ────────────────────────────────────────────────────
    local quarantine_json quarantine_count=0
    quarantine_json=$($api_curl "${MAILCOW_URL}/api/v1/get/quarantine/all" 2>/dev/null)
    if [ -n "$quarantine_json" ] && [ "$quarantine_json" != "null" ]; then
        quarantine_count=$(echo "$quarantine_json" | jq 'length' 2>/dev/null)
        quarantine_count=${quarantine_count:-0}
    else
        echo "# WARNING: could not read /api/v1/get/quarantine/all" >&2
    fi

    cat <<EOF
# HELP mailcow_quarantine_count Messages currently in quarantine
# TYPE mailcow_quarantine_count gauge
mailcow_quarantine_count $quarantine_count

EOF

    # ── Rate Limit Status ────────────────────────────────────────────
    local ratelimit_json
    ratelimit_json=$($api_curl "${MAILCOW_URL}/api/v1/get/rl-mbox" 2>/dev/null)
    if [ -n "$ratelimit_json" ] && [ "$ratelimit_json" != "null" ]; then
        echo "# Rate limit data retrieved successfully"
    else
        echo "# WARNING: could not read /api/v1/get/rl-mbox" >&2
    fi

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

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

# HELP mailcow_exporter_duration_seconds Time to generate all metrics
# TYPE mailcow_exporter_duration_seconds gauge
mailcow_exporter_duration_seconds $script_duration

# HELP mailcow_exporter_last_run_timestamp Unix timestamp of last successful run
# TYPE mailcow_exporter_last_run_timestamp gauge
mailcow_exporter_last_run_timestamp $script_end
EOF
    echo ""
}

# ── HTTP server mode ──────────────────────────────────────────────────────────

run_http_server() {
    echo "# Starting mailcow 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 mailcow 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='<html><head><title>Mailcow Exporter v1.0</title></head><body><h1>Mailcow Exporter v1.0</h1><p><a href="/metrics">Metrics</a></p></body></html>'
                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}/.mailcow_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 "$@"