import UIKit
import ExpandableDatePicker

class ViewController: UITableViewController, ExpandableDatePicker {
    // Not used directly by you, but is part of the protocol so the framework can use it.
    var edpIndexPath: IndexPath?

    // Whether or not the expansion should include a TimeZone row selector.
    var edpShowTimeZoneRow = true

    fileprivate var rowThatTogglesDatePicker: Int!

    fileprivate var selectedDate = Date()
    fileprivate var selectedTimeZone = TimeZone.autoupdatingCurrent

    override func viewDidLoad() {
        super.viewDidLoad()

        tableView.registerExpandableDatePicker()
    }
}

// MARK: - UITableViewDataSource
extension ViewController {
    override func tableView(_ tableView: UITableView, cellForRowAt indexPath: IndexPath) -> UITableViewCell {
        if edpShouldShowDatePicker(at: indexPath) {
            let cell = ExpandableDatePickerCell.reusableCell(for: indexPath, in: tableView)
            cell.onDateChanged = {
                [unowned self] date in
                self.selectedDate = date
                self.tableView.reloadRows(at: [self.edpLabelIndexPath!], with: .automatic)
            }

            cell.datePicker.date = selectedDate

            return cell
        } else if edpShouldShowTimeZoneRow(at: indexPath) {
            return ExpandableDatePickerTimeZoneCell.reusableCell(for: indexPath, in: tableView, timeZone: selectedTimeZone)
        }

        let modelIndexPath = edpUpdatedModelIndexPath(for: indexPath)

        if modelIndexPath.row == rowThatExpandsToDatePicker {
            let cell = ExpandableDatePickerSelectionCell.reusableCell(for: indexPath, in: tableView)
            cell.detailTextLabel!.text = DateFormatter.localizedString(from: selectedDate, dateStyle: .short, timeStyle: .none)

            return cell
        }

        let cell = tableView.dequeueReusableCell(withIdentifier: "normal", for: indexPath)
        cell.textLabel!.text = tableData[modelIndexPath.row]
        return cell
    }

    override func tableView(_ tableView: UITableView, numberOfRowsInSection section: Int) -> Int {
        return tableData.count + edpDatePickerRowsShowing
    }
}

// MARK: - UITableViewDelegate
extension ViewController {
    override func tableView(_ tableView: UITableView, didSelectRowAt indexPath: IndexPath) {
        guard let modelIndexPath = edpTableCellWasSelected(at: indexPath) else {
            // If tableCellWasSelected(at:) returns nil, they clicked on the time zone selector row.
            let vc = ExpandableDatePickerTimeZoneTableViewController {
                [unowned self] timeZone in
                self.selectedTimeZone = timeZone
                tableView.reloadRows(at: [indexPath], with: .automatic)
            }

            navigationController!.pushViewController(vc, animated: true)

            return
        }

        // modelIndexPath is the new indexPath you use for which row was selected.
    }
}

    override func viewDidLoad() {
        super.viewDidLoad()

        tableView.registerExpandableDatePicker()

        tableView.register(MyCoolTimeZoneSubclassCell.self, forCellReuseIdentifier: ExpandableDatePickerTimeZoneCell.identifier)
    }