#!/bin/bash

# This script sets up a comprehensive histopathology environment

# Set up colors for output
GREEN='\033[0;32m'
RED='\033[0;31m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Default tools directory is the current directory
DEFAULT_TOOLS_DIR="$(pwd)/histopath_tools"
TOOLS_DIR=""

echo -e "${YELLOW}=== HistoPath Environment Setup ===${NC}"
echo -e "${BLUE}This script will set up a comprehensive histopathology environment with various tools and packages.${NC}"

# Check if conda is installed
if ! command -v conda &> /dev/null && ! command -v micromamba &> /dev/null; then
    echo -e "${RED}Error: Conda is not installed or not in PATH.${NC}"
    echo "Please install Miniconda or Anaconda first."
    echo "Visit: https://docs.conda.io/en/latest/miniconda.html"
    exit 1
fi

# Function to handle errors
handle_error() {
    local exit_code=$1
    local error_message=$2
    local optional=${3:-false}

    if [ $exit_code -ne 0 ]; then
        echo -e "${RED}Error: $error_message${NC}"
        if [ "$optional" = true ]; then
            echo -e "${YELLOW}Continuing with setup as this component is optional.${NC}"
            return 0
        else
            if [ -z "$NON_INTERACTIVE" ]; then
                read -p "Continue with setup? (y/n) " -n 1 -r
                echo
                if [[ ! $REPLY =~ ^[Yy]$ ]]; then
                    echo -e "${RED}Setup aborted.${NC}"
                    exit 1
                fi
            else
                echo -e "${YELLOW}Non-interactive mode: continuing despite error.${NC}"
            fi
        fi
    fi
    return $exit_code
}

# Function to install a specific environment file
install_env_file() {
    local env_file=$1
    local description=$2
    local optional=${3:-false}

    echo -e "\n${BLUE}=== Installing $description ===${NC}"

    if [ "$optional" = true ]; then
        if [ -z "$NON_INTERACTIVE" ]; then
            read -p "Do you want to install $description? (y/n) " -n 1 -r
            echo
            if [[ ! $REPLY =~ ^[Yy]$ ]]; then
                echo -e "${YELLOW}Skipping $description installation.${NC}"
                return 0
            fi
        else
            echo -e "${YELLOW}Non-interactive mode: automatically installing $description.${NC}"
        fi
    fi

    echo -e "${YELLOW}Installing $description from $env_file...${NC}"
    conda env update -f $env_file
    handle_error $? "Failed to install $description." $optional

    if [ $? -eq 0 ]; then
        echo -e "${GREEN}Successfully installed $description!${NC}"
    fi
}

# Main installation process
main() {
    # Step 1: Create base conda environment
    echo -e "\n${YELLOW}Step 1: Creating base environment from environment.yml...${NC}"
    conda env create -n histopath -f environment.yml
    handle_error $? "Failed to create base conda environment."

    # Step 2: Activate the environment
    echo -e "\n${YELLOW}Step 2: Activating conda environment...${NC}"
    if command -v micromamba &> /dev/null; then
        eval "$("$MAMBA_EXE" shell hook --shell bash)"
        micromamba activate histopath
    else
        eval "$(conda shell.bash hook)"
        conda activate histopath
    fi
    handle_error $? "Failed to activate histopath environment."

    # Step 3: Install core histopathology tools
    echo -e "\n${YELLOW}Step 3: Installing core histopathology tools...${NC}"
    install_env_file "histo_env.yml" "core bioinformatics tools"
}

# Run the main installation process
main