#!/usr/bin/env bash
# SPDX-License-Identifier: (GPL-2.0-only OR MIT)
# Copyright (C) 2026 Amlogic, Inc. All rights reserved
################################################################################
# Amlbian NPU Compilation Module
#
# Purpose: NPU compilation with Hook support and aml-comp integration
# Provides: NPU compilation, driver and library preparation
################################################################################

set -e

# Get script directory
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
AMLBIAN_DIR="$(dirname "$(dirname "${SCRIPT_DIR}")")"
BASE_DIR="$(dirname "$(dirname "${AMLBIAN_DIR}")")"

display_alert "SCRIPT_DIR ${SCRIPT_DIR} AMLBIAN_DIR ${AMLBIAN_DIR} BASE_DIR ${BASE_DIR}" "info"

# Source common functions
if [[ -f "${AMLBIAN_DIR}/lib/function-common.sh" ]]; then
    source "${AMLBIAN_DIR}/lib/function-common.sh"
fi

# Source ADLA build script if it exists
ADLA_SCRIPT="${SCRIPT_DIR}/build-adla.sh"
if [[ -f "${ADLA_SCRIPT}" ]]; then
    display_alert "Loading ADLA build script" "${ADLA_SCRIPT}" "info"
    source "${ADLA_SCRIPT}"
else
    display_alert "ADLA build script not found" "${ADLA_SCRIPT}" "warn"
fi

# ==============================================================================
# Shared Helper Functions
# ==============================================================================

# ------------------------------------------------------------------------------
# Step 1: Version and Package Information Definition
# ------------------------------------------------------------------------------

# Initialize common package metadata for NPU common library
function _npu_common_init_package_info() {
    PKG_MAINTAINER="Armbian Build <build@armbian.com>"
    PKG_DESCRIPTION="Common library for Amlogic NPU"
    PKG_NAME="aml-npu-common"
    PKG_VERSION="1.0"
    PKG_ARCH="arm64"
    LIB_DIR="aarch64-linux-gnu"

    NPU_COMM_DIR="${BASE_DIR}/aml-comp/vendor/amlogic/aml_npu_common"
}

# Initialize common package metadata for NPU driver
function _npu_driver_init_package_info() {
    PKG_MAINTAINER="Armbian Build <build@armbian.com>"
    PKG_DESCRIPTION="Amlogic NPU driver and libraries"
    PKG_NAME="aml-npu"
    PKG_VERSION="1.0"
    PKG_ARCH="arm64"
    NPU_MODULE="galcore.ko"

    # Cross compilation settings
    CROSS_COMPILE="${CROSS_COMPILE:-aarch64-linux-gnu-}"
}

# ------------------------------------------------------------------------------
# Step 2: Build Directory Structure Setup
# ------------------------------------------------------------------------------

# Create layered build directory structure
# Parameters:
#   $1 - base name (e.g. "${PKG_NAME}_${PKG_VERSION}_${PKG_ARCH}")
# Sets: TMP_DIR, BASE_BUILD_DIR, SOURCE_DIR, BUILD_DIR, INSTALL_DIR, PATCH_DIR,
#       DEB_PACKAGE_DIR, DEBIAN_DIR
function _npu_create_build_dirs() {
    local base_name="$1"

    TMP_DIR="${AMLBIAN_DIR}/.tmp"
    mkdir -p "${TMP_DIR}"

    # 基础构建目录
    BASE_BUILD_DIR="${TMP_DIR}/${base_name}"

	rm -rf ${BASE_BUILD_DIR}

    # 分层目录
    SOURCE_DIR="${BASE_BUILD_DIR}/source"      # 源码目录
    BUILD_DIR="${BASE_BUILD_DIR}/build"        # 编译目录
    INSTALL_DIR="${BASE_BUILD_DIR}/install"    # 安装目录
    PATCH_DIR="${BASE_BUILD_DIR}/patch"        # 补丁目录

    # 创建所有目录
    mkdir -p "${SOURCE_DIR}"
    mkdir -p "${BUILD_DIR}"
    mkdir -p "${INSTALL_DIR}"
    mkdir -p "${PATCH_DIR}"

    # 在BUILD_DIR下创建deb/包名目录
    DEB_PACKAGE_DIR="${BUILD_DIR}/deb/${base_name}"
    mkdir -p "${DEB_PACKAGE_DIR}"

    # 安装目录的子目录
    DEBIAN_DIR="${INSTALL_DIR}/DEBIAN"
    mkdir -p "${DEBIAN_DIR}"
}

