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

# CLUSTER=0        -> singolo nodo: avvia vllm serve senza Ray (nessun cluster)
# CLUSTER=1        -> nodo head: avvia Ray head e poi vllm serve
# CLUSTER>1         -> nodo worker: si collega al Ray head e resta attivo
CLUSTER="${CLUSTER:-1}"

RAY_PORT="${RAY_PORT:-6379}"
NUM_GPUS="${NUM_GPUS:-1}"

if [ "$CLUSTER" -eq 0 ]; then
    # Nessun cluster: non deve dipendere dalla rete Thunderbolt usata per Ray.
    unset VLLM_HOST_IP NCCL_SOCKET_IFNAME GLOO_SOCKET_IFNAME
    : "${MODEL_PATH:?MODEL_PATH non impostato}"

    # TOKENIZER e TRUST_REMOTE_CODE sono opzionali: se non impostate, vllm
    # usa il tokenizer embedded nel GGUF e non esegue codice remoto.
    EXTRA_ARGS=()
    if [ -n "${TOKENIZER:-}" ]; then
        EXTRA_ARGS+=(--tokenizer "${TOKENIZER}")
    fi
    if [ "${TRUST_REMOTE_CODE:-0}" = "1" ]; then
        EXTRA_ARGS+=(--trust-remote-code)
    fi
    if [ -n "${SERVED_MODEL_NAME:-}" ]; then
        EXTRA_ARGS+=(--served-model-name "${SERVED_MODEL_NAME}")
    fi
    if [ "${ENABLE_AUTO_TOOL_CHOICE:-0}" = "1" ]; then
        EXTRA_ARGS+=(--enable-auto-tool-choice)
        : "${TOOL_CALL_PARSER:?TOOL_CALL_PARSER richiesto se ENABLE_AUTO_TOOL_CHOICE=1}"
        EXTRA_ARGS+=(--tool-call-parser "${TOOL_CALL_PARSER}")
    fi
    if [ -n "${REASONING_PARSER:-}" ]; then
        EXTRA_ARGS+=(--reasoning-parser "${REASONING_PARSER}")
    fi
    if [ -n "${ATTENTION_BACKEND:-}" ]; then
        EXTRA_ARGS+=(--attention-backend "${ATTENTION_BACKEND}")
    fi
    if [ -n "${CHAT_TEMPLATE:-}" ]; then
        EXTRA_ARGS+=(--chat-template "${CHAT_TEMPLATE}")
    fi

    echo "[entrypoint] Modalita' singolo nodo (no cluster): avvio vllm serve senza Ray"
    exec vllm serve "${MODEL_PATH}" \
        "${EXTRA_ARGS[@]}" \
        --port "${SERVE_PORT:-7000}" \
        --host "${SERVE_HOST:-0.0.0.0}" \
        --max-model-len "${MAX_MODEL_LEN:-32768}" \
        --gpu-memory-utilization "${GPU_MEMORY_UTILIZATION:-0.9}" \
        --tensor-parallel-size 1 \
        --enforce-eager \
        --dtype "${DTYPE:-float16}" \
        --kv-cache-dtype "${KV_CACHE_DTYPE:-fp8}"

elif [ "$CLUSTER" -eq 1 ]; then
    : "${VLLM_HOST_IP:?VLLM_HOST_IP non impostato}"
    echo "[entrypoint] Nodo head: avvio Ray head su ${VLLM_HOST_IP}:${RAY_PORT}"
    ray start --head --port="${RAY_PORT}" --node-ip-address="${VLLM_HOST_IP}" --num-gpus="${NUM_GPUS}" \
        --dashboard-host=0.0.0.0

    : "${MODEL_PATH:?MODEL_PATH non impostato}"

    # TOKENIZER e TRUST_REMOTE_CODE sono opzionali: se non impostate, vllm
    # usa il tokenizer embedded nel GGUF e non esegue codice remoto.
    EXTRA_ARGS=()
    if [ -n "${TOKENIZER:-}" ]; then
        EXTRA_ARGS+=(--tokenizer "${TOKENIZER}")
    fi
    if [ "${TRUST_REMOTE_CODE:-0}" = "1" ]; then
        EXTRA_ARGS+=(--trust-remote-code)
    fi
    if [ -n "${SERVED_MODEL_NAME:-}" ]; then
        EXTRA_ARGS+=(--served-model-name "${SERVED_MODEL_NAME}")
    fi
    if [ "${ENABLE_AUTO_TOOL_CHOICE:-0}" = "1" ]; then
        EXTRA_ARGS+=(--enable-auto-tool-choice)
        : "${TOOL_CALL_PARSER:?TOOL_CALL_PARSER richiesto se ENABLE_AUTO_TOOL_CHOICE=1}"
        EXTRA_ARGS+=(--tool-call-parser "${TOOL_CALL_PARSER}")
    fi
    if [ -n "${REASONING_PARSER:-}" ]; then
        EXTRA_ARGS+=(--reasoning-parser "${REASONING_PARSER}")
    fi
    if [ -n "${ATTENTION_BACKEND:-}" ]; then
        EXTRA_ARGS+=(--attention-backend "${ATTENTION_BACKEND}")
    fi
    if [ -n "${CHAT_TEMPLATE:-}" ]; then
        EXTRA_ARGS+=(--chat-template "${CHAT_TEMPLATE}")
    fi

    echo "[entrypoint] Avvio vllm serve"
    exec vllm serve "${MODEL_PATH}" \
        "${EXTRA_ARGS[@]}" \
        --port "${SERVE_PORT:-7000}" \
        --host "${SERVE_HOST:-0.0.0.0}" \
        --max-model-len "${MAX_MODEL_LEN:-32768}" \
        --gpu-memory-utilization "${GPU_MEMORY_UTILIZATION:-0.9}" \
        --tensor-parallel-size "${TENSOR_PARALLEL_SIZE:-2}" \
        --distributed-executor-backend ray \
        --distributed-timeout-seconds "${DISTRIBUTED_TIMEOUT_SECONDS:-600}" \
        --enforce-eager \
        --dtype "${DTYPE:-float16}" \
        --kv-cache-dtype "${KV_CACHE_DTYPE:-fp8}"

else
    : "${VLLM_HOST_IP:?VLLM_HOST_IP non impostato}"
    : "${RAY_HEAD_ADDRESS:?RAY_HEAD_ADDRESS non impostato (IP del nodo head)}"
    echo "[entrypoint] Nodo worker: mi collego al Ray head su ${RAY_HEAD_ADDRESS}:${RAY_PORT}"
    ray start --address="${RAY_HEAD_ADDRESS}:${RAY_PORT}" --node-ip-address="${VLLM_HOST_IP}" --num-gpus="${NUM_GPUS}"

    echo "[entrypoint] Worker connesso, container resta attivo"
    exec sleep infinity
fi