Tusker/Tusker/Screens/Conversation/ConversationTableViewController.swift

//
//  ConversationTableViewController.swift
//  Tusker
//
//  Created by Shadowfacts on 8/28/18.
//  Copyright © 2018 Shadowfacts. All rights reserved.
//

import UIKit
import Pachyderm
import CoreData

class ConversationNode {
    let status: StatusMO
    var children: [ConversationNode]
    
    init(status: StatusMO) {
        self.status = status
        self.children = []
    }
}

class ConversationTableViewController: EnhancedTableViewController {
    
    static let showPostsImage = UIImage(systemName: "eye.fill")!
    static let hidePostsImage = UIImage(systemName: "eye.slash.fill")!
    
    static let bottomSeparatorTag = 101
    
    weak var mastodonController: MastodonController!
    
    let mainStatusID: String
    let mainStatusState: StatusState
    private(set) var dataSource: UITableViewDiffableDataSource<Section, Item>!
    
    var showStatusesAutomatically = false
    var visibilityBarButtonItem: UIBarButtonItem!
    
    init(for mainStatusID: String, state: StatusState = .unknown, mastodonController: MastodonController) {
        self.mainStatusID = mainStatusID
        self.mainStatusState = state
        self.mastodonController = mastodonController
        
        super.init(style: .plain)
        
        dragEnabled = true
    }
    
    required init?(coder aDecoder: NSCoder) {
        fatalError("init(coder:) has not been implemented")
    }
    
    deinit {
        guard let persistentContainer = mastodonController?.persistentContainer else { return }
        let snapshot = dataSource.snapshot()
        for case let .status(id: id, state: _) in snapshot.itemIdentifiers {
            persistentContainer.status(for: id)?.decrementReferenceCount()
        }
    }
    
