Tusker/Tusker/MultiColumnCollectionViewLayout.swift

//
//  MultiColumnCollectionViewLayout.swift
//  Tusker
//
//  Created by Shadowfacts on 9/29/23.
//  Copyright © 2023 Shadowfacts. All rights reserved.
//

import UIKit

class MultiColumnCollectionViewLayout: UICollectionViewLayout {
    
    private let numberOfColumns: Int
    private let spacing: CGFloat
    private let minimumColumnWidth: CGFloat
    private var effectiveNumberOfColumns: Int!
    
    private var attributes: [MultiColumnLayoutAttributes] = []
    private var invalidatedItemIndices: IndexSet = []
    
    var showSectionHeader = false {
        didSet {
            if showSectionHeader != oldValue {
                invalidateLayout()
            }
        }
    }
    private var sectionHeaderAttributes: MultiColumnLayoutAttributes?
    
    init(numberOfColumns: Int, columnSpacing: CGFloat, minimumColumnWidth: CGFloat) {
        precondition(numberOfColumns >= 1)
        self.numberOfColumns = numberOfColumns
        self.effectiveNumberOfColumns = nil
        self.spacing = columnSpacing
        self.minimumColumnWidth = minimumColumnWidth
        super.init()
    }
    
    required init?(coder: NSCoder) {
        fatalError("init(coder:) has not been implemented")
    }
    
    override class var layoutAttributesClass: AnyClass {
        MultiColumnLayoutAttributes.self
    }
    
    override class var invalidationContextClass: AnyClass {
        MultiColumnLayoutInvalidationContext.self
    }
    
    override func prepare() {
        guard let collectionView else { return }
        precondition(collectionView.numberOfSections <= 1)
        
        guard collectionView.numberOfSections == 1 else {
            attributes = []
            return
        }
        
        if effectiveNumberOfColumns == nil {
            updateEffectiveNumberOfColumns()
        }
        
        var lastAttributesInEachColumn: [MultiColumnLayoutAttributes?] = Array(repeating: nil, count: effectiveNumberOfColumns)
        func columnWithMinHeight() -> Int {
            var min: Int?
            for i in 0..<effectiveNumberOfColumns {
                guard let attrs = lastAttributesInEachColumn[i] else {
                    return i
                }
                if min == nil || attrs.frame.maxY < lastAttributesInEachColumn[min!]!.frame.maxY {
                    min = i
                }
            }
            return min!
        }
        
        let columnWidth = columnWidthForColumnCount(effectiveNumberOfColumns)
        func minXForColumn(_ column: Int) -> CGFloat {
            (CGFloat(column) * columnWidth) + ((CGFloat(column) + 1) * spacing)
        }
        
        let startY: CGFloat
        if showSectionHeader {
            let indexPath = IndexPath(item: 0, section: 0)
            if sectionHeaderAttributes == nil {
                sectionHeaderAttributes = .init(forSupplementaryViewOfKind: UICollectionView.elementKindSectionHeader, with: indexPath)
            }
            sectionHeaderAttributes!.frame = CGRect(x: 0, y: 0, width: collectionView.bounds.width, height: 0)
            if let view = collectionView.supplementaryView(forElementKind: UICollectionView.elementKindSectionHeader, at: indexPath) {
                let preferred = view.preferredLayoutAttributesFitting(sectionHeaderAttributes!)
                sectionHeaderAttributes!.frame.size.height = preferred.frame.height
            }
            startY = sectionHeaderAttributes!.frame.height
        } else {
            startY = spacing
        }
        
        let numberOfItems = collectionView.numberOfItems(inSection: 0)
        for item in 0..<numberOfItems {
            let column = columnWithMinHeight()
            
            let itemAttrs: MultiColumnLayoutAttributes
            if attributes.count > item {
                itemAttrs = attributes[item]
            } else {
                itemAttrs = MultiColumnLayoutAttributes(forCellWith: IndexPath(item: item, section: 0))
                // estimate
                itemAttrs.frame.size.height = 200
                attributes.append(itemAttrs)
            }
            
            itemAttrs.column = column
            itemAttrs.frame.size.width = columnWidth
            
            if invalidatedItemIndices.contains(item) {
                if let cell = collectionView.cellForItem(at: IndexPath(item: item, section: 0)) {
                    let preferred = cell.preferredLayoutAttributesFitting(itemAttrs.copy() as! UICollectionViewLayoutAttributes)
                    itemAttrs.frame.size.height = preferred.frame.height
                }
            }
            
            if let lastInColumn = lastAttributesInEachColumn[column] {
                itemAttrs.frame.origin = CGPoint(
                    x: lastInColumn.frame.minX,
                    y: lastInColumn.frame.maxY + spacing
                )
            } else {
                itemAttrs.frame.origin = CGPoint(
                    x: minXForColumn(column),
                    y: startY
                )
            }
            
            lastAttributesInEachColumn[column] = itemAttrs
        }
        
        if attributes.count > numberOfItems {
            attributes.removeLast(attributes.count - numberOfItems)
        }
        
        invalidatedItemIndices = []
    }
    
