import UIKit

class HighlightedTextLabel: UILabel {
    var highlightColor: UIColor!
    
    override func draw(_ rect: CGRect) {
        let separatedLines = self.getSeparatedLines()
        let totalNumberOfLines = separatedLines.count
        if (totalNumberOfLines < 1) {
            super.draw(rect)
            return
        }
        
        let fontSize = self.font.pointSize
        let lineFillerSpaceWidth = CGFloat(9)
        let firstLineYOffset = CGFloat(1.0) // otherwise first line would look thinner because it would go outside of label bounds
        let interHighlightSeparation = CGFloat(3)
        let intraHighlightSeparation = CGFloat(fontSize + 3)
        
        let pi = CGFloat.pi
        
        let context = UIGraphicsGetCurrentContext()!
    
        // Draw parallel lines
        CGContext.setLineWidth(context)(1.5)
        CGContext.setStrokeColor(context)(highlightColor.cgColor)
        for (lineIndex, line) in separatedLines.enumerated() {
            let lineWidth = getLineWidth(line: line as! String)
            let lineStartX = lineIndex == 0 ? lineFillerSpaceWidth : 0
            let lineEndX = lineStartX + lineWidth
            let lineStartY = lineIndex == 0 ? firstLineYOffset : firstLineYOffset + ((intraHighlightSeparation + interHighlightSeparation) * CGFloat(lineIndex))
            let lineEndY = lineStartY + intraHighlightSeparation

            CGContext.beginPath(context)()
            CGContext.move(context)(to: CGPoint(x: lineStartX, y: lineStartY))
            CGContext.addLine(context)(to: CGPoint(x: lineEndX, y: lineStartY))
            CGContext.closePath(context)()
            CGContext.strokePath(context)()

            CGContext.beginPath(context)()
            CGContext.move(context)(to: CGPoint(x: lineStartX, y: lineEndY))
            CGContext.addLine(context)(to: CGPoint(x: lineEndX, y: lineEndY))
            CGContext.closePath(context)()
            CGContext.strokePath(context)()
        }

        // Draw circle at start of first line
        CGContext.beginPath(context)()
        let initialArcX = lineFillerSpaceWidth
        let initialArcRadius = intraHighlightSeparation / 2
        let initialArcY = firstLineYOffset + initialArcRadius
        let initialArcStartAngle = -pi/2
        let initialArcEndAngle = pi/2
        let initialCircleCenter = CGPoint(x: initialArcX, y: initialArcY)
        CGContext.addArc(context)(
            center: initialCircleCenter,
            radius: initialArcRadius,
            startAngle: initialArcStartAngle,
            endAngle: initialArcEndAngle,
            clockwise: true
        )
        CGContext.move(context)(to: CGPoint(x: 0, y: 0))
        CGContext.closePath(context)()
        CGContext.strokePath(context)()

        
        // Draw circle at end of last line
        let lastLineWidth = getLineWidth(line: separatedLines.last as! String)
        CGContext.beginPath(context)()
        let finalArcX = (totalNumberOfLines == 1 ? lineFillerSpaceWidth : 0) + lastLineWidth
        let finalArcRadius = intraHighlightSeparation / 2
        let finalArcY = firstLineYOffset + ((intraHighlightSeparation + interHighlightSeparation) * CGFloat(totalNumberOfLines - 1)) + initialArcRadius
        let finalArcStartAngle = pi/2
        let finalArcEndAngle = -pi/2
        let finalCircleCenter = CGPoint(x: finalArcX, y: finalArcY)
        CGContext.addArc(context)(
            center: finalCircleCenter,
            radius: finalArcRadius,
            startAngle: finalArcStartAngle,
            endAngle: finalArcEndAngle,
            clockwise: true
        )
        CGContext.move(context)(to: CGPoint(x: 0, y: 0))
        CGContext.closePath(context)()
        CGContext.strokePath(context)()
        
        super.draw(rect)
    }
    
    func getSeparatedLines() -> [Any] {
        if self.lineBreakMode != NSLineBreakMode.byWordWrapping {
            self.lineBreakMode = .byWordWrapping
        }
        var lines = [Any]() /* capacity: 10 */
        let wordSeparators = CharacterSet.whitespacesAndNewlines
        var currentLine: String? = self.text
        let textLength: Int = (self.text?.count ?? 0)
        var rCurrentLine = NSRange(location: 0, length: textLength)
        var rWhitespace = NSRange(location: 0, length: 0)
        var rRemainingText = NSRange(location: 0, length: textLength)
        var done: Bool = false
        while !done {
            // determine the next whitespace word separator position
            rWhitespace.location = rWhitespace.location + rWhitespace.length
            rWhitespace.length = textLength - rWhitespace.location
            rWhitespace = (self.text! as NSString).rangeOfCharacter(from: wordSeparators, options: .caseInsensitive, range: rWhitespace)
            if rWhitespace.location == NSNotFound {
                rWhitespace.location = textLength
                done = true
            }
            let rTest = NSRange(location: rRemainingText.location, length: rWhitespace.location - rRemainingText.location)
            let textTest: String = (self.text! as NSString).substring(with: rTest)
            let fontAttributes: [String: Any]? = [NSFontAttributeName: font]
            let maxWidth = (textTest as NSString).size(attributes: fontAttributes).width
            if maxWidth > self.bounds.size.width {
                lines.append(currentLine?.trimmingCharacters(in: wordSeparators) ?? "")
                rRemainingText.location = rCurrentLine.location + rCurrentLine.length
                rRemainingText.length = textLength - rRemainingText.location
                continue
            }
            rCurrentLine = rTest
            currentLine = textTest
        }
        lines.append(currentLine?.trimmingCharacters(in: wordSeparators) ?? "")
        return lines
    }

    func getLineWidth(line: String) -> CGFloat {
        let fontAttributes: [String: Any]? = [NSFontAttributeName: self.font]
        return (line as NSString).size(attributes: fontAttributes).width
    }
}