#!/usr/bin/with-contenv bashio
# shellcheck shell=bash

FAH_API="http://localhost:7396/api"
POLL_INTERVAL=30
STATE_TOPIC="foldingathome/state"
DISCOVERY_PREFIX="homeassistant"

# Abort cleanly if MQTT service is not available
if ! bashio::services.mqtt; then
    bashio::log.info "MQTT not available — skipping sensor publishing. Install the Mosquitto broker add-on to enable HA sensors."
    exec sleep infinity
fi

MQTT_HOST=$(bashio::services.mqtt 'host')
MQTT_PORT=$(bashio::services.mqtt 'port')
MQTT_USER=$(bashio::services.mqtt 'username')
MQTT_PASS=$(bashio::services.mqtt 'password')
UNIQUE_ID="foldingathome_$(hostname | tr -cd '[:alnum:]_')"

mqtt_publish() {
    mosquitto_pub \
        -h "$MQTT_HOST" -p "$MQTT_PORT" \
        -u "$MQTT_USER" -P "$MQTT_PASS" \
        -t "$1" -m "$2" -r -q 1
}

publish_discovery() {
    local device
    device=$(jq -n \
        --arg id "$UNIQUE_ID" \
        '{identifiers: [$id], name: "Folding@home", model: "FAH 8", manufacturer: "foldingathome.org"}')

    declare -A sensors
    sensors=(
        ["status"]='{"name":"Folding@home Status","value_template":"{{ value_json.status }}","icon":"mdi:dna"}'
        ["ppd"]='{"name":"Folding@home PPD","value_template":"{{ value_json.ppd }}","unit_of_measurement":"PPD","icon":"mdi:speedometer","state_class":"measurement"}'
        ["progress"]='{"name":"Folding@home Progress","value_template":"{{ value_json.progress }}","unit_of_measurement":"%","icon":"mdi:progress-clock","state_class":"measurement"}'
        ["credit"]='{"name":"Folding@home Credit","value_template":"{{ value_json.credit }}","unit_of_measurement":"points","icon":"mdi:star","state_class":"total_increasing"}'
    )

    for key in "${!sensors[@]}"; do
        local config
        config=$(echo "${sensors[$key]}" | jq \
            --arg uid "${UNIQUE_ID}_${key}" \
            --arg topic "$STATE_TOPIC" \
            --argjson device "$device" \
            '. + {unique_id: $uid, state_topic: $topic, device: $device}')
        mqtt_publish "${DISCOVERY_PREFIX}/sensor/${UNIQUE_ID}_${key}/config" "$config"
    done

    bashio::log.info "HA sensor discovery published (${#sensors[@]} sensors)"
}

# Wait for FAH API to become ready
bashio::log.info "Waiting for Folding@home API on port 7396..."
until curl -sf "${FAH_API}/info" > /dev/null 2>&1; do
    sleep 5
done
bashio::log.info "Folding@home API ready — starting MQTT sensor publisher"

publish_discovery

while true; do
    UNITS=$(curl -sf "${FAH_API}/units" 2>/dev/null || echo "[]")
    UNIT_COUNT=$(echo "$UNITS" | jq 'length')

    if [ "$UNIT_COUNT" -gt 0 ]; then
        RAW_STATE=$(echo "$UNITS" | jq -r '.[0].state // "unknown"')
        case "$RAW_STATE" in
            RUN)      STATUS="Running" ;;
            PAUSE)    STATUS="Paused" ;;
            FINISH)   STATUS="Finishing" ;;
            ASSIGN)   STATUS="Assigning" ;;
            DOWNLOAD) STATUS="Downloading" ;;
            SEND)     STATUS="Sending" ;;
            *)        STATUS="Idle" ;;
        esac
        PPD=$(echo "$UNITS" | jq '[.[].ppd // 0] | add // 0 | round')
        PROGRESS=$(echo "$UNITS" | jq '.[0].progress // 0 | . * 100 | round')
        CREDIT=$(echo "$UNITS" | jq '[.[].credit // 0] | add // 0 | round')
    else
        STATUS="Idle"
        PPD=0
        PROGRESS=0
        CREDIT=0
    fi

    PAYLOAD=$(jq -n \
        --arg status "$STATUS" \
        --argjson ppd "$PPD" \
        --argjson progress "$PROGRESS" \
        --argjson credit "$CREDIT" \
        '{status: $status, ppd: $ppd, progress: $progress, credit: $credit}')

    mqtt_publish "$STATE_TOPIC" "$PAYLOAD"

    sleep "$POLL_INTERVAL"
done