#!/bin/bash

SESSION_NAME="$1"
PYTHON_VERSION_INPUT="${2:-3.10}"  # Default to 3.8 if not specified

# Normalize the Python version input
if [[ "$PYTHON_VERSION_INPUT" =~ ^python[0-9.]+$ ]]; then
    PYTHON_VERSION="${PYTHON_VERSION_INPUT//python/}"
elif [[ "$PYTHON_VERSION_INPUT" =~ ^py[0-9.]+$ ]]; then
    PYTHON_VERSION="${PYTHON_VERSION_INPUT//py/}"
else
    PYTHON_VERSION="$PYTHON_VERSION_INPUT"
fi

# Format for Conda
MAMBA_PYTHON_VERSION="python=$PYTHON_VERSION"

# Check if Conda environment exists
if ! mamba env list | grep -q "^$SESSION_NAME\s"; then
    echo "Creating new Mamba environment: $SESSION_NAME with $MAMBA_PYTHON_VERSION"
    mamba create --name "$SESSION_NAME" "$MAMBA_PYTHON_VERSION" --yes
fi

# Find Conda env directory
CONDA_ENV_DIR=$(mamba  env list | grep "^$SESSION_NAME" | awk '{print $2}')

# Handle tmux session
if ! tmux has-session -t "$SESSION_NAME" 2>/dev/null; then
    echo "Creating new tmux session: $SESSION_NAME"
    tmux new-session -d -s "$SESSION_NAME"
    sleep 2
fi

# Attach to tmux session and update PATH before activating Conda environment
sleep 1
tmux send-keys -t "$SESSION_NAME" "export PATH=\"$MAMBA_ENV_DIR/bin:\$PATH\"" C-m
tmux send-keys -t "$SESSION_NAME" "source ~/.zshrc" C-m
tmux send-keys -t "$SESSION_NAME" "mamba activate $SESSION_NAME" C-m
tmux attach -t "$SESSION_NAME"