diff --git a/plugin.js b/plugin.js
index 802f768..d63b80c 100644
--- a/plugin.js
+++ b/plugin.js
@@ -1,15 +1,24 @@
 CKEDITOR.plugins.add( 'codeTag', {
   icons: 'code',
   init: function( editor ) {
-    editor.addCommand( 'wrapCode', {
-      exec: function( editor ) {
-        editor.insertHtml( '<code>' + editor.getSelection().getSelectedText() + '</code>' );
-      }
-    });
-    editor.ui.addButton( 'Code', {
-      label: 'Wrap code',
-      command: 'wrapCode',
-      toolbar: 'insert'
-    });
+
+    var style = new CKEDITOR.style( { element: 'code' } );
+
+    // Listen to contextual style activation.
+    editor.attachStyleStateChange( style, function (state) {
+      !editor.readOnly && editor.getCommand( 'wrapCode').setState(state);
+    } );
+
+    // Create the command.
+    editor.addCommand( 'wrapCode', new CKEDITOR.styleCommand( style ) );
+
+    // Register the button, if the button plugin is enabled.
+    if ( editor.ui.addButton ) {
+      editor.ui.addButton( 'Code', {
+        label: 'Code',
+        command: 'wrapCode',
+        toolbar: 'insert'
+      } );
+    }
   }
-});
\ No newline at end of file
+});