How to trigger UIImageView animation in a UITableViewHeaderFooterView when pressing on a tableView section?

I want to create an expandable/collapsable tableView cells behaviour with a custom header view from xib using UITableViewHeaderFooterView.

I have the following CustomHeaderView xib setup: UILabel, UIButton (to press on the header and trigger expandable/collapsable behaviour and chevron as UIImageView to represent a current section state:

Here is how a CustomHeaderView class looks. It has some outlets, method for rotating the chevron and simple protocol to tell my ViewController that certain section was tapped (button's tag == section)

import UIKit

 protocol HeaderViewDelegate: AnyObject {
    func expandedSection(button: UIButton)
 }

class CustomHeaderView: UITableViewHeaderFooterView {

weak var delegate: HeaderViewDelegate?

@IBOutlet weak var lable: UILabel!
@IBOutlet weak var imageView: UIImageView!
@IBOutlet weak var headerButton: UIButton!

func configure(title: String, section: Int) {
    lable.text = title
    headerButton.tag = section
}

func rotateImage(_ expanded: Bool) {
    if expanded {
        UIView.animate(withDuration: 0.3) {
            self.imageView.transform = CGAffineTransform(rotationAngle: .pi / 2)
        }
    } else {
        UIView.animate(withDuration: 0.3) {
            self.imageView.transform = .identity
        }
    }
}

@IBAction func tapHeader(_ sender: UIButton) {
    delegate?.expandedSection(button: sender)
}
}

Here is how my ViewController looks like:

import UIKit

struct ExpandedModel {
   var isExpanded: Bool
   let title: String
   let array: [String]
}

class TableViewController: UITableViewController {

   let headerID = String(describing: CustomHeaderView.self)
   var arrayOfData = [ExpandedModel]()

override func viewDidLoad() {
    super.viewDidLoad()
    
    arrayOfData = [
        ExpandedModel(isExpanded: true, title: "Words", array: ["One", "Two", "Three", "Four", "Five"]),
        ExpandedModel(isExpanded: true, title: "Numbers", array: ["6", "7", "8", "9", "10"])
    ]
    tableViewConfig()
}

private func tableViewConfig() {
    let nib = UINib(nibName: headerID, bundle: nil)
    tableView.register(nib, forHeaderFooterViewReuseIdentifier: headerID)
}

override func numberOfSections(in tableView: UITableView) -> Int {
    arrayOfData.count
}

override func tableView(_ tableView: UITableView, numberOfRowsInSection section: Int) -> Int {
    if !arrayOfData[section].isExpanded {
        return 0
    } else {
        return arrayOfData[section].array.count
    }
}

override func tableView(_ tableView: UITableView, cellForRowAt indexPath: IndexPath) -> UITableViewCell {
    let cell = tableView.dequeueReusableCell(withIdentifier: "dataCell", for: indexPath) as! DataTableViewCell
    cell.label.text = arrayOfData[indexPath.section].array[indexPath.row]
    return cell
}

override func tableView(_ tableView: UITableView, viewForHeaderInSection section: Int) -> UIView? {
    let header = tableView.dequeueReusableHeaderFooterView(withIdentifier: headerID) as! CustomHeaderView
    
    header.configure(title: arrayOfData[section].title, section: section)
    header.rotateImage(arrayOfData[section].isExpanded)
    header.delegate = self
    return header
}

override func tableView(_ tableView: UITableView, heightForHeaderInSection section: Int) -> CGFloat {
    return tableView.estimatedSectionHeaderHeight
}

override func tableView(_ tableView: UITableView, heightForFooterInSection section: Int) -> CGFloat {
    return tableView.estimatedSectionFooterHeight
}
}

 extension TableViewController: HeaderViewDelegate {
    func expandedSection(button: UIButton) {
       let section = button.tag
       let indexPaths = (0..<(arrayOfData[section].array.count)).map { IndexPath(row: $0, section: section)}
    
    arrayOfData[section].isExpanded.toggle()
    
    if !arrayOfData[section].isExpanded {
        tableView.beginUpdates()
        tableView.deleteRows(at: indexPaths, with: .fade)
        tableView.endUpdates()
    } else {
        tableView.beginUpdates()
        tableView.insertRows(at: indexPaths, with: .fade)
        tableView.endUpdates()
    }
}
}

The main goal is ok - I can collapse and expand cells. But the problem is how to animate the chevron in CustomHeaderView?

If I'll use a reloading section method, then the chevron animation will be skipped:

tableView.reloadSections(IndexSet(integer: section), with: .none)

If I will add a delay with DispatchQueue.main.asyncAfter, it will be animated in, but skipped when out:

        func rotateImage(_ expanded: Bool) {
        if expanded {
            DispatchQueue.main.asyncAfter(deadline: .now() + 0.01) {
                UIView.animate(withDuration: 0.2) {
                self.imageView.transform = CGAffineTransform(rotationAngle: .pi / 2)
                }
            }
        } else {
            DispatchQueue.main.asyncAfter(deadline: .now() + 0.01) {
                UIView.animate(withDuration: 0.2) {
                self.imageView.transform = .identity
                }
            }
        }
    }

What is a most sufficient way to achieve a proper chevron animation?

UPDATE

According to DonMag's answer here is the working solution:

extension TableViewController: HeaderViewDelegate {
func expandedSection(button: UIButton) {
    let section = button.tag
    let indexPaths = (0..<(arrayOfData[section].array.count)).map { IndexPath(row: $0, section: section)}
    
    arrayOfData[section].isExpanded.toggle()
    
    tableView.beginUpdates()
    
    if let header = tableView.headerView(forSection: section) as? CustomHeaderView {
                header.rotateImage(arrayOfData[section].isExpanded)
    }
    
    if !arrayOfData[section].isExpanded {
        tableView.deleteRows(at: indexPaths, with: .fade)
    } else {
        tableView.insertRows(at: indexPaths, with: .fade)
    }
    tableView.endUpdates()
}
}

Using your first block of code, do you get the animation you want if you change your expandedSection(button: UIButton) func to this:

extension TableViewController: HeaderViewDelegate {
    func expandedSection(button: UIButton) {
        let section = button.tag
        let indexPaths = (0..<(arrayOfData[section].array.count)).map { IndexPath(row: $0, section: section)}
        
        arrayOfData[section].isExpanded.toggle()

        // get a reference to the section's header view    
        if let v = tableView.headerView(forSection: section) as? SomeCustomHeaderView {
            v.rotateImage(arrayOfData[section].isExpanded)
        }

        if !arrayOfData[section].isExpanded {
            tableView.beginUpdates()
            tableView.deleteRows(at: indexPaths, with: .fade)
            tableView.endUpdates()
        } else {
            tableView.beginUpdates()
            tableView.insertRows(at: indexPaths, with: .fade)
            tableView.endUpdates()
        }
    }
}