    override func viewDidLoad() {
        super.viewDidLoad()
        
        title = NSLocalizedString("Conversation", comment: "conversation screen title")

        tableView.delegate = self
        tableView.dataSource = self
        
        tableView.register(UINib(nibName: "TimelineStatusTableViewCell", bundle: nil), forCellReuseIdentifier: "statusCell")
        tableView.register(UINib(nibName: "ConversationMainStatusTableViewCell", bundle: nil), forCellReuseIdentifier: "mainStatusCell")
        tableView.register(UINib(nibName: "ExpandThreadTableViewCell", bundle: .main), forCellReuseIdentifier: "expandThreadCell")
        
        tableView.prefetchDataSource = self
        
        tableView.backgroundColor = .secondarySystemBackground
        // separators are disabled on the table view so we can re-add them ourselves
        // so they're not inserted in between statuses in the ame sub-thread
        tableView.separatorStyle = .none
        
        dataSource = UITableViewDiffableDataSource<Section, Item>(tableView: tableView, cellProvider: { (tableView, indexPath, item) -> UITableViewCell? in
            switch item {
            case let .status(id: id, state: state):
                let rowsInSection = self.dataSource.tableView(tableView, numberOfRowsInSection: indexPath.section)
                let firstInSection = indexPath.row == 0
                let lastInSection = indexPath.row == rowsInSection - 1
                
                let identifier = id == self.mainStatusID ? "mainStatusCell" : "statusCell"
                let cell = tableView.dequeueReusableCell(withIdentifier: identifier, for: indexPath) as! BaseStatusTableViewCell
                
                if id == self.mainStatusID {
                    cell.selectionStyle = .none
                }
                
                cell.delegate = self
                cell.showStatusAutomatically = self.showStatusesAutomatically
                
                if let cell = cell as? TimelineStatusTableViewCell {
                    cell.showReplyIndicator = false
                }
                
                cell.updateUI(statusID: id, state: state)
                
                cell.setShowThreadLinks(prev: !firstInSection, next: !lastInSection)
                if lastInSection {
                    if cell.viewWithTag(ConversationTableViewController.bottomSeparatorTag) == nil {
                        let separator = UIView()
                        separator.tag = ConversationTableViewController.bottomSeparatorTag
                        separator.translatesAutoresizingMaskIntoConstraints = false
                        separator.backgroundColor = tableView.separatorColor
                        cell.addSubview(separator)
                        NSLayoutConstraint.activate([
                            separator.heightAnchor.constraint(equalToConstant: 0.5),
                            separator.bottomAnchor.constraint(equalTo: cell.bottomAnchor),
                            separator.leftAnchor.constraint(equalTo: cell.leftAnchor, constant: cell.separatorInset.left),
                            separator.rightAnchor.constraint(equalTo: cell.rightAnchor),
                        ])
                    }
                } else {
                    cell.viewWithTag(ConversationTableViewController.bottomSeparatorTag)?.removeFromSuperview()
                }
                
                return cell
                
            case let .expandThread(childThreads: childThreads):
                let cell = tableView.dequeueReusableCell(withIdentifier: "expandThreadCell", for: indexPath) as! ExpandThreadTableViewCell
                cell.updateUI(childThreads: childThreads)
                return cell
            }
        })
        
        let initialImage = showStatusesAutomatically ? ConversationTableViewController.hidePostsImage : ConversationTableViewController.showPostsImage
        visibilityBarButtonItem = UIBarButtonItem(image: initialImage, style: .plain, target: self, action: #selector(toggleVisibilityButtonPressed))
        navigationItem.rightBarButtonItem = visibilityBarButtonItem
        
        loadMainStatus()
    }
    
    private func loadMainStatus() {
        if let mainStatus = mastodonController.persistentContainer.status(for: mainStatusID) {
            self.mainStatusLoaded(mainStatus)
        } else {
            let request = Client.getStatus(id: mainStatusID)
            mastodonController.run(request) { (response) in
                switch response {
                case let .success(status, _):
                    let viewContext = self.mastodonController.persistentContainer.viewContext
                    self.mastodonController.persistentContainer.addOrUpdate(status: status, incrementReferenceCount: false, context: viewContext) { (statusMO) in
                        self.mainStatusLoaded(statusMO)
                    }
                    
                case .failure(_):
                    fatalError()
                }
            }
        }
    }
    
    private func mainStatusLoaded(_ mainStatus: StatusMO) {
        let mainStatusItem = Item.status(id: mainStatusID, state: mainStatusState)
        
        var snapshot = NSDiffableDataSourceSnapshot<Section, Item>()
        snapshot.appendSections([.statuses])
        snapshot.appendItems([mainStatusItem], toSection: .statuses)
        dataSource.apply(snapshot, animatingDifferences: false)

        let mainStatusInReplyToID = mainStatus.inReplyToID
        mainStatus.incrementReferenceCount()

        // todo: it would be nice to cache these contexts
        let request = Status.getContext(mainStatusID)
        mastodonController.run(request) { response in
            guard case let .success(context, _) = response else { fatalError() }

            let parents = self.getDirectParents(inReplyTo: mainStatusInReplyToID, from: context.ancestors)
            let parentStatuses = context.ancestors.filter { parents.contains($0.id) }
            
            // todo: should this really be blindly adding all the descendants?
            self.mastodonController.persistentContainer.addAll(statuses: parentStatuses + context.descendants) {
                DispatchQueue.main.async {
                    var snapshot = self.dataSource.snapshot()
                    snapshot.insertItems(parents.map { .status(id: $0, state: .unknown) }, beforeItem: mainStatusItem)
                    
                    // fetch all descendant status managed objects
                    let descendantIDs = context.descendants.map(\.id)
                    let request: NSFetchRequest<StatusMO> = StatusMO.fetchRequest()
                    request.predicate = NSPredicate(format: "id in %@", descendantIDs)
                    
                    if let descendants = try? self.mastodonController.persistentContainer.viewContext.fetch(request) {
                        // convert array of descendant statuses into tree of sub-threads
                        let childThreads = self.getDescendantThreads(mainStatus: mainStatus, descendants: descendants)
                        
                        // convert sub-threads into items for section and add to snapshot
                        self.addFlattenedChildThreadsToSnapshot(childThreads, mainStatus: mainStatus, snapshot: &snapshot)
                    }
                    
                    self.dataSource.apply(snapshot, animatingDifferences: false) {
                        // ensure that the main status is on-screen after newly loaded statuses are added
                        // todo: should this not happen if the user has already started scrolling (e.g. because the main status is very long)?
                        if let indexPath = self.dataSource.indexPath(for: mainStatusItem) {
                            self.tableView.scrollToRow(at: indexPath, at: .middle, animated: false)
                        }
                    }
                }
            }
        }
    }

    private func getDirectParents(inReplyTo inReplyToID: String?, from statuses: [Status]) -> [String] {
        var statuses = statuses
        var parents = [String]()
        
        var parentID: String? = inReplyToID

        while parentID != nil, let parentIndex = statuses.firstIndex(where: { $0.id == parentID }) {
            let parentStatus = statuses.remove(at: parentIndex)
            parents.insert(parentStatus.id, at: 0)
            parentID = parentStatus.inReplyToID
        }
        
        return parents
    }
    
    private func getDescendantThreads(mainStatus: StatusMO, descendants: [StatusMO]) -> [ConversationNode] {
        var descendants = descendants
        
        func removeAllInReplyTo(id: String) -> [StatusMO] {
            let statuses = descendants.filter { $0.inReplyToID == id }
            descendants.removeAll { $0.inReplyToID == id }
            return statuses
        }
        
        var nodes: [String: ConversationNode] = [
            mainStatus.id: ConversationNode(status: mainStatus)
        ]
        
        var idsToCheck = [mainStatusID]
        
        while !idsToCheck.isEmpty {
            let inReplyToID = idsToCheck.removeFirst()
            let nodeForID = nodes[inReplyToID]!
            
            let inReply = removeAllInReplyTo(id: inReplyToID)
            for reply in inReply {
                idsToCheck.append(reply.id)
                
                let replyNode = ConversationNode(status: reply)
                nodes[reply.id] = replyNode
                
                nodeForID.children.append(replyNode)
            }
        }
        
        return nodes[mainStatusID]!.children
    }
    
    private func addFlattenedChildThreadsToSnapshot(_ childThreads: [ConversationNode], mainStatus: StatusMO, snapshot: inout NSDiffableDataSourceSnapshot<Section, Item>) {
        var childThreads = childThreads
        
        // child threads by the same author as the main status come first
        let pivotIndex = childThreads.partition(by: { $0.status.account.id != mainStatus.account.id })
        
        // within each group, child threads are sorted chronologically
        childThreads[0..<pivotIndex].sort(by: { $0.status.createdAt < $1.status.createdAt })
        childThreads[pivotIndex...].sort(by: { $0.status.createdAt < $1.status.createdAt })
        
        for node in childThreads {
            snapshot.appendSections([.childThread(firstStatusID: node.status.id)])
            snapshot.appendItems([.status(id: node.status.id, state: .unknown)])
            
            var currentNode = node
            while true {
                let next: ConversationNode
                
                if currentNode.children.count == 0 {
                    break
                } else if currentNode.children.count == 1 {
                    next = currentNode.children[0]
                } else {
                    let sameAuthorStatuses = currentNode.children.filter({ $0.status.account.id == node.status.account.id })
                    if sameAuthorStatuses.count == 1 {
                        next = sameAuthorStatuses[0]
                    } else {
                        snapshot.appendItems([.expandThread(childThreads: currentNode.children)])
                        break
                    }
                }
                
                currentNode = next
                snapshot.appendItems([.status(id: next.status.id, state: .unknown)])
            }
        }
    }
    
    func item(for indexPath: IndexPath) -> (id: String, state: StatusState)? {
        return self.dataSource.itemIdentifier(for: indexPath).flatMap { (item) in
            switch item {
            case let .status(id: id, state: state):
                return (id: id, state: state)
            default:
                return nil
            }
        }
    }
    
    // MARK: - Table view delegate
    
    override func tableView(_ tableView: UITableView, didSelectRowAt indexPath: IndexPath) {
        if case .expandThread = dataSource.itemIdentifier(for: indexPath),
           case let .status(id: id, state: state) = dataSource.itemIdentifier(for: IndexPath(row: indexPath.row - 1, section: indexPath.section)) {
            self.selected(status: id, state: state)
        } else {
            super.tableView(tableView, didSelectRowAt: indexPath)
        }
    }
    
    override func tableView(_ tableView: UITableView, canEditRowAt indexPath: IndexPath) -> Bool {
        return true
    }
    
    override func tableView(_ tableView: UITableView, leadingSwipeActionsConfigurationForRowAt indexPath: IndexPath) -> UISwipeActionsConfiguration? {
        return (tableView.cellForRow(at: indexPath) as? TableViewSwipeActionProvider)?.leadingSwipeActionsConfiguration()
    }
    
    override func tableView(_ tableView: UITableView, trailingSwipeActionsConfigurationForRowAt indexPath: IndexPath) -> UISwipeActionsConfiguration? {
        return (tableView.cellForRow(at: indexPath) as? TableViewSwipeActionProvider)?.trailingSwipeActionsConfiguration()
    }
    
    @objc func toggleVisibilityButtonPressed() {
        #if SDK_IOS_15
        if #available(iOS 15.0, *) {
            visibilityBarButtonItem.isSelected = !visibilityBarButtonItem.isSelected
            showStatusesAutomatically = visibilityBarButtonItem.isSelected
        } else {
            showStatusesAutomatically = !showStatusesAutomatically
        }
        #else
        showStatusesAutomatically = !showStatusesAutomatically
        #endif
        
        let snapshot = dataSource.snapshot()
        for case let .status(id: _, state: state) in snapshot.itemIdentifiers where state.collapsible == true {
            state.collapsed = !showStatusesAutomatically
        }
        
        for cell in tableView.visibleCells {
            guard let cell = cell as? BaseStatusTableViewCell,
                cell.collapsible else { continue }
            cell.showStatusAutomatically = showStatusesAutomatically
            cell.setCollapsed(!showStatusesAutomatically, animated: false)
        }
        
        // recalculate cell heights
        tableView.beginUpdates()
        tableView.endUpdates()

        if showStatusesAutomatically {
            visibilityBarButtonItem.image = ConversationTableViewController.hidePostsImage
        } else {
            visibilityBarButtonItem.image = ConversationTableViewController.showPostsImage
        }
    }

}