# Create standard install subdirectories for NPU common library
function _npu_common_create_install_dirs() {
    mkdir -p "${INSTALL_DIR}/usr/include"
    mkdir -p "${INSTALL_DIR}/usr/lib"
}

# Create standard install subdirectories for NPU driver
function _npu_driver_create_install_dirs() {
    local kernel_version="$1"

    mkdir -p "${INSTALL_DIR}/lib/modules/${kernel_version}/kernel/amlogic/npu/"
    mkdir -p "${INSTALL_DIR}/usr/lib/"
    mkdir -p "${INSTALL_DIR}/etc/npu"
}

# ------------------------------------------------------------------------------
# Step 3: Debian Packaging Scripts Creation
# ------------------------------------------------------------------------------

# Create Debian control file
# Parameters:
#   $1 - package name
#   $2 - package version
#   $3 - package architecture
#   $4 - maintainer
#   $5 - description
#   $6 - dependencies (optional)
#   $7 - conflicts (optional)
#   $8 - provides (optional)
#   $9 - section (optional, default "libs")
function _npu_create_control_file() {
    local pkg_name="$1"
    local pkg_version="$2"
    local pkg_arch="$3"
    local maintainer="$4"
    local description="$5"
    local depends="${6:-}"
    local conflicts="${7:-}"
    local provides="${8:-}"
    local section="${9:-libs}"

    display_alert "COMPILE-NPU Creating control file DEBIAN/control" "info"
    cat > "${DEBIAN_DIR}/control" <<EOF
Package: ${pkg_name}
Version: ${pkg_version}
Architecture: ${pkg_arch}
Maintainer: ${maintainer}
Depends: ${depends}
Conflicts: ${conflicts}
Provides: ${provides}
Section: ${section}
Priority: optional
Description: ${description}
EOF
}

# Create postinst script for NPU common library
function _npu_common_create_postinst() {
    display_alert "COMPILE-NPU-COMMON Creating post-installation script DEBIAN/postinst" "info"
    cat > "${DEBIAN_DIR}/postinst" << EOF
#!/bin/sh
set -e

if [ "\$1" = "configure" ]; then
    # Run ldconfig to update library cache
    ldconfig

    echo "NPU common library installed successfully."
    echo "Header files are installed in /usr/include/"
    echo "Library files are installed in /usr/lib/"
fi

exit 0
EOF
    chmod 755 "${DEBIAN_DIR}/postinst"
}

# Create prerm script for NPU common library
function _npu_common_create_prerm() {
    display_alert "COMPILE-NPU-COMMON Creating pre-removal script DEBIAN/prerm" "info"
    cat > "${DEBIAN_DIR}/prerm" << EOF
#!/bin/sh
set -e

if [ "\$1" = "remove" ] || [ "\$1" = "deconfigure" ]; then
    # Run ldconfig to update library cache
    ldconfig

    echo "NPU common library removed successfully."
fi

exit 0
EOF
    chmod 755 "${DEBIAN_DIR}/prerm"
}

