SwiftUI UICollectionView Wrapper

//
// Copyright (c) Nathan Tannar
//

import SwiftUI
import Engine // https://github.com/nathantannar4/Engine
import Turbocharger // https://github.com/nathantannar4/Turbocharger
import Transmission // https://github.com/nathantannar4/Transmission

#if os(iOS)

@available(iOS 14.0, *)
public struct CollectionView<
    Header: View,
    Content: View,
    Footer: View,
    Data: RandomAccessCollection
>: View where
    Data.Element: RandomAccessCollection,
    Data.Index: Hashable,
    Data.Element.Element: Identifiable
{
    var data: Data
    var header: (Data.Index) -> Header
    var content: (Data.Element.Element) -> Content
    var footer: (Data.Index) -> Footer

    public init(
        _ sections: Data,
        @ViewBuilder content: @escaping (Data.Element.Element) -> Content,
        @ViewBuilder header: @escaping (Data.Index) -> Header,
        @ViewBuilder footer: @escaping (Data.Index) -> Footer
    ) {
        self.data = sections
        self.header = header
        self.content = content
        self.footer = footer
    }

    public var body: some View {
        CollectionViewBody(
            representable: Representable(),
            data: data,
            header: header,
            content: content,
            footer: footer
        )
    }

    private struct Representable: CollectionViewRepresentable {
        func makeUIView(context: Context, options: CollectionViewLayoutOptions) -> UICollectionView {
            var configuration = UICollectionLayoutListConfiguration(appearance: .plain)
            configuration.headerMode = options.contains(.header) ? .supplementary : .none
            configuration.footerMode = options.contains(.footer) ? .supplementary : .none
            configuration.showsSeparators = false
            configuration.backgroundColor = .clear
            let layout = UICollectionViewCompositionalLayout.list(using: configuration)
            let uiCollectionView = UICollectionView(frame: .zero, collectionViewLayout: layout)
            uiCollectionView.clipsToBounds = false
            uiCollectionView.keyboardDismissMode = .interactive
            // uiCollectionView.transform3D = CATransform3DMakeScale(1, -1, 1)
            return uiCollectionView
        }

        func updateUIView(_ uiView: UICollectionView, context: Context) { }
    }
}

@available(iOS 14.0, *)
extension CollectionView {
    public init<
        Items: RandomAccessCollection
    >(
        _ items: Items,
        @ViewBuilder content: @escaping (Items.Element) -> Content,
        @ViewBuilder header: @escaping (Data.Index) -> Header,
        @ViewBuilder footer: @escaping (Data.Index) -> Footer
    ) where Items: RandomAccessCollection, Items.Element: Identifiable, Data == Array<Items>
    {
        self.init([items], content: content, header: header, footer: footer)
    }

    public init<
        Items: RandomAccessCollection
    >(
        _ items: Items,
        @ViewBuilder content: @escaping (Items.Element) -> Content,
        @ViewBuilder header: @escaping (Data.Index) -> Header
    ) where Items: RandomAccessCollection, Items.Element: Identifiable, Data == Array<Items>, Footer == EmptyView
    {
        self.init(items, content: content, header: header, footer: { _ in EmptyView() })
    }
}

@available(iOS 14.0, *)
extension CollectionView {
    public init<
        Sections: RandomAccessCollection,
        ID: Hashable
    >(
        _ sections: Sections,
        id: KeyPath<Sections.Element.Element, ID>,
        @ViewBuilder content: @escaping (Sections.Element.Element) -> Content,
        @ViewBuilder header: @escaping (Data.Index) -> Header,
        @ViewBuilder footer: @escaping (Data.Index) -> Footer
    ) where Sections.Element: RandomAccessCollection, Sections.Index: Hashable, Data == Array<Array<IdentifiableBox<Sections.Element.Element, ID>>>
    {
        let data: Data = sections.compactMap { items in
            items.compactMap { IdentifiableBox($0, id: id) }
        }
        self.init(data, content: { content($0.value) }, header: header, footer: footer)
    }

