package dev.austinzhu.algods.containers.util

import dev.austinzhu.algods.containers.tree.AbstractBinaryTree
import dev.austinzhu.algods.containers.tree.Tree
import kotlin.math.max

interface TreePrinter<N : Tree.Node<*, *, N>> {
    override fun toString(): String

    fun MutableList<TreeLine>.getWidth(): Int {
        val topLine = this.firstOrNull() ?: return 0
        return topLine.lIndent + topLine.line.length + topLine.rIndent
    }

    fun MutableList<TreeLine>.halfWidths(): Pair<Int, Int> {
        val topLine = this.firstOrNull() ?: return 0 to 0
        val leftWidth = topLine.lIndent + (topLine.line.length) / 2
        val rightWidth = topLine.rIndent + (topLine.line.length - 1) / 2
        return leftWidth to rightWidth
    }

    fun TreeLine.halfBodyWidths(): Pair<Int, Int> {
        val leftWidth = (this.line.length) / 2
        val rightWidth = (this.line.length - 1) / 2
        return leftWidth to rightWidth
    }

    fun buildRootLine(
        root: String,
        fst: MutableList<TreeLine>,
        leftWidth: Int,
        snd: MutableList<TreeLine>,
        rightWidth: Int
    ): TreeLine {
        val rootLineFst = (root.length) / 2
        val rootLineSnd = (root.length - 1) / 2
        val lIndent: Int
        val rIndent: Int
        // if the root line take more space than the merged tree lines
        if (rootLineFst - leftWidth <= 0) {
            lIndent = leftWidth - rootLineFst
        } else {
            lIndent = 0
            val excessIndent = rootLineFst - leftWidth
            fst.forEach { it.lIndent += excessIndent }
        }
        if (rootLineSnd - rightWidth <= 0) {
            rIndent = rightWidth - rootLineSnd
        } else {
            rIndent = 0
            val excessIndent = rootLineSnd - rightWidth
            snd.forEach { it.rIndent += excessIndent }
        }
        return TreeLine(lIndent, StringBuilder(root), rIndent)
    }

}

data class TreeLine(var lIndent: Int, val line: StringBuilder, var rIndent: Int) {
    override fun toString(): String {
        return StringBuilder().append(" ".repeat(lIndent)).append(line).append(" ".repeat(rIndent)).toString()
    }
}

// Number of spaces that separates the nodes
const val SEP_LENGTH = 2

class NAryTreePrinter<N : Tree.Node<*, *, N>>(private val root: N?) : TreePrinter<N> {
    private fun addRootLine(treeLines: MutableList<TreeLine>, root: String) {
        val branchLine = treeLines.firstOrNull()
        // modify the branch line
        if (branchLine != null) {
            val mid = branchLine.line.length / 2
            when (branchLine.line[mid]) {
                '─' -> branchLine.line[mid] = '┴'
                '┬' -> branchLine.line[mid] = '┼'
                else -> branchLine.line[mid] = '│'
            }
        }

        val (leftWidth, rightWidth) = treeLines.halfWidths()
        val rootLine = buildRootLine(root, treeLines, leftWidth, treeLines, rightWidth)
        treeLines.add(0, rootLine)
    }

    private fun buildSingleBranchLine(line: TreeLine, isFirst: Boolean, isLast: Boolean): TreeLine {
        val branchLine = TreeLine(0, StringBuilder(), 0)

        var lIndent = line.lIndent
        var rIndent = line.rIndent
        val (fstHalf, sndHalf) = line.halfBodyWidths()
        if (isFirst) {
            lIndent += fstHalf
            rIndent += sndHalf
            branchLine.line.append('┌')
        } else if (isLast) {
            lIndent += fstHalf
            rIndent += sndHalf
            branchLine.line.append('┐')
        } else {
            branchLine.line.append("─".repeat(fstHalf)).append('┬').append("─".repeat(sndHalf))
        }
        return TreeLine(lIndent, branchLine.line, rIndent)
    }

    private fun mergeLines(fst: MutableList<TreeLine>, snd: MutableList<TreeLine>): MutableList<TreeLine> {
        val result = mutableListOf<TreeLine>()

        val fstWidth = fst.getWidth()
        val sndWidth = snd.getWidth()

        for (i in 0 until max(fst.size, snd.size)) {
            val line1 = fst.getOrNull(i)
            val line2 = snd.getOrNull(i)
            if (line1 != null && line2 != null) {
                val lineBody = StringBuilder(line1.line)
                // if it is the branch line
                if (i == 0) {
                    lineBody.append("─".repeat(line1.rIndent + SEP_LENGTH + line2.lIndent))
                } else {
                    // add space between the lines for separation
                    lineBody.append(" ".repeat(line1.rIndent + SEP_LENGTH + line2.lIndent))
                }
                val both = TreeLine(
                    line1.lIndent,
                    lineBody.append(line2.line),
                    line2.rIndent
                )
                result.add(both)
            } else if (line1 != null) {
                // add indent to compensate for the separation space
                result.add(TreeLine(line1.lIndent, line1.line, line1.rIndent + sndWidth + SEP_LENGTH))
            } else if (line2 != null) {
                result.add(TreeLine(line2.lIndent + fstWidth + SEP_LENGTH, line2.line, line2.rIndent))
            }
        }
        return result
    }

