#!/usr/bin/env bash

VIOTUCLUSTER_PYTHON=${VIOTUCLUSTER_PYTHON:-python}

distribution_version() {
    "$VIOTUCLUSTER_PYTHON" -c 'from importlib.metadata import version; import sys; print(version(sys.argv[1]))' "$1"
}

# Function to check if a command is installed and retrieve its version
check_command() {
    if ! command -v "$1" >/dev/null 2>&1; then
        local script_dir
        local candidates=()
        script_dir="$(cd "$(dirname "$0")" && pwd)"
        case "$1" in
            "viralverify")
                candidates=(
                    "${script_dir}/../envs/viralverify/bin/viralverify"
                    "${script_dir}/../../viralverify/bin/viralverify"
                )
                ;;
            "vRhyme")
                candidates=(
                    "${script_dir}/../envs/vRhyme/bin/vRhyme"
                    "${script_dir}/../../vRhyme/bin/vRhyme"
                )
                ;;
        esac
        for candidate in "${candidates[@]}"; do
            if [ -x "$candidate" ]; then
                echo -e "[✅] $1 is available via environment: $candidate"
                return 0
            fi
        done
    fi

    command -v "$1" >/dev/null 2>&1
    if [ $? -eq 0 ]; then
        if [ "$1" = "checkm" ] && ! checkm -h >/dev/null 2>&1; then
            echo -e "[❌] checkm is installed but cannot start."
            return 1
        fi

        version=""
        case "$1" in
            "genomad")
                version="$(distribution_version "genomad")"
                ;;
            "checkm")
                version="$(distribution_version "checkm-genome")"
                ;;
            "checkv")
                version="$(checkv 2>&1 | grep -oP 'CheckV v\d+\.\d+\.\d+' | sed 's/CheckV v//')"
                ;;
            "dRep")
                version="$(dRep -h 2>&1 | grep -oP '\d+\.\d+\.\d+')"
                ;;
            "bwa")
                version="$(bwa 2>&1 | grep -oP '\d+\.\d+\.\d+')"
                ;;
            "virsorter")
                version="$(virsorter -h 2>&1 | grep -oP '\d+\.\d+\.\d+\.\d+')"
                ;;
            *)
                version=$("$1" --version 2>&1 | head -n 1 | grep -oP '\d+\.\d+\.\d+')
                ;;
        esac

        if [[ -n "$version" ]]; then
            echo -e "[✅] $1 is installed. Version: $version"
        else
            echo -e "[✅] $1 is installed, but version information could not be retrieved."
        fi
        return 0
    else
        echo -e "[❌] $1 is NOT installed."
        return 1
    fi
}

get_installed_version() {
    case "$1" in
        "genomad")
            distribution_version "genomad"
            ;;
        "checkm")
            distribution_version "checkm-genome"
            ;;
        "checkv")
            checkv 2>&1 | grep -oP 'CheckV v\d+\.\d+\.\d+' | sed 's/CheckV v//'
            ;;
        "dRep")
            dRep -h 2>&1 | grep -oP '\d+\.\d+\.\d+'
            ;;
        "bwa")
            bwa 2>&1 | grep -oP '\d+\.\d+\.\d+'
            ;;
        "virsorter")
            virsorter -h 2>&1 | grep -oP '\d+\.\d+\.\d+\.\d+'
            ;;
        *)
            "$1" --version 2>&1 | head -n 1 | grep -oP '\d+\.\d+\.\d+'
            ;;
    esac
}

# Check the installation of all required dependencies
echo "Checking dependencies..."

# List of dependencies to check
dependencies=(
    "fastp"
    "megahit"
    "spades.py"
    "virsorter"
    "viralverify"
    "genomad"
    #"pyhmmer"
    "checkv"
    "vRhyme"
    "dRep"
    #"DRAM"
    #"iPhop"
    "checkm"
    "bwa"
    "sambamba"
    "parallel"
    "makeblastdb"
    "blastn"
)

# Define expected versions
declare -A expected_versions
expected_versions=(
    [genomad]="1.8.0"
    [checkm]="1.2.2"
    [dRep]="3.5.0"
    [checkv]="1.0.3"
    [virsorter]="2.2.4.2"
)

# Initialize a variable to track missing dependencies
missing_dependencies=()

# Loop through each dependency and check if it's installed
for dep in "${dependencies[@]}"; do
    if check_command "$dep"; then
        if [[ -n "${expected_versions[$dep]}" ]]; then
            installed_version="$(get_installed_version "$dep")"

            if [[ "$installed_version" != "${expected_versions[$dep]}" ]]; then
                echo -e "[⚠️] $dep version mismatch: Expected ${expected_versions[$dep]}, but found $installed_version."
            fi
        fi
    else
        missing_dependencies+=("$dep")
    fi
done

# Final message based on whether all dependencies are installed
if [ ${#missing_dependencies[@]} -eq 0 ]; then
    echo -e "All dependencies are installed."
else
    echo -e "The following dependencies are missing: "
    for missing_dep in "${missing_dependencies[@]}"; do
        echo -e "  - ❌ $missing_dep"
    done
    exit 1
fi