# Create postinst script for NPU driver
function _npu_driver_create_postinst() {
    local kernel_version="$1"
    display_alert "COMPILE-NPU-DRIVER Creating post-installation script DEBIAN/postinst" "info"
    cat > "${DEBIAN_DIR}/postinst" << EOF
#!/bin/bash
set -e

# Run depmod to update module dependencies
# Clear any KERNELDIR environment variable that might interfere
unset KERNELDIR
unset STAGING_KERNEL_DIR

if [ -d /lib/modules ]; then
    for kernel in /lib/modules/*; do
        if [ -d "\${kernel}" ]; then
            # Use depmod with explicit kernel version
            depmod -a "${kernel_version}" 2>/dev/null || {
                echo "Warning: depmod failed for kernel "${kernel_version}")"
                echo "This is normal if the kernel is not currently running"
            }
        fi
    done
fi

# Update library cache
ldconfig

echo "Amlogic NPU driver installed successfully."
EOF
    chmod 0755 "${DEBIAN_DIR}/postinst"
}

# Create prerm script for NPU driver
function _npu_driver_create_prerm() {
    display_alert "COMPILE-NPU-DRIVER Creating pre-removal script DEBIAN/prerm" "info"
    cat > "${DEBIAN_DIR}/prerm" << 'EOF'
#!/bin/bash
set -e

echo "Removing Amlogic NPU driver..."
EOF
    chmod 0755 "${DEBIAN_DIR}/prerm"
}

# ------------------------------------------------------------------------------
# Step 4: Source/Vendor File Copying
# ------------------------------------------------------------------------------

# Copy NPU common vendor files to install directory
function _npu_common_copy_files() {
    display_alert "COMPILE-NPU-COMMON Copying files from vendor binaries ${NPU_COMM_DIR}" "info"

    # Determine the correct library path based on configuration
    local AML_CUST_SPEC="${AML_CUST_SPEC:-}"
    local HOST_ARCH="${HOST_ARCH:-aarch64}"

    # Set ARM_TARGET based on architecture
    local ARM_TARGET="lib32"
    if [ "${HOST_ARCH}" = "aarch64" ]; then
        ARM_TARGET="lib64"
    fi

    # Install header files
    display_alert "COMPILE-NPU-COMMON Installing NPU common header files" "info"
    if [ -d "${NPU_COMM_DIR}/nnsdk/include" ]; then
        install -m 0644 "${NPU_COMM_DIR}"/nnsdk/include/* "${INSTALL_DIR}/usr/include/"
    fi

    if [ -d "${NPU_COMM_DIR}/llmsdk/include" ]; then
        install -m 0644 "${NPU_COMM_DIR}"/llmsdk/include/* "${INSTALL_DIR}/usr/include/"
    fi


    display_alert "COMPILE-NPU-COMMON Library source path: ${LIB_SOURCE_PATH}" "info"
    if [ "${AML_CUST_SPEC}" = "CUST_L_HUB_F" ]; then
        LIB_SOURCE_PATH="${NPU_COMM_DIR}/nnsdk/linux/starf/${ARM_TARGET}"
        safe_install -m 0644 "${LIB_SOURCE_PATH}"/*.so* "${INSTALL_DIR}/usr/lib/"
    else
        LIB_SOURCE_PATH="${NPU_COMM_DIR}/nnsdk/linux/yocto/${ARM_TARGET}"
        safe_install -m 0644 "${LIB_SOURCE_PATH}"/*.so* "${INSTALL_DIR}/usr/lib/"
        safe_install -m 0644 "${NPU_COMM_DIR}"/llmsdk/linux/yocto/${ARM_TARGET}/*.so* "${INSTALL_DIR}/usr/lib/"
    fi
}

# Copy NPU driver files to install directory
function _npu_driver_copy_files() {
    display_alert "COMPILE-NPU-DRIVER Copying NPU driver files" "info"

    local kernel_version="$1"
    local module_dir="${INSTALL_DIR}/lib/modules/${kernel_version}/kernel/amlogic/npu/"

    # Check if NPU module was built
    if [[ -f "${SOURCE_DIR}/build/sdk/drivers/${NPU_MODULE}" ]]; then
        display_alert "COMPILE-NPU-DRIVER NPU module built successfully: ${NPU_MODULE}" "info"

        # Install NPU module to INSTALL_DIR
        display_alert "COMPILE-NPU-DRIVER Installing NPU module to install directory ${module_dir}" "info"
        install -m 0644 "${SOURCE_DIR}/build/sdk/drivers/${NPU_MODULE}" "${module_dir}/"

        # Install NPU libraries based on architecture to INSTALL_DIR
        display_alert "COMPILE-NPU-DRIVER Installing NPU libraries for ${PKG_ARCH} architecture" "info"

        if [[ "${PKG_ARCH}" == "arm64" ]]; then
            # Install 64-bit libraries
            LIB_SRC_DIR="${SOURCE_DIR}/sharelib/lib64"
            if [[ -d "${LIB_SRC_DIR}" ]]; then
                for lib in "${LIB_SRC_DIR}"/*.so*; do
                    if [[ -f "${lib}" ]]; then
                        install -m 0644 "${lib}" "${INSTALL_DIR}/usr/lib/"
                    fi
                done
            fi
        else
            # Install 32-bit libraries
            LIB_SRC_DIR="${SOURCE_DIR}/sharelib/lib32"
            if [[ -d "${LIB_SRC_DIR}" ]]; then
                for lib in "${LIB_SRC_DIR}"/*.so*; do
                    if [[ -f "${lib}" ]]; then
                        install -m 0644 "${lib}" "${INSTALL_DIR}/usr/lib/"
                    fi
                done
            fi
        fi

        # Create NPU configuration directory in INSTALL_DIR
        CONFIG_DIR="${INSTALL_DIR}/etc/npu"
        mkdir -p "${CONFIG_DIR}"

        # Create a sample configuration file
        cat > "${CONFIG_DIR}/npu.conf" << EOF
# Amlogic NPU Configuration
NPU_ENABLE=1
NPU_MODE=performance
NPU_FREQUENCY=800
EOF

        display_alert "COMPILE-NPU-DRIVER NPU configuration files created" "info"
    else
        display_alert "COMPILE-NPU-DRIVER NPU module build failed. Could not find ${NPU_MODULE}" "err"
        exit 1
    fi
}

# ------------------------------------------------------------------------------
# Step 5: Debian Package Building
# ------------------------------------------------------------------------------

# Build Debian package from install directory
# Parameters:
#   $1 - log prefix (e.g. "COMPILE-NPU-COMMON" or "COMPILE-NPU-DRIVER")
#   $2 - output deb filename (without path)
# Sets: OUTPUT_DEB
function _npu_build_deb_package() {
    local log_prefix="$1"
    local output_deb_name="$2"

    # 将install目录内容拷贝到deb目录
    display_alert "Copying install directory to deb package directory ${log_prefix}" "info"
    cp -r "${INSTALL_DIR}/"* "${DEB_PACKAGE_DIR}/" 2>&1

    # 基于deb目录打包
    display_alert "Building Debian package from deb directory ${log_prefix}" "info"
    dpkg-deb --build "${DEB_PACKAGE_DIR}" 2>&1

    # Move to output directory
    OUTPUT_DEB="${SRC}/output/${output_deb_name}"
    mkdir -p "$(dirname "${OUTPUT_DEB}")"
    mv "${DEB_PACKAGE_DIR}.deb" "${OUTPUT_DEB}"

    display_alert "${log_prefix} Package built successfully ${output_deb_name}" "info"
    display_alert "${log_prefix} Build directory preserved ${BASE_BUILD_DIR}" "info"
}

# Display build directory info
# Parameters:
#   $1 - log prefix
function _npu_display_build_info() {
    local log_prefix="$1"
    display_alert "${log_prefix} Source dir: ${SOURCE_DIR}" "info"
    display_alert "${log_prefix} Build dir: ${BUILD_DIR}" "info"
    display_alert "${log_prefix} Install dir: ${INSTALL_DIR}" "info"
    display_alert "${log_prefix} Patch dir: ${PATCH_DIR}" "info"
    display_alert "${log_prefix} Deb package dir: ${DEB_PACKAGE_DIR}" "info"
}

# ==============================================================================
# Main Build Functions (Entry Points)
# ==============================================================================

# ------------------------------------------------------------------------------
# NPU Common Library Build
# Entry point: amlbian_build_npu_common
# ------------------------------------------------------------------------------

function amlbian_build_npu_common() {
    # ---- Step 1: Version and Package Information ----
    _npu_common_init_package_info

    # amlbian_checkout_toolchain

    # Get kernel version for packaging
    local KERNEL_VERSION=$(make -s -C "${BASE_DIR}/aml-comp/kernel/aml-6.12/build" kernelrelease 2>/dev/null || true)
    if [[ -z "${KERNEL_VERSION}" ]]; then
        KERNEL_VERSION="unknown-$(date +%Y%m%d)"
    fi
    display_alert "COMPILE-NPU-COMMON Building NPU common library for kernel ${KERNEL_VERSION}" "info"

    # ---- Step 2: Build Directory Setup ----
    _npu_create_build_dirs "${PKG_NAME}_${PKG_VERSION}_${PKG_ARCH}"
    _npu_common_create_install_dirs

    # Display info
    display_alert "COMPILE-NPU-COMMON Building NPU common library version ${PKG_VERSION} for ${PKG_ARCH}" "info"
    _npu_display_build_info "COMPILE-NPU-COMMON"

    # ---- Step 3: Debian Packaging Scripts ----
    _npu_create_control_file \
        "${PKG_NAME}" "${PKG_VERSION}" "${PKG_ARCH}" "${PKG_MAINTAINER}" \
        "${PKG_DESCRIPTION}" "" "" "" "libs"

    _npu_common_create_postinst
    _npu_common_create_prerm

    # ---- Step 4: Copy Source/Vendor Files ----
    _npu_common_copy_files

    # ---- Step 5: Build Debian Package ----
    _npu_build_deb_package "COMPILE-NPU-COMMON" "${PKG_NAME}_${PKG_VERSION}_${PKG_ARCH}.deb"

    return 0
}

# ------------------------------------------------------------------------------
# NPU Driver Build
# Entry point: amlbian_build_npu_driver
# ------------------------------------------------------------------------------

function amlbian_build_npu_driver() {
    set -e

    # ---- Step 1: Version and Package Information ----
    _npu_driver_init_package_info
    # amlbian_checkout_toolchain

    # Get kernel version for packaging
    local KERNEL_VERSION=$(make -s -C "${BASE_DIR}/aml-comp/kernel/aml-6.12/build" kernelrelease 2>/dev/null || true)
    if [[ -z "${KERNEL_VERSION}" ]]; then
        KERNEL_VERSION="unknown-$(date +%Y%m%d)"
    fi
    display_alert "COMPILE-NPU-DRIVER Building NPU driver for kernel ${KERNEL_VERSION}" "info"

    # Source directory
    local SRC_DIR="${BASE_DIR}/aml-comp/hardware/amlogic/npu/nanoq"

    # Check if source exists
    if [[ ! -d "${SRC_DIR}" ]]; then
        display_alert "COMPILE-NPU-DRIVER NPU driver source code not found. Expected at: ${SRC_DIR}" "error"
        exit 1
    fi

    # Check for build script
    if [[ ! -f "${SRC_DIR}/aml_buildroot.sh" ]]; then
        display_alert "COMPILE-NPU-DRIVER NPU build script not found. Expected: ${SRC_DIR}/aml_buildroot.sh" "error"
        exit 1
    fi

    # ---- Step 2: Build Directory Setup ----
    _npu_create_build_dirs "${PKG_NAME}_${PKG_VERSION}_${PKG_ARCH}"
    _npu_driver_create_install_dirs "${KERNEL_VERSION}"
    _npu_display_build_info "COMPILE-NPU-DRIVER"

    # Copy source to source directory
    display_alert "COMPILE-NPU-DRIVER Copying source files to source directory" "info"
    cp -r "${SRC_DIR}/"* "${SOURCE_DIR}/"

    # Set up environment variables for build
    export ARCH="${PKG_ARCH}"
    export CROSS_COMPILE="${CROSS_COMPILE}"
    export KERNELDIR="${BASE_DIR}/aml-comp/kernel/aml-6.12"
    export STAGING_KERNEL_DIR="${KERNELDIR}"

    # Check for kernel headers
    if [[ ! -d "${KERNELDIR}" ]]; then
        display_alert "COMPILE-NPU-DRIVER Kernel headers directory not found at: ${KERNELDIR}" "error"
        exit 1
    fi

    # Change to source directory
    cd "${SOURCE_DIR}"

    # Build the NPU driver
    display_alert "COMPILE-NPU-DRIVER Building NPU kernel module and libraries" "info"

    # Execute the build script with appropriate parameters
    ./aml_buildroot.sh "${PKG_ARCH}" "${KERNELDIR}/build" "/usr/bin/aarch64-linux-gnu-"

    # ---- Step 3: Debian Packaging Scripts ----
    _npu_create_control_file \
        "${PKG_NAME}" "${PKG_VERSION}" "${PKG_ARCH}" "${PKG_MAINTAINER}" \
        "${PKG_DESCRIPTION}" "" "" "" "kernel"

    _npu_driver_create_postinst "${KERNEL_VERSION}"
    _npu_driver_create_prerm

    # ---- Step 4: Copy Source/Vendor Files ----
    _npu_driver_copy_files "${KERNEL_VERSION}"

    # ---- Step 5: Build Debian Package ----
    _npu_build_deb_package "COMPILE-NPU-DRIVER" "${PKG_NAME}_${PKG_VERSION}_${KERNEL_VERSION}_${PKG_ARCH}.deb"

    display_alert "COMPILE-NPU-DRIVER NPU driver build completed. Successfully built for ${PKG_ARCH} with kernel ${KERNEL_VERSION}" "info"
    return 0
}

# ------------------------------------------------------------------------------
# Main NPU Build Function
# ------------------------------------------------------------------------------

function amlbian_build_npu() {
    if ! is_unit_enabled "npu"; then
        echo "[SKIP] >> NPU disable in BUILD_UNITS"
        return
    fi
    amlbian_build_npu_common
    amlbian_build_adla
}

# ==============================================================================
# Main Execution
# ==============================================================================

display_alert "AMLbian NPU Build Script" "Starting NPU build" "info"
amlbian_build_npu
display_alert "AMLOGIC: NPU build completed" "Check output directory for packages" "info"