UICollectionView with dynamically-scaling cells

My goal is to create a layout presented below:

I know how to create these custom UICollectionViewCells, but I trouble with the layout. All of the shown cells differ in width, so there can be, for instance: four in the first row, two in the second, and the last one remaining - in the third. That's just one possible configuration, and there are many more (including one shown on an image) depending on the label's width.

I'm building everything programmatically. Also I feel like using UICollectionView is the best choice, but I'm open to any suggestions.

Thanks in advance!

What I've already tried:

let collectionView = UICollectionView(frame: .zero, collectionViewLayout: TagLayout()

override func viewDidLoad() {
    super.viewDidLoad()
    setupCollectionView()
}

private func setupCollectionView() {
    collectionView.backgroundColor = .systemGray5
    collectionView.dataSource = self
    collectionView.delegate = self
    collectionView.register(SubjectCollectionViewCell.self, forCellWithReuseIdentifier: "cell")
        
    //adding a view to subview and constraining it programmatically using Stevia
}

extension ProfileVC: UICollectionViewDelegate, UICollectionViewDataSource {
    func collectionView(_ collectionView: UICollectionView, numberOfItemsInSection section: Int) -> Int {
        return 5
    }
    
    
    func collectionView(_ collectionView: UICollectionView, cellForItemAt indexPath: IndexPath) -> UICollectionViewCell {
        guard let cell = collectionView.dequeueReusableCell(withReuseIdentifier: "cell", for: indexPath) as? SubjectCollectionViewCell else { return UICollectionViewCell() }
        cell.data = SubjectTagData(emoji: "🇬🇧", subjectName: "Item I")
        
        return cell
    }
}

Use following collectionViewLayout

    // MARK: - TagLayoutDelegate
    protocol TagLayoutDelegate: class {
        func widthForItem(at indexPath: IndexPath) -> CGFloat
        func rowHeight() -> CGFloat
    }

    // MARK: - TagLayout
    class TagLayout: UICollectionViewLayout {

    // MARK: Variables
    weak var delegate       : TagLayoutDelegate?
    var cellPadding         : CGFloat = 5.0
    var deafultRowHeight    : CGFloat = 35.0
    var scrollDirection     : UICollectionView.ScrollDirection = .vertical
    
    private var contentWidth: CGFloat = 0
    private var contentHeight: CGFloat = 0
    
    private var cache: [UICollectionViewLayoutAttributes] = []
    
    // MARK: Public Functions
    func reset() {
        cache.removeAll()
        contentHeight = 0
        contentWidth = 0
    }
    
    // MARK: Override
    override var collectionViewContentSize: CGSize {
        return CGSize(width: contentWidth, height: contentHeight)
    }
    
    override func prepare() {
        super.prepare()
        if scrollDirection == .vertical {
            prepareForVerticalScroll()
        }
        else {
            prepareForHorizontalScroll()
        }
    }
    
    override func layoutAttributesForElements(in rect: CGRect) -> [UICollectionViewLayoutAttributes]? {
        var visibleLayoutAttributeElements: [UICollectionViewLayoutAttributes] = []
        for attribute in cache {
            if attribute.frame.intersects(rect) {
                visibleLayoutAttributeElements.append(attribute)
            }
        }
        return visibleLayoutAttributeElements
    }
    
    override func layoutAttributesForItem(at indexPath: IndexPath) -> UICollectionViewLayoutAttributes? {
        return cache[indexPath.item]
    }
    
    // MARK: Private Functions
    private func prepareForVerticalScroll() {
        guard cache.isEmpty, let collectionView = collectionView else {
            return
        }
        
        let noOfItems   = collectionView.numberOfItems(inSection: 0)
        var xOffset     = [CGFloat](repeating: 0.0, count: noOfItems)
        var yOffset     = [CGFloat](repeating: 0.0, count: noOfItems)
        
        let insets      = collectionView.contentInset
        contentWidth = collectionView.bounds.width - (insets.left + insets.right)
        
        var rowWidth: CGFloat = 0
        for i in 0 ..< noOfItems {
            let indexPath = IndexPath(item: i, section: 0)
            let textWidth = delegate?.widthForItem(at: indexPath) ?? 75.0
            let width = textWidth + cellPadding
            let height = delegate?.rowHeight() ?? 30.0
            let frame = CGRect(
                x: xOffset[i],
                y: yOffset[i],
                width: width,
                height: height
            )
            let insetFrame = frame.insetBy(dx: cellPadding, dy: cellPadding)
            
            let attributes = UICollectionViewLayoutAttributes(forCellWith: indexPath)
            attributes.frame = insetFrame
            cache.append(attributes)
            
            contentHeight = max(contentHeight, frame.maxY)
            
            
            rowWidth += frame.width
            
            if i < noOfItems-1 {
                let nextIP = IndexPath(item: i+1, section: 0)
                let nextWidth = delegate?.widthForItem(at: nextIP) ?? 75.0
                if rowWidth + nextWidth + cellPadding <= contentWidth {
                    xOffset[i+1] = xOffset[i] + width
                    yOffset[i+1] = yOffset[i]
                }
                else {
                    rowWidth = 0
                    yOffset[i+1] = yOffset[i] + (delegate?.rowHeight() ?? 30.0)
                }
            }
        }
    }
    
    private func prepareForHorizontalScroll() {
        guard cache.isEmpty, let collectionView = collectionView else {
            return
        }
        
        let insets = collectionView.contentInset
        contentHeight = collectionView.bounds.height - (insets.top + insets.bottom)
        
        let rowHeight = delegate?.rowHeight() ?? deafultRowHeight
        
        let noOfRows: Int = 2
        
        var yOffset: [CGFloat] = []
        for row in 0 ..< noOfRows {
            yOffset.append(CGFloat(row) * rowHeight)
        }
        
        var row = 0
        var xOffset: [CGFloat] = [CGFloat](repeating: 0.0, count: noOfRows)
        
        for i in 0 ..< collectionView.numberOfItems(inSection: 0) {
            let indexPath = IndexPath(item: i, section: 0)
            let textWidth = delegate?.widthForItem(at: indexPath) ?? 75.0
            let width = textWidth + cellPadding
            let frame = CGRect(
                x       : xOffset[row],
                y       : yOffset[row],
                width   : width,
                height  : rowHeight)
            let insetFrame = frame.insetBy(dx: cellPadding, dy: cellPadding)
            
            let attributes = UICollectionViewLayoutAttributes(forCellWith: indexPath)
            attributes.frame = insetFrame
            cache.append(attributes)
            
            contentWidth = max(contentWidth, frame.maxX)
            xOffset[row] = xOffset[row] + width
            
            row = row < (noOfRows - 1) ? row + 1 : 0
        }
     }
   }

And implement it as follows

let tagLayout = TagLayout()
let collectionView = UICollectionView(frame: .zero, collectionViewLayout: tagLayout)

private func setupCollectionView() {
  tagLayout.delegate = self
  //Your other code goes here
}

extension ProfileVC: TagLayoutDelegate {
  func widthForItem(at indexPath: IndexPath) -> CGFloat {
    return 100.0
  }
  func rowHeight() -> CGFloat {
    return 30.0
  }
}