    public init<
        Sections: RandomAccessCollection,
        ID: Hashable
    >(
        _ sections: Sections,
        id: KeyPath<Sections.Element.Element, ID>,
        @ViewBuilder content: @escaping (Sections.Element.Element) -> Content,
        @ViewBuilder header: @escaping (Data.Index) -> Header
    ) where Sections.Element: RandomAccessCollection, Sections.Index: Hashable, Data == Array<Array<IdentifiableBox<Sections.Element.Element, ID>>>, Footer == EmptyView
    {
        self.init(sections, id: id, content: content, header: header, footer: { _ in EmptyView() })
    }
}

@available(iOS 14.0, *)
extension CollectionView where Header == EmptyView, Footer == EmptyView {
    public init<
        Items: RandomAccessCollection
    >(
        _ items: Items,
        @ViewBuilder content: @escaping (Items.Element) -> Content
    ) where Items: RandomAccessCollection, Items.Element: Identifiable, Data == Array<Items>
    {
        self.init(items, content: content, header: { _ in EmptyView() }, footer: { _ in EmptyView() })
    }

    public init<
        Items: RandomAccessCollection,
        ID: Hashable
    >(
        _ items: Items,
        id: KeyPath<Items.Element, ID>,
        @ViewBuilder content: @escaping (Items.Element) -> Content
    ) where Items: RandomAccessCollection, Data == Array<Array<IdentifiableBox<Items.Element, ID>>>
    {
        self.init([items], id: id, content: content, header: { _ in EmptyView() }, footer: { _ in EmptyView() })
    }
}

public protocol CollectionViewRepresentable {
    associatedtype UIViewType: UICollectionView
    func makeUIView(context: Context, options: CollectionViewLayoutOptions) -> UIViewType
    func updateUIView(_ uiView: UIViewType, context: Context)

    typealias Context = CollectionViewRepresentableContext
}

public struct CollectionViewRepresentableContext {
    public var environment: EnvironmentValues
    public var transaction: Transaction
}

public struct CollectionViewLayoutOptions: OptionSet {
    public var rawValue: UInt8
    public init(rawValue: UInt8) {
        self.rawValue = rawValue
    }

    /// The `UICollectionViewLayout` should include a header
    public static let header = CollectionViewLayoutOptions(rawValue: 1 << 0)

    /// The `UICollectionViewLayout` should include a footer
    public static let footer = CollectionViewLayoutOptions(rawValue: 1 << 1)
}

struct CollectionViewLayoutOptionsKey: EnvironmentKey {
    static let defaultValue = CollectionViewLayoutOptions()
}

extension EnvironmentValues {
    public var collectionViewLayoutOptions: CollectionViewLayoutOptions {
        get { self[CollectionViewLayoutOptionsKey.self] }
        set { self[CollectionViewLayoutOptionsKey.self] = newValue }
    }
}

@available(iOS 14.0, *)
public struct CollectionViewAdapter<
    Header: View,
    Content: View,
    Footer: View,
    Representable: CollectionViewRepresentable,
    Data: RandomAccessCollection
>: View where
    Data.Element: RandomAccessCollection,
    Data.Index: Hashable,
    Data.Element.Element: Identifiable
{
    var representable: Representable
    var data: Data
    var header: (Data.Index) -> Header
    var content: (Data.Element.Element) -> Content
    var footer: (Data.Index) -> Footer

    public init(
        _ representable: Representable,
        sections: Data,
        @ViewBuilder content: @escaping (Data.Element.Element) -> Content,
        @ViewBuilder header: @escaping (Data.Index) -> Header,
        @ViewBuilder footer: @escaping (Data.Index) -> Footer
    ) {
        self.representable = representable
        self.data = sections
        self.header = header
        self.content = content
        self.footer = footer
    }

    public var body: some View {
        CollectionViewBody(
            representable: representable,
            data: data,
            header: header,
            content: content,
            footer: footer
        )
    }
}

@available(iOS 14.0, *)
extension CollectionViewAdapter {
    public init<
        Items: RandomAccessCollection
    >(
        _ representable: Representable,
        items: Items,
        @ViewBuilder content: @escaping (Items.Element) -> Content,
        @ViewBuilder header: @escaping (Data.Index) -> Header,
        @ViewBuilder footer: @escaping (Data.Index) -> Footer
    ) where Items: RandomAccessCollection, Items.Element: Identifiable, Data == Array<Items>
    {
        self.init(representable, sections: [items], content: content, header: header, footer: footer)
    }

