#!/bin/bash

#修改all_config.yaml中的ip配置
CONFIG_FILE="./src/all_config.yaml"

USER="root"
IMAGE="swr.cn-south-1.myhuaweicloud.com/ascendhub/mindie:2.0.T3-800I-A2-py311-openeuler24.03-lts"
MOUNT_DIR="/data01/deepseek"
CONTAINER_DIR="/deepseek"
CONTAINER_NAME="deepseek_atb"
MINDIE_ASCEND_DEPLOY_DIR="mindie-ascend-deploy"
IPS=$(awk '
BEGIN { RS = "[\n]"; FS = ":"; }
/all_ip:/ {
    while (getline) {
        if ($0 ~ /^\s+-/) {
            gsub(/^[ \t]+-\s*\[/, "");
            gsub(/\]\s*$/, "");
            split($0, ip_pair, ", ");
            gsub(/"/, "", ip_pair[1]);
            gsub(/"/, "", ip_pair[2]);
            print ip_pair[1] " " ip_pair[2];
        } else {
            break;
        }
    }
}
' "$CONFIG_FILE")


check_ssh_key() {
    local ip=$1
    ssh -o BatchMode=yes -o ConnectTimeout=5 $ip echo "SSH_OK" &> /dev/null
    if [ $? -ne 0 ]; then
        echo "Error: SSH key not set up for $ip. please use ssh_rank_table.py or manually configure the ip."
        exit 1
    fi
}

execute_container_commands() {
    local IP=$1
    local MIES_CONTAINER_IP=$2
    local use_ssh=false

    if [ "$IP" == "$(echo "$IPS" | head -n 1 | awk '{print $1}')" ]; then
        echo "Processing master node at IP: $IP with container IP: $MIES_CONTAINER_IP"
    else
        echo "Processing node at IP: $IP with container IP: $MIES_CONTAINER_IP"
        check_ssh_key $IP
        use_ssh=true
    fi

    docker_command= "docker run -itd -u 0 -e MIES_CONTAINER_IP=$MIES_CONTAINER_IP --ipc=host --network host \
        --name $CONTAINER_NAME \
        --privileged \
        --device=/dev/davinci0 \
        --device=/dev/davinci1 \
        --device=/dev/davinci2 \
        --device=/dev/davinci3 \
        --device=/dev/davinci4 \
        --device=/dev/davinci5 \
        --device=/dev/davinci6 \
        --device=/dev/davinci7 \
        --device=/dev/davinci_manager \
        --device=/dev/devmm_svm \
        --device=/dev/hisi_hdc \
        -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
        -v /usr/local/sbin:/usr/local/sbin \
        -v $MOUNT_DIR:$CONTAINER_DIR \
        $IMAGE \
        /bin/bash"

    if $use_ssh; then
        CONTAINER_ID=$(ssh $USER@$IP "$docker_command")
    else
        CONTAINER_ID=$(eval "$docker_command")
    fi
    exec_in_container "if command -v yum &> /dev/null; then yum install -y jq; elif command -v apt-get &> /dev/null; then apt-get update && apt-get install -y jq; fi"

    exec_in_container "cd $CONTAINER_DIR/$MINDIE_ASCEND_DEPLOY_DIR && if ! command -v jq &> /dev/null; then cp ./jq-linux-arm64 /usr/bin/jq && chmod +x /usr/bin/jq; fi && pip install -r ./src/requirements.txt && bash start_mindie.sh"


exec_in_container() {
    local command=$1
    if $use_ssh; then
        ssh $USER$IP "docker exec -it $CONTAINER_ID bash -c '$command'"
    else
        docker exec -it $CONTAINER_ID bash -c "$command"
    fi
}

while read -r IP MIES_CONTAINER_IP; do
    execute_container_commands "$IP" "$MIES_CONTAINER_IP"
done <<< "$IPS"