#!/bin/bash

#
# /*
#  * Copyright (c) 2025 Huawei Technologies Co., Ltd.
#  * openFuyao is licensed under Mulan PSL v2.
#  * You can use this software according to the terms and conditions of the Mulan PSL v2.
#  * You may obtain a copy of Mulan PSL v2 at:
#  *          http://license.coscl.org.cn/MulanPSL2
#  * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
#  * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
#  * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
#  * See the Mulan PSL v2 for more details.
#  */
#

set -e

export DOCKER_BUILDKIT=1

function usage() {
    echo "Usage: $0 [image_version]"
    echo
    echo "Arguments:"
    echo "  image_version  Image version to build (default: 1.0.0)"
}

function validate_version() {
    [[ "$1" =~ ^[0-9]+\.[0-9]+\.[0-9]+$ ]] || return 1
    return 0
}

image_version=${1:-1.0.0}

if ! validate_version "$image_version"; then
    echo "Error: Invalid image version format. Use x.y.z format."
    usage
    exit 1
fi

echo "Image version: $image_version"

current_dir=$(
    cd "$(dirname "$0")" || exit 1
    pwd
)
top_dir=$(dirname "${current_dir}")
pkg_dir=${top_dir}/pkg

echo "current_dir=${current_dir}"
echo "top_dir=${top_dir}"
echo "pkg_dir=${pkg_dir}"

function arch_config() {
    arch=$(uname -m)
    if [[ ${arch} == "x86_64" ]]; then
        platform="x86_64"
    elif [[ ${arch} == "aarch64" ]]; then
        platform="aarch64"
    else
        echo "incorrect arch mode"
        exit 1
    fi
}

function mk_xpu_pkg_dir() {
    [ -e "${pkg_dir}" ] && rm -rf "${pkg_dir}"
    mkdir -p "${pkg_dir}"/images
    mkdir -p "${pkg_dir}"/chart
    chmod -R 750 "${pkg_dir}"
}

function build_chart_package(){
    cd ${top_dir}/charts && sed -i "s/^version: .*/version: ${image_version}/" vnpu/Chart.yaml && helm package vnpu
    cp -P --remove-destination -rf vxpu-${image_version}.tgz ${pkg_dir}/chart
}

function build_image() {
    local dockerfile=$1
    local image_name=$2
    local version=$3
    local tag="${image_name}:${version}"
    echo "build ${image_name} image begin"
    docker build \
        --squash \
        --no-cache \
        -f "${current_dir}/pipeline/Dockerfile/${dockerfile}" \
        --build-arg GOPROXY=https://goproxy.cn,direct \
        -t "${tag}" \
        "${top_dir}"
    docker save -o "${pkg_dir}/images/${image_name}-${version}-${platform}.tar" "${tag}"
    echo "build ${image_name} image end"
}

function extract_symbols() {
    local tmp_container
    tmp_container=$(docker create acl_client_update:${image_version})
    docker cp "${tmp_container}:/root/XPU_symbols.tar.gz" "${top_dir}/XPU_symbols.tar.gz"
    docker rm "${tmp_container}" >/dev/null
}

function build_output_packages() {
    cd "${pkg_dir}"
    [ -e "${top_dir}/output/software" ] && rm -rf "${top_dir}/output/software"
    mkdir -p ${top_dir}/output/software
    upload_arch=$(echo ${arch} | sed 's/_/-/g')
    zip -1 -r -y ${top_dir}/output/software/vxpu-${image_version}-${upload_arch}.zip *

    mkdir -p ${top_dir}/output/inner
    cp -P --remove-destination -rf ${top_dir}/XPU_symbols.tar.gz \
        ${top_dir}/output/inner/vxpu-${image_version}-${upload_arch}_sym.tar.gz
}

function main() {
    arch_config
    mk_xpu_pkg_dir
    build_chart_package
    build_image "acl-client-update.Dockerfile"      "acl_client_update"      ${image_version}
    extract_symbols
    build_image "npu-device-plugin.Dockerfile"       "npu_device_plugin"      ${image_version}
    build_image "xpu-exporter.Dockerfile"            "xpu_exporter"           ${image_version}
    build_image "vc-controller-manager.Dockerfile"   "vc_controller_manager"  "1.9.0"
    build_image "vc-scheduler.Dockerfile"            "vc_scheduler"           "1.9.0"
    build_image "vc-webhook-manager.Dockerfile"      "vc_webhook_manager"     "1.9.0"
    build_output_packages
}

main "$@"