Highlighter uses last used style

Source/EPUBCore/FRHighlight.swift

@@ -141,7 +141,7 @@ class FRHighlight: NSObject {
         if let removedId = currentPage.webView.js("removeHighlightById('\(highlightId)')") {
             return removedId
         } else {
-            print("Error removing Higlight")
+            print("Error removing Higlight from page")
             return nil
         }
     }

Source/FolioReaderContainer.swift

@@ -73,7 +73,8 @@ class FolioReaderContainer: UIViewController, FolioReaderSidePanelDelegate {
             kCurrentFontFamily: 0,
             kNightMode: false,
             kCurrentFontSize: 2,
-            kCurrentAudioRate: 1
+            kCurrentAudioRate: 1,
+            kCurrentHighlightStyle: 0
             ])
     }
 

Source/FolioReaderKit.swift

@@ -25,6 +25,7 @@ internal let kApplicationDocumentsDirectory = NSSearchPathForDirectoriesInDomain
 internal let kCurrentFontFamily = "kCurrentFontFamily"
 internal let kCurrentFontSize = "kCurrentFontSize"
 internal let kCurrentAudioRate = "kCurrentAudioRate"
+internal let kCurrentHighlightStyle = "kCurrentHighlightStyle"
 internal let kNightMode = "kNightMode"
 internal let kHighlightRange = 30
 internal var kBookId: String!
@@ -75,6 +76,14 @@ public class FolioReader {
         }
     }
 
+    var currentHighlightStyle: Int {
+        get { return FolioReader.defaults.valueForKey(kCurrentHighlightStyle) as! Int }
+        set (value) {
+            FolioReader.defaults.setValue(value, forKey: kCurrentHighlightStyle)
+            FolioReader.defaults.synchronize()
+        }
+    }
+
     // MARK: - Present Folio Reader
 
     /**

Source/FolioReaderPage.swift

@@ -85,6 +85,9 @@ class FolioReaderPage: UICollectionViewCell, UIWebViewDelegate, UIGestureRecogni
                     let newRange = NSRange(location: range.location + item.contentPre.characters.count, length: item.content.characters.count)
                     html = html.stringByReplacingCharactersInRange(newRange, withString: tag)
                 }
+                else {
+                    print("highlight range not found")
+                }
             }
         }
 
@@ -380,7 +383,7 @@ extension UIWebView {
     }
 
     func highlight(sender: UIMenuController?) {
-        let highlightAndReturn = js("highlightString()")
+        let highlightAndReturn = js("highlightString('\(HighlightStyle.classForStyle(FolioReader.sharedInstance.currentHighlightStyle))')")
         let jsonData = highlightAndReturn?.dataUsingEncoding(NSUTF8StringEncoding)
 
         do {
@@ -431,36 +434,30 @@ extension UIWebView {
     // MARK: - Set highlight styles
 
     func setYellow(sender: UIMenuController?) {
-        if let updateId = js("setYellow()") {
-            Highlight.updateById(updateId, type: .Yellow)
-        }
-        colors(sender)
+        changeHighlightStyle(sender, style: .Yellow)
     }
 
     func setGreen(sender: UIMenuController?) {
-        if let updateId = js("setGreen()") {
-            Highlight.updateById(updateId, type: .Green)
-        }
-        colors(sender)
+        changeHighlightStyle(sender, style: .Green)
     }
 
     func setBlue(sender: UIMenuController?) {
-        if let updateId = js("setBlue()") {
-            Highlight.updateById(updateId, type: .Blue)
-        }
-        colors(sender)
+        changeHighlightStyle(sender, style: .Blue)
     }
 
     func setPink(sender: UIMenuController?) {
-        if let updateId = js("setPink()") {
-            Highlight.updateById(updateId, type: .Pink)
-        }
-        colors(sender)
+        changeHighlightStyle(sender, style: .Pink)
     }
 
     func setUnderline(sender: UIMenuController?) {
-        if let updateId = js("setUnderline()") {
-            Highlight.updateById(updateId, type: .Underline)
+        changeHighlightStyle(sender, style: .Underline)
+    }
+
+    func changeHighlightStyle(sender: UIMenuController?, style: HighlightStyle) {
+        FolioReader.sharedInstance.currentHighlightStyle = style.rawValue
+
+        if let updateId = js("setHighlightStyle('\(HighlightStyle.classForStyle(style.rawValue))')") {
+            Highlight.updateById(updateId, type: style)
         }
         colors(sender)
     }

Source/Resources/Bridge.js

@@ -81,7 +81,7 @@ function setFontSize(cls) {
 /*
  *	Native bridge Highlight text
  */
-function highlightString() {
+function highlightString(style) {
     var range = window.getSelection().getRangeAt(0);
     var selectionContents = range.extractContents();
     var elm = document.createElement("highlight");
@@ -90,7 +90,7 @@ function highlightString() {
     elm.appendChild(selectionContents);
     elm.setAttribute("id", id);
     elm.setAttribute("onclick","callHighlightURL(this);");
-    elm.setAttribute("class", "highlight-yellow");
+    elm.setAttribute("class", style);
 
     range.insertNode(elm);
     thisHighlight = elm;
@@ -102,28 +102,8 @@ function highlightString() {
 }
 
 // Menu colors
-function setYellow() {
-    thisHighlight.className = "highlight-yellow";
-    return thisHighlight.id;
-}
-
-function setGreen() {
-    thisHighlight.className = "highlight-green";
-    return thisHighlight.id;
-}
-
-function setBlue() {
-    thisHighlight.className = "highlight-blue";
-    return thisHighlight.id;
-}
-
-function setPink() {
-    thisHighlight.className = "highlight-pink";
-    return thisHighlight.id;
-}
-
-function setUnderline() {
-    thisHighlight.className = "highlight-underline";
+function setHighlightStyle(style) {
+    thisHighlight.className = style;
     return thisHighlight.id;
 }