Change Button image from view controller inside custom header cell - swift

Question

I have a VC that is UICollectionView that has header. This header is a custom cell, and inside that cell there is an button with image which I need to change based on info from API or user tap.
In func loadData i get info about icon and then I call function inside headerCell to change the icon. Do I need to do this in other way in order for this to work?
protocol HeaderCellDelegate: class {
func setDropdown()
func tapOnSelection()
}
class HeaderCell: UIGestureRecognizerDelegate {
// delegate
weak var hDelegate: HeaderCellDelegate?
#objc
lazy var selectionBtn: UIButton = {
let button = UIButton(type: .system)
button.setTitle("One", for: .normal)
button.addTarget(self, action: #selector(changeSelection), for: .touchUpInside)
return button
}()
#objc
lazy var oneOrTwo: UIButton = {
let button = UIButton()
button.setImage(UIImage(named: "oneOff")?.withRenderingMode(.alwaysTemplate), for: .normal)
button.addTarget(self, action: #selector(toggleOneOrTwo), for: .touchUpInside)
return button
}()
#objc
lazy var iconBtn: UIButton = {
let button = UIButton()
button.setImage(UIImage(named: "iconGreen")?.withRenderingMode(.alwaysTemplate), for: .normal)
button.addTarget(self, action: #selector(setIcon), for: .touchUpInside)
return button
}()
override init(frame: CGRect) {
super.init(frame: frame)
// I setup view helper
// Just a button and icon on the right
// button opens DropDown list
}
required init?(coder aDecoder: NSCoder) {
fatalError("init(coder:) has not been implemented")
}
// MARK: - Delegate Functions
#objc
func toggleOneOrTwo(_ sender: Any) {
self.hDelegate?.setDropdown()
}
#objc
func changeSelection(sender: UIButton) {
self.hDelegate?.tapOnSelection()
}
#objc
func setIcon() {
isIconGreen
? iconColor.setImage(UIImage(named: "iconGreen")?.withRenderingMode(.alwaysTemplate), for: .normal)
: iconColor.setImage(UIImage(named: "iconRed")?.withRenderingMode(.alwaysTemplate), for: .normal)
}
}
ViewController
class ViewController: UICollectionViewController, UIGestureRecognizerDelegate {
fileprivate let headerCellId = "headerCellId"
lazy var headerCell = HeaderCell()
override func viewDidLoad() {
super.viewDidLoad()
// Register Cell Class
self.collectionView.register(HeaderCell.self, forSupplementaryViewOfKind: UICollectionView.elementKindSectionHeader, withReuseIdentifier: headerCellId)
loadData()
}
override func collectionView(_ collectionView: UICollectionView, viewForSupplementaryElementOfKind kind: String, at indexPath: IndexPath) -> UICollectionReusableView {
switch kind {
case UICollectionView.elementKindSectionHeader:
let header = collectionView.dequeueReusableSupplementaryView(ofKind: kind, withReuseIdentifier: headerCellId, for: indexPath) as! HeaderCell
header.backgroundColor = .groupTableViewBackground
header.hDelegate = self
return header
default:
fatalError("This should never happen!!")
}
}
func loadData() {
// get data from API and isIconGreen value
let isIconGreen = true
headerCell.setIcon(isIconGreen: isIconGreen)
}
}