    public init<
        Items: RandomAccessCollection
    >(
        _ items: Items,
        representable: Representable,
        @ViewBuilder content: @escaping (Items.Element) -> Content,
        @ViewBuilder header: @escaping (Data.Index) -> Header
    ) where Items: RandomAccessCollection, Items.Element: Identifiable, Data == Array<Items>, Footer == EmptyView
    {
        self.init(representable, items: items, content: content, header: header, footer: { _ in EmptyView() })
    }
}

@available(iOS 14.0, *)
extension CollectionViewAdapter {
    public init<
        Sections: RandomAccessCollection,
        ID: Hashable
    >(
        _ representable: Representable,
        sections: Sections,
        id: KeyPath<Sections.Element.Element, ID>,
        @ViewBuilder content: @escaping (Sections.Element.Element) -> Content,
        @ViewBuilder header: @escaping (Data.Index) -> Header,
        @ViewBuilder footer: @escaping (Data.Index) -> Footer
    ) where Sections.Element: RandomAccessCollection, Sections.Index: Hashable, Data == Array<Array<IdentifiableBox<Sections.Element.Element, ID>>>
    {
        let data: Data = sections.compactMap { items in
            items.compactMap { IdentifiableBox($0, id: id) }
        }
        self.init(representable, sections: data, content: { content($0.value) }, header: header, footer: footer)
    }

    public init<
        Sections: RandomAccessCollection,
        ID: Hashable
    >(
        _ representable: Representable,
        sections: Sections,
        id: KeyPath<Sections.Element.Element, ID>,
        @ViewBuilder content: @escaping (Sections.Element.Element) -> Content,
        @ViewBuilder header: @escaping (Data.Index) -> Header
    ) where Sections.Element: RandomAccessCollection, Sections.Index: Hashable, Data == Array<Array<IdentifiableBox<Sections.Element.Element, ID>>>, Footer == EmptyView
    {
        self.init(representable, sections: sections, id: id, content: content, header: header, footer: { _ in EmptyView() })
    }
}

@available(iOS 14.0, *)
extension CollectionViewAdapter where Header == EmptyView, Footer == EmptyView {
    public init<
        Items: RandomAccessCollection
    >(
        _ representable: Representable,
        items: Items,
        @ViewBuilder content: @escaping (Items.Element) -> Content
    ) where Items: RandomAccessCollection, Items.Element: Identifiable, Data == Array<Items>
    {
        self.init(representable, items: items, content: content, header: { _ in EmptyView() }, footer: { _ in EmptyView() })
    }

    public init<
        Items: RandomAccessCollection,
        ID: Hashable
    >(
        _ representable: Representable,
        items: Items,
        id: KeyPath<Items.Element, ID>,
        @ViewBuilder content: @escaping (Items.Element) -> Content
    ) where Items: RandomAccessCollection, Data == Array<Array<IdentifiableBox<Items.Element, ID>>>
    {
        self.init(representable, sections: [items], id: id, content: content, header: { _ in EmptyView() }, footer: { _ in EmptyView() })
    }
}

@available(iOS 14.0, *)
private struct CollectionViewBody<
    Header: View,
    Content: View,
    Footer: View,
    Representable: CollectionViewRepresentable,
    Data: RandomAccessCollection
