#!/bin/bash

# Check if an argument is provided
if [ $# -eq 0 ]; then
    echo "Usage: $0 <input_path>" >&2
    exit 1
fi

input_path="$1"
input_file="${input_path}/temp-hostfile"
output_file="${input_path}/deepspeed-hostfile"

# Check if input file exists
if [ ! -f "$input_file" ]; then
    echo "Error: Input file '${input_file}' not found" >&2
    exit 1
fi

# 读取输入文件
content=$(cat "$input_file")
echo "Input content: $content" >&2

# input_file="temp-hostfile"
# output_file="deepspeed-hostfile"

# # 读取输入文件
# content=$(cat "$input_file")
# echo "Input content: $content" >&2

# 清空输出文件
> "$output_file"

# 函数：生成指定长度的数字序列，保留前导零
generate_sequence() {
    local start=$1
    local end=$2
    local length=${#start}  # 使用起始数字的长度作为填充长度

    for ((i=10#$start; i<=10#$end; i++)); do
        printf "%0${length}d\n" $i
    done
}

# 使用正则表达式匹配所有节点组
while [[ $content =~ ([a-zA-Z0-9-]+(-[a-zA-Z0-9-]+)?(\[[0-9,.-]+\])?)(,|$) ]]; do
    group="${BASH_REMATCH[1]}"
    echo "Processing group: $group" >&2
    
    if [[ $group =~ ^([a-zA-Z0-9-]+)(-[a-zA-Z0-9-]+)?(\[([0-9,.-]+)\])?$ ]]; then
        prefix="${BASH_REMATCH[1]}"
        middle="${BASH_REMATCH[2]}"
        nodes="${BASH_REMATCH[4]}"

        echo "Prefix: $prefix" >&2
        echo "Middle: $middle" >&2
        echo "Nodes: $nodes" >&2

        # 移除middle开头的连字符（如果存在）
        middle="${middle#-}"

        if [[ -n $nodes ]]; then
            echo "Processing nodes in brackets" >&2
            IFS=',' read -ra node_array <<< "$nodes"
            for node in "${node_array[@]}"; do
                if [[ $node == *-* ]]; then
                    echo "Range format detected in $node" >&2
                    IFS='-' read -ra range <<< "$node"
                    start="${range[0]}"
                    end="${range[1]}"
                    for i in $(generate_sequence "$start" "$end"); do
                        node_name="${prefix}${middle:+-$middle}-${i}"
                        node_name="${node_name/--/-}"  # 替换双连字符为单连字符
                        echo "${node_name} slots=8" >> "$output_file"
                        echo "Added: ${node_name}" >&2
                    done
                else
                    node_name="${prefix}${middle:+-$middle}-${node}"
                    node_name="${node_name/--/-}"  # 替换双连字符为单连字符
                    echo "${node_name} slots=8" >> "$output_file"
                    echo "Added: ${node_name}" >&2
                fi
            done
        else
            echo "Processing single node" >&2
            node_name="${prefix}${middle:+-$middle}"
            node_name="${node_name/--/-}"  # 替换双连字符为单连字符
            echo "${node_name} slots=8" >> "$output_file"
            echo "Added: ${node_name}" >&2
        fi
    else
        echo "Invalid format for node group: $group" >&2
    fi

    # 从内容中移除已处理的部分
    content=${content#*"${BASH_REMATCH[0]}"}
done

echo "Expanded hostfile has been written to $output_file" >&2

# 检查输出文件是否为空
if [ -s "$output_file" ]; then
    # 处理master node和slave node list 供deepspeed启动使用
    first_node=$(head -n 1 "$output_file" | cut -d ' ' -f1)
    last_node=$(tail -n 1 "$output_file" | cut -d ' ' -f1)

    echo "Master Node = ${first_node}" >&2
    echo "Slave Node = ${last_node}" >&2
    echo "${first_node} ${last_node}"
else
    echo "Error: Output file is empty" >&2
    exit 1
fi
