#!/usr/bin/env bash
set -euo pipefail

esn=""
fail() {
    jq -n --arg error "$1" --arg esn "$esn" \
        '{error: $error, esn: (if $esn == "" then null else $esn end), deviceInstances: [], deviceModels: []}'
    exit 1
}

if ! command -v jq >/dev/null; then
    echo '{"error":"jq not found in PATH","esn":null,"deviceInstances":[],"deviceModels":[]}'
    exit 1
fi
command -v kubectl >/dev/null || fail "kubectl not found in PATH"
kubectl cluster-info >/dev/null 2>&1 || fail "Cannot connect to Kubernetes cluster"
esn=$(kubectl get namespace kube-system -o jsonpath='{.metadata.uid}') || fail "Failed to get kube-system namespace UID"
[[ -n "$esn" ]] || fail "Failed to get kube-system namespace UID"
nodes=$(kubectl get nodes -o json) || fail "Failed to get nodes"
[[ -n "$nodes" ]] || fail "No node list returned"

# Each annotation is a string containing the vendor device list.
annotation() {
    jq -c --arg key "$1" '.items[].metadata.annotations[$key] // empty' <<< "$nodes"
}

json_devices() {
    jq -c 'fromjson? | arrays | .[] | objects'
}

collect_instances() {
    # NVIDIA, AMD, Alibaba and Ascend register JSON arrays with id/type fields.
    jq -c '.items[] | (.metadata.annotations // {}) | to_entries[]
        | select(.key == "hami.io/node-nvidia-register" or .key == "hami.io/node-amd-register"
            or .key == "hami.io/node-alibaba-register"
            or ((.key | startswith("hami.io/node-register-")) and .key != "hami.io/node-register-xpu"))
        | .value' <<< "$nodes" | json_devices | jq -c '{uuid: .id, deviceType: .type}'

    # HCU and Iluvatar register colon-separated devices. Fields 1 and 5 hold UUID and model.
    jq -c '.items[] | (.metadata.annotations // {}) | to_entries[]
        | select(.key | startswith("hami.io/node-") and endswith("-register"))
        | .value | select(contains(":") and (startswith("[") | not))
        | split(":")[] | split(",")
        | select((length == 7 or length == 9) and (.[1] | test("^[0-9]+$")))
        | {uuid: .[0], deviceType: .[4]}' <<< "$nodes"

    annotation metax-tech.com/node-gpu-devices | json_devices | jq -c '
        (.model // "" | gsub("^\\s+|\\s+$"; "")) as $model
        | {uuid, deviceType: (if $model == "" then "Metax-SGPU" else $model end)}'
}

registered_models() {
    annotation "$1" | json_devices | jq -c --arg fallback "$2" '
        (.type // "") as $model
        | if $model == "" or ($fallback == "XPU" and ($model | test("^\\s*$")))
          then $fallback else $model end | {deviceType: ., count: 1}'
}

model_count() {
    local model=$1 count=$2
    if (( count > 0 )); then
        jq -cn --arg model "$model" --argjson count "$count" '{deviceType: $model, count: $count}'
    fi
}

node_capacities() {
    jq -r --arg mlu "${HAMI_MLU_CORE_RESOURCE:-cambricon.com/mlu.smlu.vcore}" \
          --arg kunlun "${HAMI_KUNLUN_COUNT_RESOURCE:-kunlunxin.com/xpu}" '
        # Node capacity can use scaled integers such as 1k or 1Ki.
        def quantity:
            tonumber? // (capture("^(?<n>[0-9]+)(?<unit>[kKMGTPE])(?<binary>i?)$")
                | (.unit | ascii_upcase) as $unit
                | (.n | tonumber) * pow(if .binary == "i" then 1024 else 1000 end;
                                         ("KMGTPE" | index($unit)) + 1));
        .items[] | (.status.capacity // {}) as $cap
        | def count($key): (($cap[$key] // "0" | quantity) // 0);
        (.metadata.labels["node.kubernetes.io/instance-type"] // "" | gsub("^\\s+|\\s+$"; "")) as $label
        | [count("enflame.com/gcu"), count("metax-tech.com/gpu"), count($kunlun),
           count("mthreads.com/sgpu-core"), count($mlu), count("cambricon.com/mlu"),
           count("aws.amazon.com/neuron"), $label] | @tsv' <<< "$nodes"
}

collect_models() {
    registered_models hami.io/node-register-xpu XPU
    registered_models hami.io/node-biren-register Biren
    registered_models hami.io/node-va-register Vastai

    # DRS lists physical cards in devices; each card can advertise several slices.
    annotation enflame.com/gcu-drs-capacity | jq -c '
        fromjson? | select(.devices | type == "array")
        | select(any((.profiles // {} | keys[]); test("^\\s*[+]?[1-9][0-9]*g?\\s*\\."; "i")))
        | select(all(.devices[]; ((.capacity | tonumber?) // 0) >= 1))
        | {deviceType: "Enflame", count: (.devices | length)}'

    local gcu metax kunlun mt_cores mlu_cores mlu_cards neuron instance_type
    node_capacities | while IFS=$'\t' read -r gcu metax kunlun mt_cores mlu_cores mlu_cards neuron instance_type; do
        model_count GCU "$gcu"
        model_count Metax-GPU "$metax"
        model_count kunlun "$kunlun"
        model_count Mthreads "$(( (mt_cores + 15) / 16 ))"
        if (( mlu_cores > 0 )); then
            mlu_cards=$(( (mlu_cores + 99) / 100 ))
        fi
        model_count MLU "$mlu_cards"
        model_count "${instance_type:-AWSNeuron}" "$neuron"
    done
}

# Deduplicate only the stable UUID list. Every model-only card counts on its own node.
{
    collect_instances | jq -s 'map(select((.uuid | type == "string" and length > 0) and
        (.deviceType | type == "string" and length > 0))) | unique_by(.uuid)'
    collect_models | jq -s 'map(select(.count > 0)) | group_by(.deviceType)
        | map({deviceType: .[0].deviceType, count: (map(.count) | add)})'
} | jq -es --arg esn "$esn" 'select(length == 2) | {esn: $esn, deviceInstances: .[0], deviceModels: .[1]}'