#!/usr/bin/env bash
#
# emergetools - project maintenance CLI for EMerge
#
# Usage:
#   ./emergetools <command> <subcommand> [args...]
#
# Run `./emergetools help` for a list of commands.
#
# CLAUDE CODE GENERATED CLI
# ./emergetools clean numba-cache --dry-run   # preview
# ./emergetools clean numba-cache             # actually delete .nbi/.nbc
# ./emergetools clean pycache                 # nukes __pycache__ entirely (also clears numba cache, since they share dirs)
# ./emergetools help
set -euo pipefail

# Resolve repo root as the directory this script lives in, so it behaves
# the same regardless of the caller's current working directory.
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
SRC_DIR="${SCRIPT_DIR}/src"

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

log()  { echo "[emergetools] $*"; }
err()  { echo "[emergetools] ERROR: $*" >&2; }

require_src_dir() {
    if [[ ! -d "$SRC_DIR" ]]; then
        err "Expected source directory not found: $SRC_DIR"
        err "(edit SRC_DIR at the top of this script if your layout differs)"
        exit 1
    fi
}

detect_platform() {
    local os
    local arch
    os="$(uname -s)"
    arch="$(uname -m)"
    case "$os" in
        Darwin)
            if [[ "$arch" == "x86_64" ]] && [[ "$(sysctl -n sysctl.proc_translated 2>/dev/null)" == "1" ]]; then
                arch="arm64"
            fi

            if [[ "$arch" == "arm64" ]]; then
                echo "arm-macos"
            else
                echo "x86-macos"
            fi
            ;;
        Linux)
            echo "linux"
            ;;
        *)
            echo "unknown"
            ;;
    esac
}

has_accelerate() {

    for py in python3 python; do
        if command -v "$py" &>/dev/null; then
            if "$py" -c "import emerge_aasds" $>/dev/null; then
                return 0
            fi
        fi
    done
    return 1
}
set_threads_armmacos() {
    local nthreads="$1"
    echo "Setting threads to $nthreads"
    echo " - Checking if Accellerate solver is installed ..."
    accelerate_installed="$(has_accelerate)"
    if [[ accelerate_installed==0 ]]; then
        echo " - Accellerate solver found!"
        echo " - Setting Veclib thread count to ${nthreads}"
        export VECLIB_NUM_THREADS="${nthreads}"
        export VECLIB_MAXIMUM_THREADS="${nthreads}"
    else
        echo " - Accellerate solver not found :("
        echo " - Consider installing using:"
        echo " > pip install git+https://github.com/FennisRobert/emerge-aasds"
    fi
}

set_threads_linux() {
    local nthreads="$1"
    echo "Setting threads to $nthreads"

    echo " - Setting PARDISO/MKL thread count to ${nthreads}"

    export MKL_NUM_THREADS="${nthreads}"
    export OMP_NUM_THREADS="${nthreads}"
}

# ---------------------------------------------------------------------------
# clean numba-cache
# ---------------------------------------------------------------------------

cmd_clean_numba_cache() {
    require_src_dir

    local dry_run=false
    for arg in "$@"; do
        case "$arg" in
            --dry-run|-n) dry_run=true ;;
            *) err "Unknown option for 'clean numba-cache': $arg"; exit 1 ;;
        esac
    done

    log "Scanning for Numba cache artifacts under: $SRC_DIR"

    local files
    files="$(find "$SRC_DIR" \( -name "*.nbi" -o -name "*.nbc" \) 2>/dev/null || true)"

    if [[ -z "$files" ]]; then
        log "No cache files found. Nothing to do."
        return 0
    fi

    local count
    count="$(echo "$files" | wc -l | tr -d ' ')"

    if $dry_run; then
        log "Would remove $count file(s):"
        echo "$files" | sed 's/^/  /'
        return 0
    fi

    echo "$files" | while IFS= read -r f; do
        rm -f -- "$f"
    done
    log "Removed $count Numba cache file(s). Next import will trigger a full recompile."
}

# ---------------------------------------------------------------------------
# clean pycache  (bonus: also clears regular .pyc / __pycache__ dirs)
# ---------------------------------------------------------------------------

cmd_clean_pycache() {
    require_src_dir

    local dry_run=false
    for arg in "$@"; do
        case "$arg" in
            --dry-run|-n) dry_run=true ;;
            *) err "Unknown option for 'clean pycache': $arg"; exit 1 ;;
        esac
    done

    log "Scanning for __pycache__ directories under: $SRC_DIR"

    local dirs
    dirs="$(find "$SRC_DIR" -type d -name "__pycache__" 2>/dev/null || true)"

    if [[ -z "$dirs" ]]; then
        log "No __pycache__ directories found. Nothing to do."
        return 0
    fi

    local count
    count="$(echo "$dirs" | wc -l | tr -d ' ')"

    if $dry_run; then
        log "Would remove $count directory(ies):"
        echo "$dirs" | sed 's/^/  /'
        return 0
    fi

    echo "$dirs" | while IFS= read -r d; do
        rm -rf -- "$d"
    done
    log "Removed $count __pycache__ director(y/ies) (this also clears .pyc and Numba cache files)."
}

# ---------------------------------------------------------------------------
# Dispatch
# ---------------------------------------------------------------------------

print_help() {
    cat <<EOF
emergetools - project maintenance CLI for EMerge

Usage:
  ./emergetools clean numba-cache [--dry-run]
      Remove all Numba *.nbi/*.nbc cache files under src/, forcing a full
      recompile of every @njit(cache=True) function on next import.

  ./emergetools clean pycache [--dry-run]
      Remove all __pycache__ directories under src/ (clears .pyc files too,
      as well as Numba cache files since they share the same directories).

  ./emergetools help
      Show this message.

Options:
  --dry-run, -n   List what would be removed without deleting anything.
EOF
}

main() {
    if [[ $# -eq 0 ]]; then
        print_help
        exit 0
    fi

    local command="$1"; shift || true

    case "$command" in
        clean)
            if [[ $# -eq 0 ]]; then
                err "Missing subcommand for 'clean'. Try: numba-cache, pycache"
                exit 1
            fi
            local subcommand="$1"; shift || true
            case "$subcommand" in
                numba-cache|numba) cmd_clean_numba_cache "$@" ;;
                pycache)           cmd_clean_pycache "$@" ;;
                *)
                    err "Unknown 'clean' subcommand: $subcommand"
                    err "Try: numba-cache, pycache"
                    exit 1
                    ;;
            esac
            ;;
        help|-h|--help)
            print_help
            ;;
        threads)
            if [[ $# -eq 0 ]]; then
                err "Missing thread count specifier. Use an integer value above 0."
                exit 1
            fi

            nthreads="$1"
            my_os="$(detect_platform)"
            case "$my_os" in
                arm-macos)
                    set_threads_armmacos "$nthreads"
                    ;;
                x86-macos)
                    set_threads_armmacos "$nthreads"
                    ;;
                linux)
                    set_threads_linux "$nthreads"
                    ;;
                *)
                    echo "Unkonwn architecture :("
                    ;;
            esac
            ;;
        *)
            err "Unknown command: $command"
            print_help
            exit 1
            ;;

    esac
}

main "$@"