    override var collectionViewContentSize: CGSize {
        guard let collectionView else {
            return .zero
        }
        let maxY = attributes.lazy.map(\.frame.maxY).max()
        return CGSize(width: collectionView.bounds.width, height: (maxY ?? 0) + spacing)
    }
    
    override func layoutAttributesForItem(at indexPath: IndexPath) -> UICollectionViewLayoutAttributes? {
        return attributes[indexPath.item]
    }
    
    override func layoutAttributesForElements(in rect: CGRect) -> [UICollectionViewLayoutAttributes]? {
        // TODO: optimize this
        var attributes = attributes.filter { $0.frame.intersects(rect) }
        if showSectionHeader,
           let sectionHeaderAttributes,
           rect.minY <= sectionHeaderAttributes.frame.maxY {
            attributes.append(sectionHeaderAttributes)
        }
        return attributes
    }
    
    override func layoutAttributesForSupplementaryView(ofKind elementKind: String, at indexPath: IndexPath) -> UICollectionViewLayoutAttributes? {
        if elementKind == UICollectionView.elementKindSectionHeader,
           showSectionHeader {
            return sectionHeaderAttributes
        } else {
            return nil
        }
    }
    
    override func indexPathsToInsertForSupplementaryView(ofKind elementKind: String) -> [IndexPath] {
        if elementKind == UICollectionView.elementKindSectionHeader {
            return [IndexPath(item: 0, section: 0)]
        } else {
            return []
        }
    }
    
    override func shouldInvalidateLayout(forBoundsChange newBounds: CGRect) -> Bool {
        return newBounds.width != collectionView?.bounds.width
    }
    
    override func invalidationContext(forBoundsChange newBounds: CGRect) -> UICollectionViewLayoutInvalidationContext {
        let context = super.invalidationContext(forBoundsChange: newBounds) as! MultiColumnLayoutInvalidationContext
        context.updateEffectiveNumberOfColumns = true
        return context
    }
    
    // On landscape-to-portrait device rotations (though seemingly only when inside a UISplitViewController), the collection view
    // doesn't go through the bounds-change codepath. But we may still need to update the effective number of columns,
    // so override this private method (which is called) and configure the invalidation context appropriately.
    override func _invalidationContext(forUpdatedLayoutMargins margins: UIEdgeInsets) -> UICollectionViewLayoutInvalidationContext {
        let context = super._invalidationContext(forUpdatedLayoutMargins: margins) as! MultiColumnLayoutInvalidationContext
        context.updateEffectiveNumberOfColumns = true
        return context
    }
    
    override func shouldInvalidateLayout(forPreferredLayoutAttributes preferredAttributes: UICollectionViewLayoutAttributes, withOriginalAttributes originalAttributes: UICollectionViewLayoutAttributes) -> Bool {
        return preferredAttributes.frame.height != originalAttributes.frame.height
    }
    
    override func invalidationContext(forPreferredLayoutAttributes preferredAttributes: UICollectionViewLayoutAttributes, withOriginalAttributes originalAttributes: UICollectionViewLayoutAttributes) -> UICollectionViewLayoutInvalidationContext {
        let context = super.invalidationContext(forPreferredLayoutAttributes: preferredAttributes, withOriginalAttributes: originalAttributes)
        if let kind = preferredAttributes.representedElementKind {
            context.invalidateSupplementaryElements(ofKind: kind, at: [preferredAttributes.indexPath])
        } else {
            context.invalidateItems(at: [preferredAttributes.indexPath])
        }
        return context
    }
    
    override func invalidateLayout(with context: UICollectionViewLayoutInvalidationContext) {
        let context = context as! MultiColumnLayoutInvalidationContext
        
        for indexPath in context.invalidatedItemIndexPaths ?? [] {
            invalidatedItemIndices.insert(indexPath.item)
        }
        
        if context.invalidateEverything || context.updateEffectiveNumberOfColumns {
            updateEffectiveNumberOfColumns()
        }
        
        super.invalidateLayout(with: context)
    }
    
    private func columnWidthForColumnCount(_ count: Int) -> CGFloat {
        guard let collectionView else { return 0 }
        let spacingTotal = spacing * (CGFloat(count) + 1)
        return (collectionView.bounds.width - spacingTotal) / CGFloat(count)
    }
    
    private func updateEffectiveNumberOfColumns() {
        var n = numberOfColumns
        while n > 1 && columnWidthForColumnCount(n) < minimumColumnWidth {
            n -= 1
        }
        effectiveNumberOfColumns = n
    }
    
}

private class MultiColumnLayoutAttributes: UICollectionViewLayoutAttributes {
    var column: Int = -1
    
    override func copy(with zone: NSZone? = nil) -> Any {
        let copy = super.copy(with: zone) as! MultiColumnLayoutAttributes
        copy.column = column
        return copy
    }
}

private class MultiColumnLayoutInvalidationContext: UICollectionViewLayoutInvalidationContext {
    var updateEffectiveNumberOfColumns: Bool = false
}