extension ConversationTableViewController {
    enum Section: Hashable {
        case statuses
        case childThread(firstStatusID: String)
    }
    enum Item: Hashable {
        case status(id: String, state: StatusState)
        case expandThread(childThreads: [ConversationNode])
        
        static func == (lhs: Item, rhs: Item) -> Bool {
            switch (lhs, rhs) {
            case let (.status(id: a, state: _), .status(id: b, state: _)):
                return a == b
            case let (.expandThread(childThreads: a), .expandThread(childThreads: b)):
                return zip(a, b).allSatisfy { $0.status.id == $1.status.id }
            default:
                return false
            }
        }
        
        func hash(into hasher: inout Hasher) {
            switch self {
            case let .status(id: id, state: _):
                hasher.combine("status")
                hasher.combine(id)
            case let .expandThread(childThreads: children):
                hasher.combine("expandThread")
                hasher.combine(children.map(\.status.id))
            }
        }
    }
}

extension ConversationTableViewController: TuskerNavigationDelegate {
    func conversation(mainStatusID: String, state: StatusState) -> ConversationTableViewController {
        let vc = ConversationTableViewController(for: mainStatusID, state: state, mastodonController: mastodonController)
        // transfer show statuses automatically state when showing new conversation
        vc.showStatusesAutomatically = self.showStatusesAutomatically
        return vc
    }
}

extension ConversationTableViewController: StatusTableViewCellDelegate {
    var apiController: MastodonController { mastodonController }
    func statusCellCollapsedStateChanged(_ cell: BaseStatusTableViewCell) {
        // causes the table view to recalculate the cell heights
        tableView.beginUpdates()
        tableView.endUpdates()
    }
}

extension ConversationTableViewController: UITableViewDataSourcePrefetching, StatusTablePrefetching {
    func tableView(_ tableView: UITableView, prefetchRowsAt indexPaths: [IndexPath]) {
        let ids: [String] = indexPaths.compactMap { item(for: $0)?.id }
        prefetchStatuses(with: ids)
    }
    
    func tableView(_ tableView: UITableView, cancelPrefetchingForRowsAt indexPaths: [IndexPath]) {
        let ids: [String] = indexPaths.compactMap { item(for: $0)?.id }
        cancelPrefetchingStatuses(with: ids)
    }
}