Tusker/Tusker/Screens/Conversation/ConversationTree.swift

//
//  ConversationTree.swift
//  Tusker
//
//  Created by Shadowfacts on 2/4/23.
//  Copyright © 2023 Shadowfacts. All rights reserved.
//

import Foundation
import Pachyderm

@MainActor
class ConversationNode {
    let statusID: String
    let status: StatusMO
    var children: [ConversationNode]
    
    init(status: StatusMO) {
        self.statusID = status.id
        self.status = status
        self.children = []
    }
}

@MainActor
struct ConversationTree {
    let ancestors: [ConversationNode]
    let mainStatus: ConversationNode
    var descendants: [ConversationNode] {
        mainStatus.children
    }
    
    init(ancestors: [ConversationNode], mainStatus: ConversationNode) {
        self.ancestors = ancestors
        self.mainStatus = mainStatus
    }
    
    static func build(for mainStatus: StatusMO, ancestors: [StatusMO], descendants: [StatusMO]) -> ConversationTree {
        let mainStatusNode = ConversationNode(status: mainStatus)
        let ancestors = buildAncestorNodes(mainStatusNode: mainStatusNode, ancestors: ancestors)
        buildDescendantNodes(mainStatusNode: mainStatusNode, descendants: descendants)
        return ConversationTree(ancestors: ancestors, mainStatus: mainStatusNode)
    }
    
    private static func buildAncestorNodes(mainStatusNode: ConversationNode, ancestors: [StatusMO]) -> [ConversationNode] {
        var statuses = ancestors
        var parents = [ConversationNode]()
        
        var parentID: String? = mainStatusNode.status.inReplyToID
        
        while let currentParentID = parentID,
              let parentIndex = statuses.firstIndex(where: { $0.id == currentParentID }) {
            let parentStatus = statuses.remove(at: parentIndex)
            
            let node = ConversationNode(status: parentStatus)
            parents.insert(node, at: 0)
            
            parentID = parentStatus.inReplyToID
        }
        
        // once the parents list is built and in-order, then we walk through and set each node's children
        for (index, node) in parents.enumerated() {
            if index == parents.count - 1 {
                // the last parent is the direct parent of the main status
                node.children = [mainStatusNode]
            } else {
                // otherwise, it's the parent of the status that comes immediately after it in the parents list
                node.children = [parents[index + 1]]
            }
        }
        
        return parents
    }
    
    // doesn't return anything, since we're modifying the main status node in-place
    private static func buildDescendantNodes(mainStatusNode: ConversationNode, descendants: [StatusMO]) {
        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] = [
            mainStatusNode.status.id: mainStatusNode
        ]
        
        var idsToCheck = [mainStatusNode.status.id]
        
        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)
            }
        }
    }
}