#!/bin/bash

# model-checker.sh - Check and manage ONNX model files for Kylin AI
#
# This script verifies the presence and integrity of ONNX model files
# required for Kylin AI functionality. It can also copy these files
# from a reference location when needed.
#
# Usage:
#   ./model-checker.sh          # Check model files
#   ./model-checker.sh -c       # Copy model files from reference location

set -e
set -o pipefail

# Configuration
SOURCE_DIR="/opt/appdata/kylin-ai/triton_cpu"

# Model files and their expected MD5 checksums
declare -A MODEL_FILES_MD5=(
    ["model_repo_paraformer_offline/asr_encoder_nonstreaming/1/model.onnx"]="88f0dd75be4815345833deedd48fffc3"
    ["model_repo_paraformer_online/asr_encoder_streaming/1/model.onnx"]="f2f1559529a29696bb004873774eb23c"
    ["model_repo_paraformer_online/asr_lfr-cmvn-pe_streaming/1/lfr_cmvn_pe.onnx"]="4a71760be89fcc0da10406dbf9117800"
    ["model_repo_paraformer_online/asr_decoder_streaming/1/decoder.onnx"]="ee0271138e4b5f58c4f1856d58d28e7e"
    ["model_repo_punc_dynamic_batching_small/punc_ct-transformer_nonstreaming/1/model.onnx"]="3cbe569249b95a43d79b9d02b4e5a6bb"
    ["model_repo_vad_dynamic_batching/asr_fsmn-vad_streaming/1/model.onnx"]="c0224e566d4330bb0dfb031d064d3abe"
)

# Color codes for output
readonly RED='\033[0;31m'
readonly GREEN='\033[0;32m'
readonly YELLOW='\033[1;33m'
readonly NC='\033[0m' # No Color

# Print colored output
print_message() {
    local color="$1"
    local level="$2"
    local message="$3"
    
    echo -e "${color}${level}${message}${NC}" >&2
}

error() {
    print_message "${RED}" "ERROR: " "$1"
}

warning() {
    print_message "${YELLOW}" "WARNING: " "$1"
}

success() {
    print_message "${GREEN}" "" "$1"
}

info() {
    echo "$1"
}

# Copy model files from reference location
copy_model_files() {
    local file_count=0
    local success_count=0
    local errors=()
    
    info "Copying ONNX model files from ${SOURCE_DIR}..."
    
    for file in "${!MODEL_FILES_MD5[@]}"; do
        local source_file="${SOURCE_DIR}/${file}"
        local dest_file="${file}"
        local dest_dir=$(dirname "${dest_file}")
        
        ((file_count++))
        
        # Check source file existence
        if [[ ! -f "${source_file}" ]]; then
            errors+=("${source_file}")
            error "Source file not found: ${source_file}"
            continue
        fi
        
        # Create destination directory
        if ! mkdir -p "${dest_dir}" 2>/dev/null; then
            errors+=("${dest_file}")
            error "Failed to create directory: ${dest_dir}"
            continue
        fi
        
        # Copy file
        if cp "${source_file}" "${dest_file}" 2>/dev/null; then
            info "Copied: ${file}"
            ((success_count++))
        else
            errors+=("${file}")
            error "Failed to copy ${file}"
        fi
    done
    
    # Report results
    if (( ${#errors[@]} > 0 )); then
        echo
        error "Failed to copy ${#errors[@]} file(s):"
        for err in "${errors[@]}"; do
            error "  - ${err}"
        done
        return 1
    else
        echo
        success "Successfully copied ${success_count}/${file_count} model files."
        return 0
    fi
}

# Check model files integrity
check_model_files() {
    local missing_files=()
    local mismatch_files=()
    local file_count=0
    local verified_count=0
    
    info "Checking for required ONNX model files..."
    
    for file in "${!MODEL_FILES_MD5[@]}"; do
        local expected_md5="${MODEL_FILES_MD5[${file}]}"
        local actual_md5
        
        ((file_count++))
        
        # Check file existence
        if [[ ! -f "${file}" ]]; then
            missing_files+=("${file}")
            error "Missing required file: ${file}"
            continue
        fi
        
        # Calculate MD5 checksum
        actual_md5=$(md5sum "${file}" 2>/dev/null | cut -d ' ' -f 1)
        if [[ -z "${actual_md5}" ]]; then
            missing_files+=("${file}")
            error "Failed to calculate MD5 for file: ${file}"
            continue
        fi
        
        # Verify MD5
        if [[ "${actual_md5}" != "${expected_md5}" ]]; then
            mismatch_files+=("${file}")
            error "MD5 mismatch for file: ${file}"
            info "  Local file MD5:  ${actual_md5}"
            info "  Expected MD5:    ${expected_md5}"
        else
            success "Found: ${file} (MD5 verified)"
            local size=$(du -h "${file}" 2>/dev/null | awk '{print $1}')
            [[ -n "${size}" ]] && info "  Size: ${size}"
            ((verified_count++))
        fi
    done
    
    # Report results
    if (( ${#missing_files[@]} > 0 )) || (( ${#mismatch_files[@]} > 0 )); then
        echo
        
        if (( ${#missing_files[@]} > 0 )); then
            warning "${#missing_files[@]} required model file(s) missing:"
            for file in "${missing_files[@]}"; do
                warning "  - ${file}"
            done
        fi
        
        if (( ${#mismatch_files[@]} > 0 )); then
            warning "${#mismatch_files[@]} model file(s) with MD5 mismatch:"
            for file in "${mismatch_files[@]}"; do
                warning "  - ${file}"
            done
        fi
        
        echo
        error "Please ensure all ONNX model files are present and have correct MD5 before building the package."
        return 1
    else
        echo
        success "All required ONNX model files are present and MD5 verified (${verified_count}/${file_count})."
        info "Ready to build package."
        return 0
    fi
}

# Main function
main() {
    case "${1:-}" in
        -c|--copy)
            copy_model_files && exit 0 || exit 1
            ;;
        -h|--help)
            info "Usage: $0 [OPTION]"
            info "Check and manage ONNX model files for Kylin AI."
            info ""
            info "Options:"
            info "  -c, --copy    Copy model files from reference location"
            info "  -h, --help    Display this help and exit"
            info ""
            info "Without options, checks the presence and integrity of model files."
            ;;
        *)
            # Default behavior: check files
            check_model_files && exit 0 || exit 1
            ;;
    esac
}

# Run main function with all arguments
main "$@"