package dev.kaizen.app.chat

import dev.kaizen.app.net.StoredMessage

data class ChatNode(
    val id: String,
    val role: String,
    val parentId: String?,
    val childIds: List<String> = emptyList(),
)

data class BranchInfo(val index: Int, val total: Int)

data class VisibleEntry(
    val messageId: String,
    val branchInfo: BranchInfo?,
)

data class ChatTree(
    val nodes: Map<String, ChatNode>,
    val rootIds: List<String>,
    val activeRootId: String?,
    val activeChild: Map<String, String>,
) {
    companion object {
        val EMPTY = ChatTree(emptyMap(), emptyList(), null, emptyMap())
    }
}

fun buildChatTree(
    messages: List<StoredMessage>,
    savedActiveRootId: String?,
    savedActiveChild: Map<String, String>,
): ChatTree {
    val nodes = mutableMapOf<String, ChatNode>()
    val rootIds = mutableListOf<String>()

    for (msg in messages) {
        nodes[msg.id] = ChatNode(
            id = msg.id,
            role = msg.role,
            parentId = msg.parentId,
        )
    }

    for (msg in messages) {
        val pid = msg.parentId
        if (pid != null && nodes.containsKey(pid)) {
            val parent = nodes[pid]!!
            nodes[pid] = parent.copy(childIds = parent.childIds + msg.id)
        } else if (pid == null) {
            rootIds.add(msg.id)
        }
    }

    val activeRootId = if (savedActiveRootId != null && nodes.containsKey(savedActiveRootId))
        savedActiveRootId
    else rootIds.lastOrNull()

    return ChatTree(nodes, rootIds, activeRootId, savedActiveChild)
}

fun getVisiblePath(tree: ChatTree): List<VisibleEntry> {
    if (tree.rootIds.isEmpty()) return emptyList()
    val activeRootId = tree.activeRootId ?: tree.rootIds.last()
    val result = mutableListOf<VisibleEntry>()
    var currentId: String? = activeRootId

    while (currentId != null) {
        val node = tree.nodes[currentId] ?: break
        val siblings = if (node.parentId == null) {
            tree.rootIds
        } else {
            tree.nodes[node.parentId]?.childIds ?: listOf(node.id)
        }
        val branchInfo = if (siblings.size > 1) {
            BranchInfo(index = siblings.indexOf(node.id), total = siblings.size)
        } else null

        result.add(VisibleEntry(node.id, branchInfo))

        if (node.childIds.isEmpty()) break
        currentId = tree.activeChild[node.id] ?: node.childIds.last()
    }

    return result
}

fun treeNavigate(tree: ChatTree, nodeId: String, direction: Int): ChatTree {
    val node = tree.nodes[nodeId] ?: return tree
    val siblings = if (node.parentId == null) {
        tree.rootIds
    } else {
        tree.nodes[node.parentId]?.childIds ?: return tree
    }
    val idx = siblings.indexOf(nodeId)
    val newIdx = idx + direction
    if (newIdx < 0 || newIdx >= siblings.size) return tree
    val newActiveId = siblings[newIdx]

    return if (node.parentId == null) {
        tree.copy(activeRootId = newActiveId)
    } else {
        tree.copy(activeChild = tree.activeChild + (node.parentId to newActiveId))
    }
}