    private fun mergeAll(childTreeLines: MutableList<MutableList<TreeLine>>): MutableList<TreeLine> {
        if (childTreeLines.size == 0) {
            return mutableListOf()
        }
        if (childTreeLines.size == 1) {
            return childTreeLines.first()
        }
        val result = mutableListOf<TreeLine>()
        val mid = childTreeLines.size / 2
        val left = mergeAll(childTreeLines.subList(0, mid))
        val right = mergeAll(childTreeLines.subList(mid, childTreeLines.size))

        val treeLines = mergeLines(left, right)
        result.addAll(treeLines)
        return result
    }

    private fun printHelper(node: N?): MutableList<TreeLine> {
        if (node == null) {
            return mutableListOf()
        }

        val children = node.getChildren().filterNotNull()
        val childrenLines = mutableListOf<MutableList<TreeLine>>()
        for (i in children.indices) {
            val childLines = printHelper(children[i])
            val branchLine = buildSingleBranchLine(childLines[0], i == 0, i == children.size - 1)
            childLines.add(0, branchLine)
            childrenLines.add(childLines)
        }

        val result = mergeAll(childrenLines)
        addRootLine(result, node.toString())
        return result
    }

    override fun toString(): String {
        val treeLines = printHelper(root)
        val sb = StringBuilder()
        treeLines.forEach {
            sb.append(it).append("\n")
        }
        return sb.toString()
    }
}

class BinaryTreePrinter<N : AbstractBinaryTree<*, *, N>.Node>(private val root: N?) : TreePrinter<N> {
    private fun mergeSubTreeLine(left: MutableList<TreeLine>, right: MutableList<TreeLine>): MutableList<TreeLine> {
        val result = mutableListOf<TreeLine>()
        val width1 = left.getWidth()
        val width2 = right.getWidth()

        for (i in 0 until max(left.size, right.size)) {
            val line1 = left.getOrNull(i)
            val line2 = right.getOrNull(i)
            if (line1 != null && line2 != null) {
                val both = TreeLine(
                    line1.lIndent,
                    line1.line.append(" ".repeat(line1.rIndent + 1 + line2.lIndent)).append(line2.line),
                    line2.rIndent
                )
                result.add(both)
            } else if (line1 != null) {
                result.add(TreeLine(line1.lIndent, line1.line, line1.rIndent + width2 + 1))
            } else if (line2 != null) {
                result.add(TreeLine(line2.lIndent + width1 + 1, line2.line, line2.rIndent))
            }
        }
        return result
    }

    private fun buildBranchLine(leftRootLine: TreeLine?, rightRootLine: TreeLine?): TreeLine {
        if (leftRootLine == null && rightRootLine == null) {
            return TreeLine(0, StringBuilder(), 0)
        }
        val branch = StringBuilder()
        val lIndent: Int
        val rIndent: Int

        val rootChar: Char = if (leftRootLine != null && rightRootLine != null) {
            '┴'
        } else if (leftRootLine != null) {
            '┘'
        } else {
            '└'
        }

        if (leftRootLine != null) {
            val (fstHalfT1, sndHalfT1) = leftRootLine.halfBodyWidths()
            lIndent = leftRootLine.lIndent + fstHalfT1
            branch.append('┌').append("─".repeat(sndHalfT1 + leftRootLine.rIndent))
        } else {
            lIndent = 0
        }
        branch.append(rootChar)
        if (rightRootLine != null) {
            val (fstHalfT2, sndHalfT2) = rightRootLine.halfBodyWidths()
            rIndent = rightRootLine.rIndent + sndHalfT2
            branch.append("─".repeat(rightRootLine.lIndent + fstHalfT2)).append('┐')
        } else {
            rIndent = 0
        }

        return TreeLine(lIndent, branch, rIndent)
    }

    private fun printHelper(node: N?): MutableList<TreeLine> {
        if (node == null) {
            return mutableListOf()
        }

        val treeLines = mutableListOf<TreeLine>()

        val left = printHelper(node.left)
        val right = printHelper(node.right)

        if (left.isEmpty() && right.isEmpty()) {
            return mutableListOf(TreeLine(0, StringBuilder(node.toString()), 0))
        }

        val rootLine = buildRootLine(node.toString(), left, left.getWidth(), right, right.getWidth())
        treeLines.add(rootLine)
        val branchLine = buildBranchLine(left.firstOrNull(), right.firstOrNull())
        treeLines.add(branchLine)
        val subTreeLines = mergeSubTreeLine(left, right)
        treeLines.addAll(subTreeLines)

        return treeLines
    }

    override fun toString(): String {
        val treeLines = printHelper(root)
        val sb = StringBuilder()
        treeLines.forEach {
            sb.append(it).append("\n")
        }
        return sb.toString()
    }
}
