public protocol CellDelegate: class {}

public protocol UpdatableCell: class {
    var delegate: CellDelegate? { get set }
}

public protocol FirstRespondableCell: class {
    var shouldBecomeFirstResponder: Bool { get set }
    func becomeFirstResponder()
}

enum NotificationID {
    case wawa
}

/// Protocol that needs to be implemented if the cell supports notification.
protocol NotificationableCell: UpdatableCell {

    var cellNotificationID: NotificationID { get set }
    var notificationToBeRemoved: NotificationID? { get set }
}

/// Protocol that needs to be implemented by the objects that is managing the adding and removal of cells. Example ViewModel
protocol NotificationRemovalObserver {

    func removeNotification(with id: NotificationID)

}

public class CellDelegateTableDataSource: BaseTableViewDataSource {

    var cellDelegate: CellDelegate

    init(cellDelegate: CellDelegate, tableDataSource: BaseTableSource) {
        self.cellDelegate = cellDelegate
        super.init(vm: tableDataSource)
    }

    override public func tableView(_ tableView: UITableView, cellForRowAt indexPath: IndexPath) -> UITableViewCell {
        let cell = super.tableView(tableView, cellForRowAt: indexPath)
        (cell as? UpdatableCell)?.delegate = cellDelegate
        if let selectionStyle = (cell as? UnselectableCell)?.unselectableStyle {
            cell.selectionStyle = selectionStyle
        }

        /// Removal of notification
        if let notificationableCell = cell as? NotificationableCell,
            let notificationRemovalObserver = notificationableCell.delegate as? NotificationRemovalObserver {
            if notificationableCell.cellNotificationID == notificationableCell.notificationToBeRemoved {
                notificationRemovalObserver.removeNotification(with: notificationableCell.cellNotificationID)
            }
        }


        return cell
    }
}