 //
 //
 
-import Foundation
 import UIKit
+import BSGridCollectionViewLayout
 
 open class CustomPhotoDataSource<T: NSObject>: NSObject, UICollectionViewDataSource, UICollectionViewDelegate {
     
-    var selections = [T]()
-    var items: [T]!
-    var viewController: UIViewController
+    public var selections = [T]()
+    public var items: [T]!
+    public var changeSelections: (([T]) -> Void)? = nil
+    
+    var collectionView: UICollectionView!
+    
     var settings: Settings = Settings()
     let rowHandler: ((T) -> UIImage)
     
     fileprivate let photoCellIdentifier = "photoCellIdentifier"
-    //    fileprivate let imageContentMode: PHImageContentMode = .aspectFill
     
     let bundle: Bundle = Bundle(path: Bundle(for: PhotosViewController.self).path(forResource: "BSImagePicker", ofType: "bundle")!)!
     
-    public init(_ items: [T], viewController: UIViewController, rowHandler: @escaping ((T) -> UIImage), selections: [T]? = nil) {
+    public init(_ items: [T], collectionView: UICollectionView!, rowHandler: @escaping ((T) -> UIImage), selections: [T]? = nil) {
         self.items = items
-        self.viewController = viewController
-        //        self.settings = settings
+        self.collectionView = collectionView
         self.rowHandler = rowHandler
         if let selections = selections {
             self.selections = selections
         }
+        super.init()
+        self.collectionView.delegate = self
+        self.collectionView.dataSource = self
+        self.collectionView.collectionViewLayout = GridCollectionViewLayout()
     }
     
     public func numberOfSections(in collectionView: UICollectionView) -> Int {

            UIView.setAnimationsEnabled(true)
            
            cell.photoSelected = false
        } else {
            self.selections.append(item)
            
            if let selectionCharacter = self.settings.selectionCharacter {

            cell.photoSelected = true
        }
        
        if let changeSelectionsHandler = self.changeSelections {
            changeSelectionsHandler(self.selections)
        }
        
        return false
    }
    
    public func registerCellIdentifiersForCollectionView(_ collectionView: UICollectionView?) {
        collectionView?.register(UINib(nibName: "PhotoCell", bundle: self.bundle), forCellWithReuseIdentifier: photoCellIdentifier)
    }
    
    
}

// MARK: CustomPhotoDataSource proxy
extension CustomPhotoDataSource: BSImagePickerSettings {
    
    public var maxNumberOfSelections: Int {
        get {
            return settings.maxNumberOfSelections
        }
        set {
            settings.maxNumberOfSelections = newValue
        }
    }
    
    public var selectionCharacter: Character? {
        get {
            return settings.selectionCharacter
        }
        set {
            settings.selectionCharacter = newValue
        }
    }
    
    public var selectionFillColor: UIColor {
        get {
            return settings.selectionFillColor
        }
        set {
            settings.selectionFillColor = newValue
        }
    }
    
    public var selectionStrokeColor: UIColor {
        get {
            return settings.selectionStrokeColor
        }
        set {
            settings.selectionStrokeColor = newValue
        }
    }
    
    public var selectionShadowColor: UIColor {
        get {
            return settings.selectionShadowColor
        }
        set {
            settings.selectionShadowColor = newValue
        }
    }
    
    public var selectionTextAttributes: [String: AnyObject] {
        get {
            return settings.selectionTextAttributes
        }
        set {
            settings.selectionTextAttributes = newValue
        }
    }
    
    public var backgroundColor: UIColor {
        get {
            return settings.backgroundColor
        }
        set {
            settings.backgroundColor = newValue
        }
    }
    
    public var cellsPerRow: (_ verticalSize: UIUserInterfaceSizeClass, _ horizontalSize: UIUserInterfaceSizeClass) -> Int {
        get {
            return settings.cellsPerRow
        }
        set {
            settings.cellsPerRow = newValue
        }
    }
    
    public var takePhotos: Bool {
        get {
            return settings.takePhotos
        }
        set {
            settings.takePhotos = newValue
        }
    }
    
    public var takePhotoIcon: UIImage? {
        get {
            return settings.takePhotoIcon
        }
        set {
            settings.takePhotoIcon = newValue
        }
    }
    
}

