#!/bin/bash
################################################################################
# Script Name: plex-exporter.sh
# Version: 1.0
# Description: Prometheus exporter for Plex Media Server providing operational
#              metrics via the Plex API — library counts, active sessions,
#              transcoding stats, bandwidth usage, server info, and storage
#
# Author: Phil Connor
# Contact: contact@mylinux.work
# Website: https://mylinux.work
# License: MIT
#
# Prerequisites:
#   - Plex Media Server running with API access
#   - Plex authentication token (X-Plex-Token)
#   - curl for API calls
#   - netcat (nc) for HTTP mode
#
# Usage:
#   # Output to stdout
#   ./plex-exporter.sh --token YOUR_PLEX_TOKEN
#
#   # HTTP server mode
#   ./plex-exporter.sh --http --token YOUR_PLEX_TOKEN
#
#   # Textfile collector mode
#   ./plex-exporter.sh --textfile --token YOUR_PLEX_TOKEN
#
#   # Custom Plex URL
#   ./plex-exporter.sh --url http://10.0.0.50:32400 --token YOUR_PLEX_TOKEN
#
# Metrics Exported:
#   - plex_up - API reachability (1=up, 0=down)
#   - plex_info{version,platform,hostname} - Server info
#   - plex_library_items_total{library,type} - Items per library
#   - plex_sessions_active - Active streams
#   - plex_sessions_transcode - Active transcode sessions
#   - plex_sessions_direct_play - Active direct play sessions
#   - plex_sessions_direct_stream - Active direct stream sessions
#   - plex_session_info{user,player,title,state,decision} - Per-session details
#   - plex_bandwidth_bytes - Current total bandwidth
#   - plex_transcode_speed - Current transcode speed ratio
#   - plex_library_size_bytes{library} - Library storage usage
#   - plex_server_update_available - Update available (1/0)
#   - plex_exporter_duration_seconds - Script execution time
#   - plex_exporter_last_run_timestamp - Last run timestamp
#
# Configuration:
#   Default HTTP port: 9595
#   Default Plex URL: http://localhost:32400
#   Textfile directory: /var/lib/node_exporter
#
################################################################################

set -uo pipefail

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

TEXTFILE_DIR="/var/lib/node_exporter"
OUTPUT_FILE=""
HTTP_MODE=false
HTTP_PORT=9595
PLEX_URL="http://localhost:32400"
PLEX_TOKEN=""
TAUTULLI_URL=""
TAUTULLI_API_KEY=""

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

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

Export Plex Media Server statistics as Prometheus metrics via the Plex API.

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

OPTIONS:
    -p, --port        HTTP port (default: 9595)
    -o, --output      Output file path
    --url             Plex server URL (default: http://localhost:32400)
    --token           Plex authentication token (required)
    --tautulli-url    Tautulli URL for enhanced metrics (optional)
    --tautulli-key    Tautulli API key (optional)

EXAMPLES:
    $0 --textfile --token mytoken                   # Write to textfile collector
    $0 --http --port 9595 --token mytoken           # Run HTTP server
    $0 --token mytoken                              # Output to stdout
    $0 --url http://10.0.0.50:32400 --token mytoken # Custom Plex URL

EOF
    exit 0
}

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

# Check prerequisites
# Returns: 0 if OK, 1 if error
check_prerequisites() {
    if ! command -v curl >/dev/null 2>&1; then
        echo "ERROR: curl not found" >&2
        return 1
    fi

    if [ -z "$PLEX_TOKEN" ]; then
        echo "ERROR: --token is required (Plex authentication token)" >&2
        return 1
    fi

    return 0
}

# Escape special characters in Prometheus label values
# Args: $1 - string to escape
# Returns: escaped string safe for Prometheus labels
prom_escape() {
    local val="$1"
    val="${val//\\/\\\\}"
    val="${val//\"/\\\"}"
    val="${val//$'\n'/}"
    echo "$val"
}

# Make an authenticated Plex API call (XML response)
# Args: $1 - API endpoint path (e.g., /status/sessions)
# Returns: XML response on stdout
plex_api() {
    local endpoint="$1"
    curl -s -X GET \
        -H "X-Plex-Token: ${PLEX_TOKEN}" \
        -H "Accept: application/xml" \
        "${PLEX_URL}${endpoint}" 2>/dev/null
}

# Make an authenticated Plex API call (JSON response)
# Args: $1 - API endpoint path
# Returns: JSON response on stdout
plex_api_json() {
    local endpoint="$1"
    curl -s -X GET \
        -H "X-Plex-Token: ${PLEX_TOKEN}" \
        -H "Accept: application/json" \
        "${PLEX_URL}${endpoint}" 2>/dev/null
}