>: UIViewRepresentable where
    Data.Element: RandomAccessCollection,
    Data.Index: Hashable,
    Data.Element.Element: Identifiable
{

    var representable: Representable
    var data: Data
    var header: (Data.Index) -> Header
    var content: (Data.Element.Element) -> Content
    var footer: (Data.Index) -> Footer

    func makeUIView(context: Context) -> Representable.UIViewType {
        var layoutOptions = CollectionViewLayoutOptions()
        if Header.self != EmptyView.self {
            layoutOptions.update(with: .header)
        }
        if Footer.self != EmptyView.self {
            layoutOptions.update(with: .footer)
        }
        let uiView = representable.makeUIView(
            context: CollectionViewRepresentableContext(
                environment: context.environment,
                transaction: context.transaction
            ),
            options: layoutOptions
        )
        if #available(iOS 16.0, *) {
            uiView.selfSizingInvalidation = .enabled
        }
        context.coordinator.bindDataSource(to: uiView)

        return uiView
    }

    func updateUIView(_ uiView: Representable.UIViewType, context: Context) {
        context.coordinator.update(body: self, uiView: uiView, transaction: context.transaction)
        representable.updateUIView(
            uiView,
            context: CollectionViewRepresentableContext(
                environment: context.environment,
                transaction: context.transaction
            )
        )
    }

    func makeCoordinator() -> Coordinator {
        Coordinator(body: self)
    }

    final class Coordinator: NSObject {

        private var body: CollectionViewBody
        private var dataSourceTransaction: Transaction?
        private var dataSource: UICollectionViewDiffableDataSource<Section, Item>!

        init(body: CollectionViewBody) {
            self.body = body
            super.init()
        }

        func bindDataSource(to uiView: UICollectionView) {
            let cellRegistration = UICollectionView.CellRegistration<
                UICollectionViewCell, Item
            > { [unowned self] cellView, indexPath, id in
                cellView.contentConfiguration = self.makeContent(
                    indexPath: indexPath
                )
            }
            let headerRegistration = UICollectionView.SupplementaryRegistration<UICollectionViewCell>(
                elementKind: UICollectionView.elementKindSectionHeader
            ) { [unowned self] headerView, _, indexPath in
                headerView.contentConfiguration = self.makeHeaderContent(
                    indexPath: indexPath
                )
            }
            let footerRegistration = UICollectionView.SupplementaryRegistration<UICollectionViewCell>(
                elementKind: UICollectionView.elementKindSectionFooter
            ) { [unowned self] footerView, _, indexPath in
                footerView.contentConfiguration = self.makeFooterContent(
                    indexPath: indexPath
                )
            }
            dataSource = UICollectionViewDiffableDataSource<Section, Item>(
                collectionView: uiView
            ) { (collectionView: UICollectionView, indexPath: IndexPath, item: Item) -> UICollectionViewCell? in
                let cell = collectionView.dequeueConfiguredReusableCell(
                    using: cellRegistration,
                    for: indexPath,
                    item: item
                )
                cell.automaticallyUpdatesContentConfiguration = false
                cell.contentView.transform3D = collectionView.transform3D
                cell.clipsToBounds = collectionView.clipsToBounds
                return cell
            }
            dataSource.supplementaryViewProvider = { (collectionView: UICollectionView, kind: String, indexPath: IndexPath) -> UICollectionReusableView? in
                switch kind {
                case UICollectionView.elementKindSectionHeader:
                    let headerView = collectionView.dequeueConfiguredReusableSupplementary(
                        using: headerRegistration,
                        for: indexPath
                    )
                    headerView.automaticallyUpdatesContentConfiguration = false
                    headerView.transform3D = collectionView.transform3D
                    headerView.clipsToBounds = collectionView.clipsToBounds
                    return headerView

                case UICollectionView.elementKindSectionFooter:
                    let footerView = collectionView.dequeueConfiguredReusableSupplementary(
                        using: footerRegistration,
                        for: indexPath
                    )
                    footerView.automaticallyUpdatesContentConfiguration = false
                    footerView.transform3D = collectionView.transform3D
                    footerView.clipsToBounds = collectionView.clipsToBounds
                    return footerView

                default:
                    return nil
                }
            }
        }

        func update(
            body newValue: CollectionViewBody,
            uiView: UICollectionView,
            transaction: Transaction
        ) {
            body = newValue

            let updated = updateDataSource(transaction: transaction)
            if transaction.isAnimated {
                UIView.animate(withDuration: 0.35, delay: 0, options: [.curveEaseInOut]) {
                    self.updateVisibleViews(uiView, updated: updated, transaction: transaction)
                }
            } else {
                if #available(iOS 16.0, *) {
                    uiView.selfSizingInvalidation = .disabled
                }
                updateVisibleViews(uiView, updated: updated, transaction: transaction)
                if #available(iOS 16.0, *) {
                    withCATransaction {
                        uiView.selfSizingInvalidation = .enabled
                    }
                }
            }
        }

        private func makeContent(
            indexPath: IndexPath,
            transaction: Transaction? = nil
        ) -> UIContentConfiguration {
            let section = body.data.index(body.data.startIndex, offsetBy: indexPath.section)
            let item = body.data[section].index(body.data[section].startIndex, offsetBy: indexPath.item)
            let value = body.data[section][item]
            return makeContent(
                value: value,
                transaction: transaction
            )
        }

        private func makeContent(
            value: Data.Element.Element,
            transaction: Transaction? = nil
        ) -> UIContentConfiguration {
            HostingConfiguration(
                id: value.id,
                transaction: transaction ?? dataSourceTransaction
            ) {
                body.content(value)
            }
        }

        private func makeHeaderContent(
            indexPath: IndexPath,
            transaction: Transaction? = nil
        ) -> UIContentConfiguration {
            let section = body.data.index(body.data.startIndex, offsetBy: indexPath.section)
            return HostingConfiguration(
                id: section,
                transaction: transaction ?? dataSourceTransaction
            ) {
                body.header(section)
            }
        }

        private func makeFooterContent(
            indexPath: IndexPath,
            transaction: Transaction? = nil
        ) -> UIContentConfiguration {
            let section = body.data.index(body.data.startIndex, offsetBy: indexPath.section)
            return HostingConfiguration(
                id: section,
                transaction: transaction ?? dataSourceTransaction
            ) {
                body.footer(section)
            }
        }

        private func updateDataSource(transaction: Transaction) -> Set<Item> {
            let oldValue = dataSource.snapshot().itemIdentifiers
            var updated = Set<Item>()

            var snapshot = NSDiffableDataSourceSnapshot<Section, Item>()
            snapshot.appendSections(Array(body.data.indices))
            for section in body.data.indices {
                let ids = body.data[section].map(\.id)
                snapshot.appendItems(ids, toSection: section)
                updated.formUnion(ids)
            }
            updated.formIntersection(oldValue)

            dataSourceTransaction = transaction
            dataSource.applySnapshot(snapshot, animated: transaction.isAnimated) {
                self.dataSourceTransaction = nil
            }
            return updated
        }

        private func updateVisibleViews(
            _ uiView: UICollectionView,
            updated: Set<Item>,
            transaction: Transaction
        ) {
            for indexPath in uiView.indexPathsForVisibleItems {
                if let cellView = uiView.cellForItem(at: indexPath) {
                    let section = body.data.index(body.data.startIndex, offsetBy: indexPath.section)
                    let item = body.data[section].index(body.data[section].startIndex, offsetBy: indexPath.item)
                    let value = body.data[section][item]
                    if updated.contains(value.id) {
                        cellView.contentConfiguration = self.makeContent(
                            value: value,
                            transaction: transaction
                        )
                    }
                }
            }
            for indexPath in uiView.indexPathsForVisibleSupplementaryElements(
                ofKind: UICollectionView.elementKindSectionHeader
            ) {
                let headerView = uiView.supplementaryView(
                    forElementKind: UICollectionView.elementKindSectionHeader,
                    at: indexPath
                ) as! UICollectionViewCell
                headerView.contentConfiguration = self.makeHeaderContent(
                    indexPath: indexPath,
                    transaction: transaction
                )
            }
            for indexPath in uiView.indexPathsForVisibleSupplementaryElements(
                ofKind: UICollectionView.elementKindSectionFooter
            ) {
                let footerView = uiView.supplementaryView(
                    forElementKind: UICollectionView.elementKindSectionFooter,
                    at: indexPath
                ) as! UICollectionViewCell
                footerView.contentConfiguration = self.makeFooterContent(
                    indexPath: indexPath,
                    transaction: transaction
                )
            }
        }
    }

    typealias Section = Data.Index
    typealias Item = Data.Element.Element.ID
}

