#!/usr/bin/env bash

set -euo pipefail

# This script builds a tarball of all the dependencies for the project
# It is meant to be run on a Linux machine without internet access

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)"
PYTHON_BIN="/usr/bin/python3.9"
REQUIREMENTS_FILE="requirements.txt"
DEPENDENCIES_DIR="${SCRIPT_DIR}/dependencies"
DEPENDENCIES_TARBALL="${SCRIPT_DIR}/dependencies.tar.gz"
LOG_FILE="${SCRIPT_DIR}/setup-deps.log"

# Check Python version
check_python_version() {
    if ! command -v "${PYTHON_BIN}" >/dev/null 2>&1; then
        echo "Error: Python is not installed at ${PYTHON_BIN}"
        exit 1
    fi
    
    local version
    version=$("${PYTHON_BIN}" -c 'import sys; print(".".join(map(str, sys.version_info[:2])))')
    if [[ "${version}" != "3.9" ]]; then
        echo "Error: Required Python version is 3.9, but found ${version}"
        exit 1
    fi
}

# Check pip version and upgrade if needed
check_pip_version() {
    "${PYTHON_BIN}" -m pip install --upgrade pip >/dev/null 2>&1
}

# Setup logging
log() {
    local level="$1"
    shift
    echo "[$(date +'%Y-%m-%d %H:%M:%S')] [${level}] $*" | tee -a "${LOG_FILE}"
}

# Function to resolve absolute paths
resolve_path() {
    local path="$1"
    if [[ -f "$path" ]]; then
        if command -v realpath >/dev/null 2>&1; then
            realpath "$path"
        else
            readlink -f "$path"
        fi
    else
        echo ""
    fi
}

# Function to change directories upwards
updir() {
    local levels="$1"
    local i
    for ((i=0; i<levels; i++)); do
        cd .. || exit
    done
}

# Function to clean up dependencies
cleanup() {
    echo "Cleaning up dependencies directory and tarball..."
    rm -rf "${DEPENDENCIES_DIR}"
    rm -f "${DEPENDENCIES_TARBALL}"
}

# Function to find requirements file
find_requirements_file() {
    local current_dir parent_dir
    current_dir="$(pwd)"
    parent_dir="$(dirname "${current_dir}")"
    
    requirements_file=$(find "${current_dir}" "${parent_dir}" "${SCRIPT_DIR}" -maxdepth 1 -name "${REQUIREMENTS_FILE}" -type f)
    echo "${requirements_file}"
}

# Improved package_dependencies function
package_dependencies() {
    local abs_requirements_file
    abs_requirements_file="$(find_requirements_file)"

    if [[ -z "${abs_requirements_file}" ]]; then
        log "ERROR" "Requirements file not found in current, parent, or script directory"
        exit 1
    fi

    log "INFO" "Using requirements file: ${abs_requirements_file}"
    
    mkdir -p "${DEPENDENCIES_DIR}"
    
    log "INFO" "Downloading dependencies..."
    if ! "${PYTHON_BIN}" -m pip download -r "${abs_requirements_file}" \
        --platform manylinux1_x86_64 \
        --only-binary=:all: \
        -d "${DEPENDENCIES_DIR}" 2>&1 | tee -a "${LOG_FILE}"; then
        log "ERROR" "Failed to download dependencies"
        cleanup
        exit 1
    fi
    
    log "INFO" "Creating tarball..."
    tar czf "${DEPENDENCIES_TARBALL}" -C "$(dirname "${DEPENDENCIES_DIR}")" "$(basename "${DEPENDENCIES_DIR}")"
    log "INFO" "Dependencies packaged successfully"
}

# Improved install_dependencies function
install_dependencies() {
    if [[ ! -d "${DEPENDENCIES_DIR}" ]]; then
        if [[ -f "${DEPENDENCIES_TARBALL}" ]]; then
            log "INFO" "Extracting dependencies from tarball..."
            tar xzf "${DEPENDENCIES_TARBALL}"
        else
            log "ERROR" "Neither dependencies directory nor tarball found"
            exit 1
        fi
    fi

    local abs_requirements_file
    abs_requirements_file="$(find_requirements_file)"

    if [[ -z "${abs_requirements_file}" ]]; then
        log "ERROR" "Requirements file not found"
        exit 1
    fi

    log "INFO" "Installing dependencies..."
    "${PYTHON_BIN}" -m pip install -r "${abs_requirements_file}" \
        --no-index \
        --find-links "${DEPENDENCIES_DIR}" \
        --no-deps
}

# Function to uninstall dependencies
uninstall_dependencies() {
    local abs_requirements_file
    abs_requirements_file="$(resolve_path "${REQUIREMENTS_FILE}")"

    if [[ -z "${abs_requirements_file}" ]]; then
        log "ERROR" "Requirements file not found"
        exit 1
    fi

    "${PYTHON_BIN}" -m pip uninstall -y -r "${abs_requirements_file}"
}

# Function to print help message
print_help() {
    echo "Usage: $0 [options]"
    echo ""
    echo "Options:"
    echo "  -h, --help                     Print this help message"
    echo "  -p, --package_dependencies     Package dependencies"
    echo "  -r, --requirements FILE        Full path to requirements.txt file"
    echo "  -c, --cleanup                  Cleanup dependencies"
    echo "  -i, --install_dependencies     Install dependencies"
    echo "  -u, --uninstall_dependencies   Uninstall dependencies"
}

# Function to parse arguments
args() {
    if [[ $# -eq 0 ]]; then
        print_help
        exit 0
    fi

    while [[ $# -gt 0 ]]; do
        key="$1"
        case $key in
            -h|--help)
                print_help
                exit 0
                ;;
            -r|--requirements)
                if [[ -n "${2-}" ]]; then
                    REQUIREMENTS_FILE="$2"
                    REQUIREMENTS_FILE="$(resolve_path "${REQUIREMENTS_FILE}")"
                    shift
                else
                    echo "Error: --requirements requires a file path argument."
                    exit 1
                fi
                shift
                ;;
            -p|--package_dependencies)
                package_dependencies
                exit 0
                ;;
            -c|--cleanup)
                cleanup
                exit 0
                ;;
            -i|--install_dependencies)
                install_dependencies
                exit 0
                ;;
            -u|--uninstall_dependencies)
                uninstall_dependencies
                exit 0
                ;;
            *)
                echo "Unknown argument: $key"
                print_help
                exit 1
                ;;
        esac
    done
}

# Main execution
main() {
    check_python_version
    check_pip_version
    args "$@"
}

main "$@"