#!/usr/bin/env bash
# add_kernel.sh — create or use Python env, ensure ipykernel, register Jupyter kernel
#
# Usage:
#   ./add_kernel.sh ENV_NAME [--create] [--verbose] [pkg1 pkg2 ...]
#
# Examples:
#   ./add_kernel.sh ds_env --create python=3.11 pandas
#   ./add_kernel.sh ds_env --verbose

set -euo pipefail

print_help() {
	cat <<EOF
Usage:
  $0 ENV_NAME [--create] [--verbose] [pkg1 pkg2 ...]

Examples:
  $0 ds_env --create python=3.11 pandas
  $0 ds_env --verbose
EOF
	exit 0
}

if [[ $# -eq 0 ]] || [[ "$1" == "--help" ]] || [[ "$1" == "-h" ]]; then
	print_help
fi

ENV_NAME="$1"
shift

CREATE=0
VERBOSE=0
PKGS=()

while (("$#")); do
	case "$1" in
	--create) CREATE=1 ;;
	--verbose) VERBOSE=1 ;;
	--help | -h) print_help ;;
	*) PKGS+=("$1") ;;
	esac
	shift
done

log() {
	if ((VERBOSE)); then
		echo "$@"
	fi
}

detect_package_manager() {
	if conda info --envs 2>/dev/null | grep -qE "^$ENV_NAME\s"; then
		echo "conda"
	elif [ -f "environment.yml" ] || [ -f "conda.yaml" ]; then
		echo "conda"
	elif [ -f "requirements.txt" ] || [ -f "pyproject.toml" ]; then
		echo "pip"
	else
		if [ -n "${VIRTUAL_ENV:-}" ]; then
			echo "pip"
		else
			echo "conda"
		fi
	fi
}

PKG_MANAGER=$(detect_package_manager)
log "Detected package manager: $PKG_MANAGER"

if [[ "$PKG_MANAGER" == "conda" ]]; then
	# Initialize conda inside this script's shell so 'conda activate' works
	CONDA_BASE=$(conda info --base)
	# shellcheck source=/dev/null
	. "$CONDA_BASE/etc/profile.d/conda.sh"

	if ((CREATE)); then
		log "Creating conda env '$ENV_NAME' with packages: ${PKGS[*]:-python=3.11 ipykernel}"
		conda create -y -n "$ENV_NAME" "${PKGS[@]:-python=3.11}" ipykernel
	else
		if ! conda info --envs 2>/dev/null | grep -qE "^$ENV_NAME\s"; then
			echo "Error: Conda environment '$ENV_NAME' does not exist. Use --create to create it."
			exit 1
		fi
		log "Using existing conda env '$ENV_NAME'"
	fi

	log "Activating conda env '$ENV_NAME'"
	conda activate "$ENV_NAME"
else
	if ((CREATE)); then
		log "Creating python venv '$ENV_NAME'"
		python3 -m venv "$ENV_NAME"
		source "$ENV_NAME/bin/activate"
		log "Installing packages: ${PKGS[*]:-ipykernel}"
		pip install --upgrade pip
		pip install "${PKGS[@]:-ipykernel}"
	else
		if [[ ! -f "$ENV_NAME/bin/activate" ]]; then
			echo "Error: Python venv '$ENV_NAME' does not exist. Use --create to create it."
			exit 1
		fi
		log "Activating existing python venv '$ENV_NAME'"
		source "$ENV_NAME/bin/activate"
	fi
fi

# Ensure ipykernel installed
if ! python -c "import ipykernel" &>/dev/null; then
	echo "ipykernel not found, installing..."
	if [[ "$PKG_MANAGER" == "conda" ]]; then
		log "Installing ipykernel with conda..."
		conda install -y ipykernel
	else
		log "Installing ipykernel with pip..."
		pip install ipykernel
	fi
fi

echo "Using python executable: $(which python)"

echo "Registering kernel '$ENV_NAME' with Jupyter"
python -m ipykernel install --user --name "$ENV_NAME" --display-name "Python ($ENV_NAME)"

echo "✅ Kernel registered! You can now select 'Python ($ENV_NAME)' in Jupyter."