#!/bin/bash
# Display SafeLine ASCII Art Logo
echo "
  ____             __          _       _
 / ___|    __ _   / _|   ___  | |     (_)  _ __     ___
 \___ \   / _\` | | |_   / _ \ | |     | | | '_ \   / _ \\
  ___) | | (_| | |  _| |  __/ | |___  | | | | | | |  __/
 |____/   \__,_| |_|    \___| |_____| |_| |_| |_|  \___|
"

ARCH=$(uname -m)
SAFELINE_PATH="/data/safeline"
RESOLV_CONF="/etc/resolv.conf"
COMPOSE_CMD="docker compose"

# 设置频道版本
RELEASE="${RELEASE}"
# 如果用户提供了 RELEASE 值，则在前面加上 "-"
if [ -n "$RELEASE" ]; then
    RELEASE="-$RELEASE"
fi

# Display QR code function
qrcode() {
    echo "█████████████████████████████████████████"
    echo "█████████████████████████████████████████"
    echo "████ ▄▄▄▄▄ █▀ █▀▀██▀▄▀▀▄▀▄▀▄██ ▄▄▄▄▄ ████"
    echo "████ █   █ █▀ ▄ █▀▄▄▀▀ ▄█▄  ▀█ █   █ ████"
    echo "████ █▄▄▄█ █▀█ █▄█▄▀▀▄▀▄ ▀▀▄▄█ █▄▄▄█ ████"
    echo "████▄▄▄▄▄▄▄█▄█▄█ █▄▀ █ ▀▄▀ █▄█▄▄▄▄▄▄▄████"
    echo "████▄ ▄▄ █▄▄  ▄█▄▄▄▄▀▄▀▀▄██ ▄▄▀▄█▄▀ ▀████"
    echo "████▄ ▄▀▄ ▄▀▄ ▀ ▄█▀ ▀▄ █▀▀ ▀█▀▄██▄▀▄█████"
    echo "█████ ▀▄█ ▄ ▄▄▀▄▀▀█▄▀▄▄▀▄▀▄ ▄ ▀▄▄▄█▀▀████"
    echo "████ █▀▄▀ ▄▀▄▄▀█▀ ▄▄ █▄█▀▀▄▀▀█▄█▄█▀▄█████"
    echo "████ █ ▀  ▄▀▀ ██▄█▄▄▄▄▄▀▄▀▀▀▄▄▀█▄▀█ ▀████"
    echo "████ █ ▀▄ ▄██▀▀ ▄█▀ ▀███▄  ▀▄▀▄▄ ▄▀▄█████"
    echo "████▀▄▄█  ▄▀▄▀ ▄▀▀▀▄▀▄▀ ▄▀▄  ▄▀ ▄▀█ ▀████"
    echo "████ █ █ █▄▀ █▄█▀ ▄▄███▀▀▀▄█▀▄ ▀  ▀▄█████"
    echo "████▄███▄█▄▄▀▄ █▄█▄▄▄▄▀▀▄█▀▀ ▄▄▄  ▀█ ████"
    echo "████ ▄▄▄▄▄ █▄▀█ ▄█▀▄ █▀█▄ ▀  █▄█  ▀▄▀████"
    echo "████ █   █ █  █▄▀▀▀▄▄▄▀▀▀▀▀▀ ▄▄  ▀█  ████"
    echo "████ █▄▄▄█ █  ▀█▀ ▄▄▄▄ ▀█ ▀▀▄▀ ▀▀ ▀██████"
    echo "████▄▄▄▄▄▄▄█▄▄██▄█▄▄█▄██▄██▄▄█▄▄█▄█▄█████"
    echo "█████████████████████████████████████████"
    echo "█████████████████████████████████████████"
    echo
    echo "微信扫描上方二维码加入雷池项目讨论组"
}

# 获取当前安装目录
get_installed_directory() {
    container_id=$(docker ps -n 1 --filter name=.*safeline-mgt.* --format '{{.ID}}')
    SAFELINE_PATH=$(docker inspect --format '{{index .Config.Labels "com.docker.compose.project.working_dir"}}' $container_id)
    if [[ -z "$SAFELINE_PATH" || ! -d "$SAFELINE_PATH" ]]; then
        warning "未能找到安装目录"
        echo -e -n "\033[34m[SafeLine] 请手动输入安装目录: \033[0m"
        read input_path
        SAFELINE_PATH=$input_path
    fi
    echo "$SAFELINE_PATH"
}

# Display messages
info() {
   echo -e "\033[32m[SafeLine] $*\033[0m"
}

warning() {
    echo -e "\033[33m[SafeLine] $*\033[0m"
}

abort() {
    qrcode
    echo -e "\033[31m[SafeLine] $*\033[0m"
    exit 1
}

trap 'onexit' INT
onexit() {
    echo
    abort "用户手动结束升级"
}

# Start Docker service
start_docker() {
    systemctl enable docker
    systemctl daemon-reload
    systemctl restart docker || abort "Docker 服务启动失败"
}

# Check Docker Compose availability
check_docker_compose() {
    COMPOSE_CMD="docker compose"
    $COMPOSE_CMD version || {
        COMPOSE_CMD="docker-compose"
        $COMPOSE_CMD version || start_docker
    }
}

get_average_delay() {
    local source=$1
    local total_delay=0
    local iterations=3

    for ((i = 0; i < iterations; i++)); do
        # check timeout
        if ! curl -o /dev/null -m 1 -s -w "%{http_code}\n" "$source" > /dev/null; then
            delay=999
        else
            delay=$(curl -o /dev/null -s -w "%{time_total}\n" "$source")
        fi
        total_delay=$(awk "BEGIN {print $total_delay + $delay}")
    done

    average_delay=$(awk "BEGIN {print $total_delay / $iterations}")
    echo "$average_delay"
}

# 选择最佳的镜像源
choose_best_image_source() {
    sources=("https://registry-1.docker.io" "https://swr.cn-east-3.myhuaweicloud.com")
    best_source=""
    lowest_delay=999  # 初始为一个很大的延迟值

    for source in "${sources[@]}"; do
        delay=$(get_average_delay "$source")

        if (( $(awk 'BEGIN {print '"$delay"' < '"$lowest_delay"' }') )); then
            lowest_delay=$delay
            best_source=$source
        fi
    done

    if [ -z "$best_source" ]; then
        abort "无法连接到任何镜像源"
    fi

    if [ "$best_source" == "https://registry-1.docker.io" ]; then
        image_prefix="chaitin"
    fi
    if [ "$best_source" == "https://swr.cn-east-3.myhuaweicloud.com" ]; then
        image_prefix="swr.cn-east-3.myhuaweicloud.com/chaitin-safeline"
    fi

    echo "$image_prefix"
}

# Download and set up SafeLine files
setup_safeline_files() {
    cd "$SAFELINE_PATH" || exit
    info "正在下载 compose.yaml 脚本..."
    curl -4sSLk https://waf-ce.chaitin.cn/release/latest/compose.yaml -o compose.yaml || abort "下载 compose.yaml 脚本失败"
    info "正在下载 reset_tengine.sh 脚本..."
    curl -4sSLk https://waf-ce.chaitin.cn/release/latest/reset_tengine.sh -o reset_tengine.sh || abort "下载 reset_tengine.sh 脚本失败"

    # 选择最合适的镜像源
    IMAGE_PREFIX=$(choose_best_image_source)

    # 读取 .env 文件并设置默认值
    if [ ! -f .env ]; then
        touch .env
    fi

    # 用默认值更新 .env 文件中的变量，如果没有定义
    grep -q '^SAFELINE_DIR=' .env || echo -e -n "\nSAFELINE_DIR=$SAFELINE_PATH" >> .env
    grep -q '^IMAGE_TAG=' .env || echo -e -n "\nIMAGE_TAG=latest" >> .env
    grep -q '^MGT_PORT=' .env || echo -e -n "\nMGT_PORT=9443" >> .env
    grep -q '^POSTGRES_PASSWORD=' .env || echo -e -n "\nPOSTGRES_PASSWORD=$(LC_ALL=C tr -dc A-Za-z0-9 </dev/urandom | head -c 32)" >> .env
    grep -q '^SUBNET_PREFIX=' .env || echo -e -n "\nSUBNET_PREFIX=172.22.222" >> .env
    grep -q '^ARCH_SUFFIX=' .env || echo -e -n "\nARCH_SUFFIX=$([[ "$ARCH" =~ "aarch" || "$ARCH" =~ "arm" ]] && echo "-arm" || echo "")" >> .env
    grep -q '^IMAGE_PREFIX=' .env || echo -e -n "\nIMAGE_PREFIX=$IMAGE_PREFIX" >> .env
    grep -q '^RELEASE=' .env || echo -e -n "\nRELEASE=$RELEASE" >> .env

    info "配置 .env 文件完成"
}


# Run Docker Compose and validate container health
run_compose() {
    info "正在拉取雷池镜像..."
    $COMPOSE_CMD pull || abort "拉取 Docker 镜像失败"
    info "正在升级雷池服务..."
    $COMPOSE_CMD up -d --remove-orphans || abort "启动 Docker 容器失败"
    for service in safeline-pg safeline-mgt; do
        retry=0
        until docker inspect --format='{{.State.Health.Status}}' $service 2>/dev/null | grep -q "healthy" || ((retry++ > 30)); do
            sleep 5
        done
        [[ $retry -gt 30 ]] && abort "$service 容器未正常启动"
    done
    info "雷池服务升级成功"
}

# Main script execution
get_installed_directory
check_docker_compose
setup_safeline_files
run_compose
qrcode

docker exec safeline-mgt /app/mgt-cli reset-admin --once

# Display final access URLs
info "雷池 WAF 社区版升级成功，请访问以下地址访问控制台"
for ip in $(hostname -I); do
    warning "https://$ip:9443/"
done
