#!/bin/bash
set -e

GREEN='\033[0;32m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
NC='\033[0m'

log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

init() {
    log_info "Initializing BtToxin_Digger environment..."

    if ! BtToxin_Digger --version &>/dev/null; then
        log_error "BtToxin_Digger not found!"
        exit 1
    fi

    log_info "BtToxin_Digger $(BtToxin_Digger --version)"
    log_info "Python $(python --version)"
    log_info "Perl $(perl --version | head -2 | tail -1)"
}

update_db() {
    log_info "Updating BtToxin_Digger database..."
    if BtToxin_Digger --update-db; then
        log_info "Database updated successfully"
    else
        log_warn "Database update failed, using existing database"
    fi
}

run_analysis() {
    log_info "Starting toxin mining analysis..."
    log_info "Input: $INPUT_PATH"
    log_info "Type: $SEQUENCE_TYPE"
    log_info "Threads: $THREADS"

    BtToxin_Digger \
        --SeqPath "$INPUT_PATH" \
        --SequenceType "$SEQUENCE_TYPE" \
        --Scaf_suffix "$SCAF_SUFFIX" \
        --threads "$THREADS" \
        2>&1 | tee /workspace/logs/digger.log

    if [ $? -eq 0 ]; then
        log_info "Analysis completed successfully"
    else
        log_error "Analysis failed!"
        exit 1
    fi
}

main() {
    INPUT_PATH="${INPUT_PATH:-/workspace/input}"
    SEQUENCE_TYPE="${SEQUENCE_TYPE:-nucl}"
    SCAF_SUFFIX="${SCAF_SUFFIX:-.fna}"
    THREADS="${THREADS:-4}"
    UPDATE_DB="${UPDATE_DB:-false}"

    init

    if [ "$UPDATE_DB" = "true" ]; then
        update_db
    fi

    if [ $# -gt 0 ]; then
        exec BtToxin_Digger "$@"
    else
        run_analysis
    fi
}

main "$@"