#!/bin/bash
torch_version='2.5'
require_dtk_version='25042'






# Display help information
show_help() {
    echo "Usage: $0 [options]"
    echo
    echo "Options:"
    echo "  -c, --compile   Control compilation component, set USE_FASTPT_CUDA and load CUDA environment"
    echo "  -e, --execute   Control usage component, initialize fastpt.cuda and reload .bashrc"
    echo "  -t, --test      Control transcoding compilation component and initialize compilation environment"
    echo "  -h, --help      Display this help information"
    echo
    echo "Note: -t cannot be set with -c or -e simultaneously."
    return 0
}

check_dtk() {
    if [ -z "$ROCM_PATH" ]; then
        #echo "WARNING: ROCM_PATH environment variable is not set, check and use dtk path"
        if [ -d "/opt/dtk" ]; then
            echo "dtk path is /opt/dtk"
            version_file="/opt/dtk/.info/rocm_version"
            current_dtk_version=$(cat "$version_file" | tr -d '.')
            echo "current_dtk_version: $current_dtk_version"
            if [[ "$current_dtk_version" == "$require_dtk_version" ]]; then
            #    return 0  # Return true
                echo "SUCCESS: The current dtk version is consistent with the required dtk version."
            else
            #    return 1  # Return false
                echo "WARNING: The current dtk version $current_dtk_version, is inconsistent with the require dtk version $require_dtk_version"
            fi
        else
            echo "Error: ROCM_PATH environment variable is not set, and dtk path is not valid"
            return 1 # Return false
        fi
    else
        version_file="$ROCM_PATH/.info/rocm_version"
        current_dtk_version=$(cat "$version_file" | tr -d '.')
        echo "current_dtk_version: $current_dtk_version"
        if [[ "$current_dtk_version" == "$require_dtk_version" ]]; then
            #    return 0  # Return true
            echo "WARNING: The current dtk version is consistent with the required dtk version."
        else
            #    return 1  # Return false
            echo "WARNING: The current dtk version $current_dtk_version, is inconsistent with the require dtk version $require_dtk_version"
        fi
    fi
}

compile=false
execute=false
test_env=false
error_occurred=false
TORCH_PATH=$(python3 -c "import torch; print(torch.__path__[0])")

report_error() {
    echo "ERROR: $1"
    error_occurred=true
}

if [[ $# -eq 0 ]]; then
    report_error "Parameter must be specified"
    show_help
    return 1
else
    while [[ $# -gt 0 ]]; do
        case "${1,,}" in  # Convert parameter to lowercase
            -c|--compile|-C|--COMPILE)
                compile=true
                shift
                ;;
            -e|--execute|-E|--EXECUTE)
                execute=true
                shift
                ;;
            -t|--test|-T|--TEST)
                test_env=true
                shift
                ;;
            -h|--help|-H|--HELP)
                show_help
                shift
                break
                ;;
            *)
                report_error "Unknown parameter '$1'"
                echo "Use 'source /usr/local/bin/fastpt -h' to view help information"
                return 1
                ;;
        esac
    done
fi

# Check if both -t and -c parameters are set
if [[ "$compile" = true && "$test_env" = true ]]; then
    report_error "-t and -c parameters cannot be set simultaneously"
    show_help
    return 1
fi

# Check if both -e and -t parameters are set
if [[ "$execute" = true && "$test_env" = true ]]; then
    report_error "-e and -t parameters cannot be set simultaneously"
    show_help
    return 1
fi

check_torch_version() {
    
    # Get current PyTorch version and extract the first two digits
    current_torch_version=$(python -c "import torch; print(torch.__version__)" | cut -d '.' -f 1,2)
    #current_torch_hcu_version=$(python -c "import torch; print(torch.__hcu_version__)")

    # Perform matching comparison
    if [[ "$current_torch_version" == "$torch_version" ]]; then
        return 0  # Return true
    else
        return 1  # Return false
    fi
}

if [[ "$compile" = true ]]; then
    if check_dtk; then
        # Check PyTorch version
        if check_torch_version; then
            # Torch version matches, execute subsequent code
            if [[ -n "$USE_FASTPT" && "$USE_FASTPT" -eq 1 ]]; then
                echo "WARNING: Check USE_FASTPT environment variable is set, currently using USE_FASTPT_CUDA for compilation"
            fi
          
            cuda_env_path="$ROCM_PATH/cuda/env.sh"

            if [[ ! -f "$cuda_env_path" ]]; then
                report_error "Cannot find cuda env file $cuda_env_path"
            else
                source "$cuda_env_path"
                fastpt.cuda.init
                export USE_FASTPT_CUDA=1
                export LD_LIBRARY_PATH=$TORCH_PATH/lib:$LD_LIBRARY_PATH
                echo "USE_FASTPT_CUDA is set, and CUDA environment is loaded"
            fi
        else
            # Torch version does not match, output error message
            report_error "Torch version is not supported, current torch version is $current_torch_version.x, required torch version is $torch_version.x"
        fi
    else
        report_error "ROCM_PATH is not set and the /opt/dtk path  does not exist"
    fi
fi

if [[ "$execute" = true ]]; then
    unset USE_FASTPT_CUDA
    if check_dtk; then
        # Check PyTorch version
        if check_torch_version; then
            # Torch version matches, execute subsequent code
            fastpt.cuda.exe
            echo "current_dtk_version: $current_dtk_version"
            export LD_LIBRARY_PATH=$TORCH_PATH/lib:$LD_LIBRARY_PATH
            if [[ "$current_dtk_version" == "2504" ]]; then
                export LD_LIBRARY_PATH=$ROCM_PATH/cuda/lib64:$LD_LIBRARY_PATH
            elif [[ "$current_dtk_version" == "25041" ]]; then
                export LD_LIBRARY_PATH=$ROCM_PATH/cuda/cuda-11/lib64:$LD_LIBRARY_PATH
            else
                export LD_LIBRARY_PATH=$ROCM_PATH/cuda/cuda-12/lib64:$LD_LIBRARY_PATH
                echo "Default LD_LIBRARY_PATH: $ROCM_PATH/cuda/cuda-12/lib64/"
                return 1
            fi

            echo "fastpt.cuda has been initialized"
        else
            # Torch version does not match, output error message
            report_error "Torch version is not supported, current torch version is $current_torch_version.x, required torch version is $torch_version.x"
        fi
    else
        report_error "ROCM_PATH is not set and the /opt/dtk path  does not exist"
    fi
fi

# Handle test parameter
if [[ "$test_env" = true ]]; then
    if check_dtk; then
        # Check PyTorch version
        if check_torch_version; then
            # Torch version matches, continue checking other conditions
            if [[ (-n "$USE_FASTPT_CUDA" && "$USE_FASTPT_CUDA" -eq 1) || -n "$CUDA_PATH" ]]; then
                report_error "GPUFusion environment is detected, currently using transcoding compilation. GPUFusion should not be set. Please try opening a new terminal."
            else
                # All conditions are met, execute subsequent code
                export USE_FASTPT=1
                export USE_FASTPT_CPPEXTENSION=1
                echo "Transcoding environment settings have been initialized"
            fi
        else
            # Torch version does not match, output error message
            report_error "Torch version is not supported, current torch version is $current_torch_version.x, required torch version is $torch_version.x"
        fi
    else
        report_error "ROCM_PATH is not set and the /opt/dtk path  does not exist"
    fi
fi
