# update mar 12, 2026: Added ~/.codex to the share-state option
# to enable using codex directly instead of through opencode

llmsandbox() {
    local SHARE_STATE=0
    local OPENCODE_CONFIG_PATH=""
    local RUN_OPENCODE=0  # New variable to handle the -O flag

    # Reset getopts state (important in functions)
    OPTIND=1
    local opt
    while getopts ":SC:OE:" opt; do
        case "$opt" in
            S) SHARE_STATE=1 ;;
            C) OPENCODE_CONFIG_PATH="$OPTARG" ;;
            O) RUN_OPENCODE=1 ;;  # Set RUN_OPENCODE if -O is provided
            E) CONDA_ENVS_PATH="$OPTARG" ;;  # Set CONDA_ENVS_PATH if -E is provided
            :)  echo "Error: -$OPTARG requires a path."
                echo "Usage: llmsandbox2 -S [-C /path/to/opencode.json] [-O] [-E /path/to/conda/envs] /path/to/proj1 [/path/to/proj2 ...]"
                return 1
                ;;
            \?) echo "Error: Unknown option -$OPTARG"
                echo "Usage: llmsandbox2 -S [-C /path/to/opencode.json] [-O] [-E /path/to/conda/envs] /path/to/proj1 [/path/to/proj2 ...]"
                return 1
                ;;
        esac
    done
    shift $((OPTIND - 1))

    # Enforce: must provide -S or -C
    if [ "$SHARE_STATE" -eq 0 ] && [ -z "$OPENCODE_CONFIG_PATH" ]; then
        echo "Error: You must provide at least one of: -S (share host state) or -C <opencode.json>."
        echo "Usage: llmsandbox2 -S [-C /path/to/opencode.json] [-O] /path/to/proj1 [/path/to/proj2 ...]"
        echo "   or: llmsandbox2 -C /path/to/opencode.json /path/to/proj1 [/path/to/proj2 ...]"
        return 1
    fi

    # Create temporary session directory with full permissions
    local SESSION_DIR
    SESSION_DIR="$(mktemp -d /tmp/opencode_session-XXXXXX)" || return 1
    chmod 1777 "$SESSION_DIR" || return 1

    # Docker image name
    local IMAGE_NAME="andresfr/opencode-ubuntu-sandbox:conda_codex"

    # Get WORKDIR from the Docker image
    local WORKDIR
    WORKDIR="$(docker inspect --format '{{.Config.WorkingDir}}' "$IMAGE_NAME" 2>/dev/null)"
    [ -n "$WORKDIR" ] || WORKDIR="/workspace"

    # Build docker args (use -v as requested)
    local -a DOCKER_ARGS
    DOCKER_ARGS+=( --rm -it -p 1455:1456 --name opencode-ubuntu-sandbox )
    DOCKER_ARGS+=( -v "$SESSION_DIR:/tmp" )

    # expose host-side session dir path to container
    DOCKER_ARGS+=( -e "TMP_SESSION_DIR=$SESSION_DIR" )

    # Optional: config via -C
    if [ -n "$OPENCODE_CONFIG_PATH" ]; then
        OPENCODE_CONFIG_PATH="$(realpath -- "$OPENCODE_CONFIG_PATH")" || return 1
        if [ ! -f "$OPENCODE_CONFIG_PATH" ]; then
            echo "Error: config file not found: $OPENCODE_CONFIG_PATH"
            return 1
        fi
        DOCKER_ARGS+=( -v "$OPENCODE_CONFIG_PATH:$WORKDIR/opencode.json:ro" )
        DOCKER_ARGS+=( -e "OPENCODE_CONFIG=$WORKDIR/opencode.json" )
    fi

    # Optional: share host OpenCode state via -S
    if [ "$SHARE_STATE" -eq 1 ]; then
        local HOST_CACHE HOST_CONFIG HOST_DATA HOST_CODEX
        HOST_CACHE="${XDG_CACHE_HOME:-$HOME/.cache}/opencode"
        HOST_CONFIG="${XDG_CONFIG_HOME:-$HOME/.config}/opencode"
        HOST_DATA="${XDG_DATA_HOME:-$HOME/.local/share}/opencode"
        HOST_CODEX="$HOME/.codex"

        # Ensure host dirs exist before mounting
        mkdir -p "$HOST_CACHE" "$HOST_CONFIG" "$HOST_DATA" "$HOST_CODEX" || return 1

        DOCKER_ARGS+=( -v "$HOST_CACHE:/root/.cache/opencode" )
        DOCKER_ARGS+=( -v "$HOST_CONFIG:/root/.config/opencode" )
        DOCKER_ARGS+=( -v "$HOST_DATA:/root/.local/share/opencode" )
        DOCKER_ARGS+=( -v "$HOST_CODEX:/root/.codex" )
        DOCKER_ARGS+=( -e CODEX_HOME=/root/.codex )
    fi

    # Optional: Mount Conda environments from host via -E
    if [ -n "$CONDA_ENVS_PATH" ]; then
        if [ ! -d "$CONDA_ENVS_PATH" ]; then
            echo "Error: Conda environments directory not found: $CONDA_ENVS_PATH"
            return 1
        fi
        DOCKER_ARGS+=( -v "$CONDA_ENVS_PATH:/opt/miniconda/envs" )
    fi

    # If not sharing state, avoid autoupdate because exiting "too early"
    # while opencode is loading+updating may break it
    if [ "$SHARE_STATE" -eq 0 ]; then
        DOCKER_ARGS+=( -e OPENCODE_DISABLE_AUTOUPDATE=1 )
    fi

    # All remaining args are project dirs
    local dir abs name
    for dir in "$@"; do
        abs="$(realpath -- "$dir")" || return 1
        if [ ! -d "$abs" ]; then
            echo "Error: project dir not found: $abs"
            return 1
        fi
        name="$(basename "$abs")"
        DOCKER_ARGS+=( -v "$abs:$WORKDIR/$name" )
    done

    # Run the docker command
    if [ "$RUN_OPENCODE" -eq 1 ]; then
        docker run "${DOCKER_ARGS[@]}" "$IMAGE_NAME" opencode
    else
        docker run "${DOCKER_ARGS[@]}" "$IMAGE_NAME"
    fi
}