# Extract XML attribute value using grep/sed (no xmllint dependency)
# Args: $1 - attribute name, $2 - XML string
# Returns: attribute value
xml_attr() {
    local attr="$1"
    local xml="$2"
    echo "$xml" | grep -oP "${attr}=\"[^\"]*\"" | head -1 | sed "s/${attr}=\"//;s/\"//"
}

# Extract all values of an XML attribute across multiple elements
# Args: $1 - attribute name, $2 - XML string
# Returns: one value per line
xml_attr_all() {
    local attr="$1"
    local xml="$2"
    echo "$xml" | grep -oP "${attr}=\"[^\"]*\"" | sed "s/${attr}=\"//;s/\"//"
}

# ============================================================================
# METRIC GENERATION
# ============================================================================

# Generate all Prometheus metrics
# Returns: Prometheus text format metrics on stdout
generate_metrics() {
    local script_start
    script_start=$(date +%s%N)

    # Check prerequisites
    if ! check_prerequisites; then
        cat <<EOF
# HELP plex_up Plex API reachability (1=up, 0=down)
# TYPE plex_up gauge
plex_up 0
EOF
        return
    fi

    # ========================================================================
    # SERVER INFO
    # ========================================================================

    local identity_response
    identity_response=$(plex_api "/identity" 2>/dev/null)

    if [ -z "$identity_response" ]; then
        cat <<EOF
# HELP plex_up Plex API reachability (1=up, 0=down)
# TYPE plex_up gauge
plex_up 0
EOF
        return
    fi

    local server_response
    server_response=$(plex_api "/" 2>/dev/null)

    local plex_version
    # Skip XML declaration (<?xml version="1.0"?>) — match version on MediaContainer element
    plex_version=$(echo "$identity_response" | grep -oP 'MediaContainer[^>]*version="[^"]*"' | grep -oP 'version="\K[^"]+' | head -1)

    if [ -z "$plex_version" ]; then
        cat <<EOF
# HELP plex_up Plex API reachability (1=up, 0=down)
# TYPE plex_up gauge
plex_up 0
EOF
        return
    fi

    local plex_platform plex_hostname plex_myPlex update_available
    plex_platform=$(xml_attr "platform" "$server_response")
    plex_hostname=$(xml_attr "friendlyName" "$server_response")
    plex_myPlex=$(xml_attr "myPlex" "$server_response")
    update_available=$(xml_attr "updatedAt" "$server_response")

    plex_platform="${plex_platform:-unknown}"
    plex_hostname="${plex_hostname:-unknown}"

    cat <<EOF
# HELP plex_up Plex API reachability (1=up, 0=down)
# TYPE plex_up gauge
plex_up 1

# HELP plex_info Plex Media Server information
# TYPE plex_info gauge
plex_info{version="$(prom_escape "$plex_version")",platform="$(prom_escape "$plex_platform")",hostname="$(prom_escape "$plex_hostname")"} 1
EOF

    echo ""

    # Check for updates
    local update_check
    update_check=$(plex_api_json "/updater/status" 2>/dev/null)

    local has_update=0
    if [ -n "$update_check" ]; then
        # Plex updater/status returns canInstall:true when a newer version is available
        if echo "$update_check" | grep -qP '"canInstall"\s*:\s*true' 2>/dev/null; then
            has_update=1
        fi
    fi

    cat <<EOF
# HELP plex_server_update_available Whether a Plex update is available (1=yes, 0=no)
# TYPE plex_server_update_available gauge
plex_server_update_available $has_update
EOF

    echo ""

    # ========================================================================
    # LIBRARY METRICS
    # ========================================================================

    local libraries_response
    libraries_response=$(plex_api "/library/sections" 2>/dev/null)

    if [ -n "$libraries_response" ]; then
        cat <<EOF
# HELP plex_library_items_total Total items in each Plex library
# TYPE plex_library_items_total gauge
EOF

        # Extract library keys and types
        local lib_keys lib_titles lib_types
        lib_keys=$(echo "$libraries_response" | grep -oP 'Directory[^>]*key="[^"]*"' | grep -oP 'key="[^"]*"' | sed 's/key="//;s/"//' || true)
        lib_titles=$(echo "$libraries_response" | grep -oP 'Directory[^>]*title="[^"]*"' | grep -oP 'title="[^"]*"' | sed 's/title="//;s/"//' || true)
        lib_types=$(echo "$libraries_response" | grep -oP 'Directory[^>]*type="[^"]*"' | grep -oP 'type="[^"]*"' | sed 's/type="//;s/"//' || true)

        # Convert to arrays
        local -a keys=() titles=() types=()
        while IFS= read -r line; do keys+=("$line"); done <<< "$lib_keys"
        while IFS= read -r line; do titles+=("$line"); done <<< "$lib_titles"
        while IFS= read -r line; do types+=("$line"); done <<< "$lib_types"

        for i in "${!keys[@]}"; do
            local key="${keys[$i]}"
            local title="${titles[$i]}"
            local type="${types[$i]}"

            if [ -z "$key" ] || [ -z "$title" ]; then
                continue
            fi

            # Get item count for this library
            local count_response
            count_response=$(plex_api "/library/sections/${key}/all?X-Plex-Container-Start=0&X-Plex-Container-Size=0" 2>/dev/null)

            local item_count=0
            if [ -n "$count_response" ]; then
                item_count=$(xml_attr "totalSize" "$count_response")
                item_count="${item_count:-0}"
            fi

            echo "plex_library_items_total{library=\"$(prom_escape "$title")\",type=\"$(prom_escape "$type")\"} $item_count"
        done

        echo ""

        # Library storage usage (if available)
        cat <<EOF
# HELP plex_library_size_bytes Storage used per library section (bytes)
# TYPE plex_library_size_bytes gauge
EOF

        for i in "${!keys[@]}"; do
            local key="${keys[$i]}"
            local title="${titles[$i]}"

            if [ -z "$key" ] || [ -z "$title" ]; then
                continue
            fi

            # Get total storage for library from media info
            local size_response
            size_response=$(plex_api_json "/library/sections/${key}/totalSize" 2>/dev/null)

            if [ -n "$size_response" ]; then
                local total_bytes
                total_bytes=$(echo "$size_response" | grep -oP '"totalStorage"\s*:\s*[0-9]+' | head -1 | grep -oP '[0-9]+$')
                if [ -n "$total_bytes" ]; then
                    echo "plex_library_size_bytes{library=\"$(prom_escape "$title")\"} $total_bytes"
                fi
            fi
        done
    fi

    echo ""

    # ========================================================================
    # SESSION METRICS
    # ========================================================================

    local sessions_response
    sessions_response=$(plex_api "/status/sessions" 2>/dev/null)

    local total_sessions=0
    local transcode_sessions=0
    local direct_play_sessions=0
    local direct_stream_sessions=0
    local total_bandwidth=0

    if [ -n "$sessions_response" ]; then
        total_sessions=$(xml_attr "size" "$sessions_response")
        total_sessions="${total_sessions:-0}"

        if [ "$total_sessions" -gt 0 ]; then
            # Parse individual sessions
            local session_blocks
            session_blocks=$(echo "$sessions_response" | grep -oP '<Video[^/]*/>|<Video[^>]*>.*?</Video>|<Track[^/]*/>|<Track[^>]*>.*?</Track>' 2>/dev/null || true)

            # Count by transcode decision
            # Note: Plex may omit TranscodeSession entirely for direct play
            transcode_sessions=$(echo "$sessions_response" | grep -cP 'videoDecision="transcode"' || true)
            direct_stream_sessions=$(echo "$sessions_response" | grep -cP 'videoDecision="copy"' || true)
            local has_decision
            has_decision=$(echo "$sessions_response" | grep -cP 'videoDecision=' || true)
            direct_play_sessions=$((total_sessions - transcode_sessions - direct_stream_sessions))

            # Calculate total bandwidth from Session elements
            local bandwidths
            bandwidths=$(echo "$sessions_response" | grep -oP 'bandwidth="[0-9]+"' | grep -oP '[0-9]+' || true)
            if [ -n "$bandwidths" ]; then
                while IFS= read -r bw; do
                    total_bandwidth=$((total_bandwidth + bw))
                done <<< "$bandwidths"
                # Plex reports bandwidth in kbps, convert to bytes/sec
                total_bandwidth=$((total_bandwidth * 1000 / 8))
            fi

            # Per-session info metrics
            cat <<EOF
# HELP plex_session_info Per-session information
# TYPE plex_session_info gauge
EOF
            # Extract per-session info by splitting on Video/Track blocks
            # Plex XML nests User, Player, TranscodeSession as child elements
            local session_xml
            session_xml=$(echo "$sessions_response" | tr '\n' ' ')

            # Extract each Video/Track block with children
            echo "$session_xml" | grep -oP '<Video\b[^>]*>.*?</Video>|<Track\b[^>]*>.*?</Track>' 2>/dev/null | while IFS= read -r block; do
                local s_title s_user s_player s_state s_decision

                s_title=$(echo "$block" | grep -oP '(?<=<(Video|Track)\s)[^>]*' | grep -oP 'title="[^"]*"' | head -1 | sed 's/title="//;s/"//' || true)
                s_user=$(echo "$block" | grep -oP '<User[^>]*title="[^"]*"' | head -1 | sed 's/.*title="//;s/"//' || true)
                s_player=$(echo "$block" | grep -oP '<Player[^>]*product="[^"]*"' | head -1 | sed 's/.*product="//;s/"//' || true)
                s_state=$(echo "$block" | grep -oP '<Player[^>]*state="[^"]*"' | head -1 | sed 's/.*state="//;s/"//' || true)
                s_decision=$(echo "$block" | grep -oP 'videoDecision="[^"]*"' | head -1 | sed 's/videoDecision="//;s/"//' || true)

                s_title="${s_title:-unknown}"
                s_user="${s_user:-unknown}"
                s_player="${s_player:-unknown}"
                s_state="${s_state:-unknown}"
                s_decision="${s_decision:-directplay}"

                echo "plex_session_info{user=\"$(prom_escape "$s_user")\",player=\"$(prom_escape "$s_player")\",title=\"$(prom_escape "$s_title")\",state=\"$(prom_escape "$s_state")\",decision=\"$(prom_escape "$s_decision")\"} 1"
            done
        fi
    fi

    cat <<EOF

# HELP plex_sessions_active Total active streaming sessions
# TYPE plex_sessions_active gauge
plex_sessions_active $total_sessions

# HELP plex_sessions_transcode Active transcode sessions
# TYPE plex_sessions_transcode gauge
plex_sessions_transcode $transcode_sessions

# HELP plex_sessions_direct_play Active direct play sessions
# TYPE plex_sessions_direct_play gauge
plex_sessions_direct_play $direct_play_sessions

# HELP plex_sessions_direct_stream Active direct stream sessions (copy)
# TYPE plex_sessions_direct_stream gauge
plex_sessions_direct_stream $direct_stream_sessions

# HELP plex_bandwidth_bytes Current total streaming bandwidth in bytes per second
# TYPE plex_bandwidth_bytes gauge
plex_bandwidth_bytes $total_bandwidth
EOF

    echo ""

    # ========================================================================
    # TRANSCODE METRICS
    # ========================================================================

    local transcode_response
    transcode_response=$(plex_api "/transcode/sessions" 2>/dev/null)

    local transcode_count=0
    local avg_speed=0

    if [ -n "$transcode_response" ]; then
        transcode_count=$(xml_attr "size" "$transcode_response")
        transcode_count="${transcode_count:-0}"

        if [ "$transcode_count" -gt 0 ]; then
            local speeds total_speed=0 speed_count=0
            speeds=$(echo "$transcode_response" | grep -oP 'speed="[0-9.]*"' | grep -oP '[0-9.]+' || true)
            if [ -n "$speeds" ]; then
                while IFS= read -r spd; do
                    total_speed=$(awk "BEGIN {printf \"%.2f\", $total_speed + $spd}")
                    speed_count=$((speed_count + 1))
                done <<< "$speeds"
                if [ "$speed_count" -gt 0 ]; then
                    avg_speed=$(awk "BEGIN {printf \"%.2f\", $total_speed / $speed_count}")
                fi
            fi
        fi
    fi

    cat <<EOF
# HELP plex_transcode_sessions_active Active transcode sessions from transcode API
# TYPE plex_transcode_sessions_active gauge
plex_transcode_sessions_active $transcode_count

# HELP plex_transcode_speed Average transcode speed ratio (>1.0 = faster than realtime)
# TYPE plex_transcode_speed gauge
plex_transcode_speed $avg_speed
EOF

    echo ""

    # ========================================================================
    # TAUTULLI ENHANCED METRICS (optional)
    # ========================================================================

    if [ -n "$TAUTULLI_URL" ] && [ -n "$TAUTULLI_API_KEY" ]; then
        local tautulli_response
        tautulli_response=$(curl -s "${TAUTULLI_URL}/api/v2?apikey=${TAUTULLI_API_KEY}&cmd=get_libraries" 2>/dev/null)

        if [ -n "$tautulli_response" ] && echo "$tautulli_response" | grep -q '"result":"success"' 2>/dev/null; then
            cat <<EOF
# HELP plex_tautulli_library_duration_seconds Total play duration tracked by Tautulli per library (seconds)
# TYPE plex_tautulli_library_duration_seconds counter
EOF
            echo "$tautulli_response" | grep -oP '\{[^}]*"section_name"[^}]*\}' | while IFS= read -r lib; do
                local lib_name lib_duration
                lib_name=$(echo "$lib" | grep -oP '"section_name"\s*:\s*"[^"]*"' | sed 's/.*"section_name"\s*:\s*"//;s/"//')
                lib_duration=$(echo "$lib" | grep -oP '"duration"\s*:\s*[0-9]+' | grep -oP '[0-9]+$')
                lib_name="${lib_name:-unknown}"
                lib_duration="${lib_duration:-0}"
                echo "plex_tautulli_library_duration_seconds{library=\"$(prom_escape "$lib_name")\"} $lib_duration"
            done

            echo ""

            # Get history stats
            local history_response
            history_response=$(curl -s "${TAUTULLI_URL}/api/v2?apikey=${TAUTULLI_API_KEY}&cmd=get_home_stats&stat_id=top_users&stats_count=10" 2>/dev/null)

            if [ -n "$history_response" ] && echo "$history_response" | grep -q '"result":"success"'; then
                cat <<EOF
# HELP plex_tautulli_user_plays_total Total plays per user tracked by Tautulli
# TYPE plex_tautulli_user_plays_total counter
EOF
                echo "$history_response" | grep -oP '\{[^}]*"friendly_name"[^}]*\}' | while IFS= read -r user; do
                    local user_name user_plays
                    user_name=$(echo "$user" | grep -oP '"friendly_name"\s*:\s*"[^"]*"' | sed 's/.*"friendly_name"\s*:\s*"//;s/"//')
                    user_plays=$(echo "$user" | grep -oP '"total_plays"\s*:\s*[0-9]+' | grep -oP '[0-9]+$')
                    user_name="${user_name:-unknown}"
                    user_plays="${user_plays:-0}"
                    echo "plex_tautulli_user_plays_total{user=\"$(prom_escape "$user_name")\"} $user_plays"
                done
            fi
        fi
    fi

    echo ""

    # ========================================================================
    # EXPORTER META METRICS
    # ========================================================================

    local script_end script_duration
    script_end=$(date +%s)
    local script_end_ns
    script_end_ns=$(date +%s%N)
    script_duration=$(awk "BEGIN {printf \"%.3f\", ($script_end_ns - $script_start) / 1000000000}")

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

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

    echo ""
}

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