extension UICollectionViewDiffableDataSource {
    func applySnapshot(
        _ snapshot: NSDiffableDataSourceSnapshot<SectionIdentifierType, ItemIdentifierType>,
        animated: Bool,
        completion: (() -> Void)? = nil
    ) {
        if #available(iOS 15.0, *) {
            apply(snapshot, animatingDifferences: animated, completion: completion)
        } else {
            if animated {
                apply(snapshot, animatingDifferences: true, completion: completion)
            } else {
                UIView.performWithoutAnimation {
                    self.apply(snapshot, animatingDifferences: true, completion: completion)
                }
            }
        }
    }
}

@available(iOS 14.0, *)
@available(macOS, unavailable)
@available(tvOS, unavailable)
@available(watchOS, unavailable)
public func HostingConfiguration<
    ID: Hashable,
    Content: View
>(
    id: ID?,
    transaction: Transaction?,
    @ViewBuilder content: () -> Content
) -> UIContentConfiguration {
    if #available(iOS 16.0, *) {
        return HostingConfiguration {
            content()
                .contentTransition(.identity)
                .transaction {
                    if let transaction {
                        $0 = transaction
                    }
                }
                .id(id)
        }
    } else {
        return HostingConfiguration {
            content()
                .transaction {
                    if let transaction {
                        $0 = transaction
                    }
                }
                .id(id)
        }
    }
}

