Minimalistic badge extension for UIView and UIBarButtonItem

UIView+Badge.swift

//
//  UIView+Badge.swift
//  Created by Gökhan KOCA on 4.11.2019.
//

import UIKit

public extension UIView {
    
    var badgeCount: Int {
        get {
            if let badgeText = getBadgeLabel()?.text, let count = Int(badgeText) {
                return count
            }
            return 0
        }
        set {
            self.setBadge(newValue)
        }
    }
    
    fileprivate func getBadgeLabel() -> BadgeLabel? {
        guard let superView = superview else { return nil }
        for case let view in superView.subviews where view is BadgeLabel && view.tag == 64463 {
            return (view as! BadgeLabel)
        }
        return nil
    }
    
    func setBadge(_ value: String, to position: Pose = .topRight, badgeSize: BadgeSize = .normal, badgeBackgroundColor: UIColor = .systemRed, badgeTextColor: UIColor = .white) {
        guard let superView = superview else { return }
        var acceptedValue = ""
        if value.count > 6 { acceptedValue = String(value.prefix(3)) + "..." }
        else { acceptedValue = value }
        func removeBadge() {
            if let view = getBadgeLabel() {
                view.removeFromSuperview()
            }
        }
        
        func addBadge(with text: String) {
            removeBadge()
            if !text.isEmpty {
                let badgeLabel = BadgeLabel(frame: .zero)
                badgeLabel.text = acceptedValue
                badgeLabel.badgeStyle(backgroundColor: badgeBackgroundColor, textColor: badgeTextColor, badgeSize: badgeSize)
                let position = position.getValue()
                superView.addSubview(badgeLabel)
                superView.bringSubviewToFront(badgeLabel)
                var horizontalValue: CGFloat = 0
                var verticalValue: CGFloat = 0
                let badgeOffset = badgeLabel.frame.size.height / 2
                let badgeWidth = badgeLabel.frame.size.width
                verticalValue = position.top ? self.frame.origin.y : self.frame.origin.y + self.frame.size.height
                horizontalValue = position.right ? self.frame.origin.x + self.frame.size.width : self.frame.origin.x
                horizontalValue = position.right ? horizontalValue - (badgeWidth / 2 - badgeOffset) : horizontalValue + (badgeWidth / 2 - badgeOffset)
                badgeLabel.center = CGPoint(x: horizontalValue, y: verticalValue)
                badgeLabel.autoresizingMask = [.flexibleWidth]
            }
        }
        addBadge(with: acceptedValue)
    }
    
    func setBadge(_ value: Int, to position: Pose = .topRight, badgeSize: BadgeSize = .normal, badgeBackgroundColor: UIColor = .systemRed, badgeTextColor: UIColor = .white) {
        var badgeValue = "\(value)"
        if value < 1 {
            badgeValue = ""
        } else if value > 99999 {
            badgeValue = "99999+"
        }
        setBadge(badgeValue, to: position, badgeSize: badgeSize, badgeBackgroundColor: badgeBackgroundColor, badgeTextColor: badgeTextColor)
    }
}

extension UIBarButtonItem {
    
    var badgeCount: Int {
        get {
            if let holder = badgeViewHolder {
                return getView(in: holder).badgeCount
            }
            return 0
        }
        set {
            self.setBadge(newValue)
        }
    }
    
    private var badgeViewHolder: UIView? {
        return value(forKey: "view") as? UIView
    }
    
    private func getView(in holder: UIView) -> UIView {
        for sub in holder.subviews {
            if "\(type(of: sub))" == "_UIModernBarButton" {
                return sub
            }
        }
        return holder
    }
    
    func setBadge(_ value: String, to position: Pose = .topRight, badgeSize: BadgeSize = .little, badgeBackgroundColor: UIColor = .systemRed, badgeTextColor: UIColor = .white) {
        if let view = badgeViewHolder {
            getView(in: view).setBadge(value, to: position, badgeSize: badgeSize, badgeBackgroundColor: badgeBackgroundColor, badgeTextColor: badgeTextColor)
        } else {
            print("UIBarButtonItem is nil.")
        }
    }
    
    func setBadge(_ value: Int, to position: Pose = .topRight, badgeSize: BadgeSize = .little, badgeBackgroundColor: UIColor = .systemRed, badgeTextColor: UIColor = .white) {
        if let view = badgeViewHolder {
            getView(in: view).setBadge(value, to: position, badgeSize: badgeSize, badgeBackgroundColor: badgeBackgroundColor, badgeTextColor: badgeTextColor)
        } else {
            print("UIBarButtonItem is nil.")
        }
    }
}

private extension UILabel {
    func badgeStyle(backgroundColor: UIColor = .systemRed, textColor: UIColor = .white, badgeSize: BadgeSize = .normal) {
        var fontSize: CGFloat = 0
        var verticalPadding: CGFloat = 0
        var labelHeight: CGFloat = 0
        switch badgeSize {
        case .little:
            fontSize = 8
            verticalPadding = 3
            labelHeight = 15
        case .normal:
            fontSize = 12
            verticalPadding = 4
            labelHeight = 20
        case .big:
            fontSize = 18
            verticalPadding = 6
            labelHeight = 30
        }
        self.backgroundColor = backgroundColor
        self.textColor = textColor
        self.textAlignment = .center
        self.numberOfLines = 1
        self.font = UIFont.systemFont(ofSize: fontSize)
        self.clipsToBounds = true
        self.tag = 64463
        self.sizeToFit()
        let labelWidth = self.frame.size.width + verticalPadding < labelHeight ? labelHeight : self.frame.size.width + verticalPadding + verticalPadding
        let newFrame = CGRect(origin: self.frame.origin, size: CGSize(width: labelWidth, height: labelHeight))
        self.frame = newFrame
        self.layer.cornerRadius = self.frame.size.height / 2
    }
}

public enum BadgeSize {
    case little
    case normal
    case big
}

public enum Pose {
    case topLeft
    case topRight
    case bottomRight
    case bottomLeft
    
    func getValue() -> (top: Bool, right: Bool) {
        switch self {
        case .topLeft:
            return (true, false)
        case .topRight:
            return (true, true)
        case .bottomRight:
            return (false, true)
        case .bottomLeft:
            return (false, false)
        }
    }
}

private class BadgeLabel: UILabel {
    public override init(frame: CGRect) {
        super.init(frame: frame)
    }
    
    required init?(coder: NSCoder) {
        fatalError("init(coder:) has not been implemented")
    }
}