#!/bin/sh
# Clear stale SpacemiT TCM ownership and run the K3 startup preflight.
#
# `spacemit-tcm-smi -c` is intentionally guarded: it is only used when every
# busy block belongs to a PID that no longer exists. Live owners are left intact
# and reported as a failed preflight.
set -eu

SMI="${K3_TCM_SMI:-spacemit-tcm-smi}"
MIN_CMA_FREE_KB="${K3_MIN_CMA_FREE_KB:-65536}"
MEMINFO="${K3_MEMINFO_PATH:-/proc/meminfo}"
TCM_DEVICE="${K3_TCM_DEVICE:-/dev/tcm}"
REQUIRE_TCM_DEVICE="${K3_REQUIRE_TCM_DEVICE:-1}"
REQUIRE_TCM_FREE="${K3_REQUIRE_TCM_FREE:-1}"

snapshot="$(mktemp "${TMPDIR:-/tmp}/k3-tcm-smi.XXXXXX")"
owners="$(mktemp "${TMPDIR:-/tmp}/k3-tcm-owners.XXXXXX")"
cleanup() {
    rm -f "$snapshot" "$owners"
}
trap cleanup EXIT INT TERM

json_bool() {
    case "${1:-0}" in
        0|false|False|FALSE|'')
            printf 'false'
            ;;
        *)
            printf 'true'
            ;;
    esac
}

run_smi_snapshot() {
    if "$SMI" >"$snapshot"; then
        return 0
    fi
    busy=0
    unknown=0
    emit_json 0 "tcm_smi_failed" '["spacemit_tcm_smi"]' 1
}

if_true() {
    if [ "${1:-0}" = "1" ]; then
        printf 'true'
    else
        printf 'false'
    fi
}

mem_kb() {
    key="$1"
    awk -v k="$key" '$1 == k ":" {print $2; found=1; exit} END {if (!found) print 0}' "$MEMINFO" 2>/dev/null
}

emit_json() {
    ok="$1"
    reason="$2"
    checks="$3"
    rc="$4"
    cma_free_kb="$(mem_kb CmaFree)"
    cma_total_kb="$(mem_kb CmaTotal)"
    mem_available_kb="$(mem_kb MemAvailable)"
    tcm_device_ok=0
    [ -e "$TCM_DEVICE" ] && tcm_device_ok=1
    smi_path="$(command -v "$SMI" 2>/dev/null || printf '%s' "$SMI")"
    printf '{"cma_free_kb":%s,"cma_total_kb":%s,"dev_tcm":%s,"failed_checks":%s,"mem_available_kb":%s,"min_cma_free_kb":%s,"ok":%s,"reason":"%s","require_tcm_device":%s,"spacemit_tcm_smi":"%s","tcm_busy_count_heuristic":%s,"tcm_busy_unknown":%s,"tcm_device":"%s"}\n' \
        "$cma_free_kb" \
        "$cma_total_kb" \
        "$(json_bool "$tcm_device_ok")" \
        "$checks" \
        "$mem_available_kb" \
        "$MIN_CMA_FREE_KB" \
        "$(json_bool "$ok")" \
        "$reason" \
        "$(if_true "$REQUIRE_TCM_DEVICE")" \
        "$smi_path" \
        "$busy" \
        "$(json_bool "$unknown")" \
        "$TCM_DEVICE"
    exit "$rc"
}

if ! command -v "$SMI" >/dev/null 2>&1; then
    busy=0
    unknown=0
    emit_json 0 "tcm_smi_unavailable" '["spacemit_tcm_smi"]' 1
fi

run_smi_snapshot
awk '$1 ~ /^[0-9]+$/ && $2 == "busy" {print $6}' "$snapshot" >"$owners"

busy=0
stale=0
live=0
unknown=0

while IFS= read -r pid; do
    [ -n "$pid" ] || continue
    busy=$((busy + 1))
    case "$pid" in
        *[!0-9]*|'')
            unknown=$((unknown + 1))
            ;;
        *)
            if kill -0 "$pid" 2>/dev/null; then
                live=$((live + 1))
            else
                stale=$((stale + 1))
            fi
            ;;
    esac
done <"$owners"

if [ "$stale" -gt 0 ] && [ "$live" -eq 0 ] && [ "$unknown" -eq 0 ]; then
    echo "k3-tcm-preflight-clean: releasing $stale stale TCM block(s)" >&2
    "$SMI" -c >/dev/null
    run_smi_snapshot
    awk '$1 ~ /^[0-9]+$/ && $2 == "busy" {print $6}' "$snapshot" >"$owners"
    busy=0
    unknown=0
    live=0
    while IFS= read -r pid; do
        [ -n "$pid" ] || continue
        busy=$((busy + 1))
        case "$pid" in
            *[!0-9]*|'')
                unknown=$((unknown + 1))
                ;;
            *)
                if kill -0 "$pid" 2>/dev/null; then
                    live=$((live + 1))
                fi
                ;;
        esac
    done <"$owners"
elif [ "$live" -gt 0 ] || [ "$unknown" -gt 0 ]; then
    echo "k3-tcm-preflight-clean: TCM busy with live/unknown owner(s); strict preflight will decide" >&2
fi

failed='[]'
if [ "$REQUIRE_TCM_DEVICE" = "1" ] && [ ! -e "$TCM_DEVICE" ]; then
    failed='["tcm_device"]'
    emit_json 0 "tcm_device_missing" "$failed" 1
fi

if [ "$REQUIRE_TCM_FREE" = "1" ] && { [ "$live" -gt 0 ] || [ "$unknown" -gt 0 ]; }; then
    failed='["tcm_busy"]'
    emit_json 0 "tcm_busy" "$failed" 1
fi

cma_free="$(mem_kb CmaFree)"
case "$cma_free" in
    *[!0-9]*|'')
        cma_free=0
        ;;
esac
if [ "$MIN_CMA_FREE_KB" -gt 0 ] && [ "$cma_free" -lt "$MIN_CMA_FREE_KB" ]; then
    failed='["cma_free_below_min"]'
    emit_json 0 "cma_free_below_min" "$failed" "${K3_TCM_CMA_SOFT_FAIL_EXIT:-86}"
fi

emit_json 1 "ok" "$failed" 0