@available(iOS 14.0, *)
@available(macOS, unavailable)
@available(tvOS, unavailable)
@available(watchOS, unavailable)
public func HostingConfiguration<Content: View>(
    @ViewBuilder content: () -> Content
) -> UIContentConfiguration {
    if #available(iOS 16.0, *) {
        return UIHostingConfiguration(content: content).margins(.all, 0)
    } else {
        return _UIHostingConfiguration(content: content)
    }
}

@available(iOS 14.0, *)
@available(macOS, unavailable)
@available(tvOS, unavailable)
@available(watchOS, unavailable)
private struct _UIHostingConfiguration<Content: View>: UIContentConfiguration {
    var content: Content

    init(@ViewBuilder content: () -> Content) {
        self.content = content()
    }

    func makeContentView() -> UIView & UIContentView {
        _UIHostingConfigurationContentView(configuration: self)
    }

    func updated(for state: UIConfigurationState) -> _UIHostingConfiguration<Content> {
        self
    }
}

@available(iOS 14.0, *)
@available(macOS, unavailable)
@available(tvOS, unavailable)
@available(watchOS, unavailable)
private class _UIHostingConfigurationContentView<Content: View>: HostingView<ModifiedContent<Content, SizeObserver>>, UIContentView {
    var configuration: UIContentConfiguration {
        didSet {
            update(for: configuration as! _UIHostingConfiguration<Content>)
        }
    }

    private var layoutInvalidationSize: CGSize?

    init(configuration: _UIHostingConfiguration<Content>) {
        self.configuration = configuration
        super.init(content: configuration.content.modifier(SizeObserver(onChange: { _ in })))
        content.modifier.onChange = { [unowned self] newValue in
            self.layoutInvalidationSize = newValue
        }
    }

    required init?(coder aDecoder: NSCoder) {
        fatalError("init(coder:) has not been implemented")
    }

    private func update(for configuration: _UIHostingConfiguration<Content>) {
        content.content = configuration.content
    }

    override func layoutSubviews() {
        super.layoutSubviews()
        if let layoutInvalidationSize {
            invalidateLayout(size: layoutInvalidationSize)
        }
    }

    private func invalidateLayout(size: CGSize) {
        layoutInvalidationSize = nil

        guard
            let collectionViewCell = superview as? UICollectionViewCell,
            let collectionView = collectionViewCell.superview as? UICollectionView,
            let indexPath = collectionView.indexPath(for: collectionViewCell)
        else {
            return
        }
        let ctx = UICollectionViewLayoutInvalidationContext()
        ctx.invalidateItems(at: [indexPath])
        collectionView.collectionViewLayout.invalidateLayout(with: ctx)
    }
}

@available(iOS 14.0, *)
@available(macOS, unavailable)
@available(tvOS, unavailable)
@available(watchOS, unavailable)
private struct SizeObserver: ViewModifier {
    var onChange: (CGSize) -> Void

    init(onChange: @escaping (CGSize) -> Void) {
        self.onChange = onChange
    }

    func body(content: Content) -> some View {
        content
            .background(
                GeometryReader { proxy in
                    Color.clear
                        .hidden()
                        .onChange(of: proxy.size, perform: onChange)
                }
            )
    }
}

#endif