Updated Return to Selection Screen

-- collision with bad

local flagged_sprites = get_sprites_with_flag(1)

for sprite in all(flagged_sprites) do

    if collide_bbox(snake[1], sprite, 8, 8) then

        -- only restart the level if snake is moving

        if snake_dir.x ~= 0 or snake_dir.y ~= 0 then

            if current_level == "air" then

                reset_snake_position()

            elseif game_state == "final" then

                restart_final_level()

            else

                restart_level()

                break -- exit the loop early as we've already found a collision

            end

        end

    end

end

CHANGED TO

-- collision with bad

local flagged_sprites = get_sprites_with_flag(1)

for sprite in all(flagged_sprites) do

    if collide_bbox(snake[1], sprite, 8, 8) then

        if game_state == "selection" then

            -- Logic to move snake back to center without resetting anything

            reset_snake_position()  -- Assuming this function moves the snake to the center

        elseif snake_dir.x ~= 0 or snake_dir.y ~= 0 then

            if current_level == "air" then

                reset_snake_position()

            elseif game_state == "final" then

                restart_final_level()

            else

                restart_level()

                break -- exit the loop early as we've already found a collision

            end

        end

    end

end