# Run simple HTTP server using netcat
# Serves metrics on /metrics endpoint
run_http_server() {
    echo "Starting Plex 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

    # Infinite loop accepting HTTP requests
    while true; do
        {
            read -r request
            # Check if request is for /metrics endpoint
            if [[ "$request" =~ ^GET\ /metrics ]]; then
                echo -e "HTTP/1.1 200 OK\r\nContent-Type: text/plain; version=0.0.4\r\n\r"
                generate_metrics
            else  # Serve HTML landing page for other requests
                echo -e "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r"
                cat <<EOF
<!DOCTYPE html>
<html>
<head><title>Plex Exporter v1.0</title></head>
<body>
<h1>Plex Prometheus Exporter v1.0</h1>
<p><a href="/metrics">Metrics</a></p>
<p>Operational metrics from Plex Media Server API.</p>
</body>
</html>
EOF
            fi
        } | nc -l -p "$HTTP_PORT" -q 1 2>/dev/null
    done
}

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

# Main entry point - routes to appropriate output mode
main() {
    parse_args "$@"

    if [ "$HTTP_MODE" = true ]; then
        # Run HTTP server (blocks until killed)
        run_http_server
    elif [ -n "$OUTPUT_FILE" ]; then
        # Textfile collector mode: write atomically using temp file
        local output_dir
        output_dir="$(dirname "$OUTPUT_FILE")"
        mkdir -p "$output_dir"

        # Create temp file in SAME directory for atomic rename (same filesystem)
        local temp_file
        temp_file=$(mktemp "${output_dir}/.plex_metrics.XXXXXX")

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

        # Validate: file must exist, have content
        local file_lines
        file_lines=$(wc -l < "$temp_file" 2>/dev/null || echo 0)

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

        # Set permissions before move
        chmod 644 "$temp_file"

        # Atomic rename - no gap where file is missing
        mv -f "$temp_file" "$OUTPUT_FILE"

        echo "Metrics written to $OUTPUT_FILE ($file_lines lines)" >&2
    else
        # Default: output to stdout
        generate_metrics
    fi
}

# Execute main function with all script arguments
main "$@"