package com.h0tk3y.jsonwalk

import kotlinx.serialization.encodeToString
import kotlinx.serialization.json.*
import java.io.File
import java.util.*

fun main(args: Array<String>) {
    val jsonFilePath = when {
        args.getOrNull(0) == "--file" -> args.getOrNull(1) ?: error("expected filename after --file")
        args.getOrNull(0) == "--new" -> null
        else -> error("expected args: --file <filename> or --new")
    }

    val rootElement = jsonFilePath?.let {
        MutableElement.fromJsonElement(Json.parseToJsonElement(File(it).readText()))
    } ?: MutableElement.ObjectElement(mutableMapOf())

    val state = State(rootElement)

    loop(state)
}

fun loop(state: State) {
    while (true) {
        val line = readLine()
        val lineWord = line?.substringBefore(" ") ?: return
        val command = when (lineWord) {
            "path" -> Command.Path
            "print" -> Command.Print
            "up" -> Command.Up
            "go" -> Command.Go(ElementPath.parse(line.substringAfter(" ")))
            "findpaths" -> Command.FindPaths(ElementPath.parse(line.substringAfter(" ")))
            "delete" -> Command.Delete(ElementPath.parse(line.substringAfter(" ")))
            "replace", "insert" -> run f@{
                val argsLine = line.substringAfter(" ")
                if (" " !in argsLine) return@f null
                val (pathString, elementString) = argsLine.split(" ", limit = 2)
                val path = ElementPath.parse(pathString)
                val createElement = { MutableElement.fromJsonElement(Json.parseToJsonElement(elementString)) }
                when (lineWord) {
                    "replace" -> Command.Replace(path, createElement)
                    "insert" -> Command.Insert(path, createElement)
                    else -> error("unexpected")
                }
            }
            "copy" -> Command.Copy(ElementPath.parse(line.substringAfter(" ")))
            "paste" -> Command.Paste
            "exit" -> return
            else -> null
        }
        val result = if (command != null) state.applyCommand(command) else CommandResult.Failure
        when (result) {
            CommandResult.Success -> println("OK")
            CommandResult.Failure -> println("fail")
            is CommandResult.ResultValue -> println("OK ${result.value}")
            is CommandResult.ResultValues -> {
                println("OK ${result.values.size}")
                result.values.forEach { println(it) }
            }
            is CommandResult.InvalidPath -> println("bad-path ${result.atIndex}")
        }
    }
}

class State(
    val rootElement: MutableElement
) {
    private var currentPath: ElementPath

    init {
        require(rootElement is MutableElement.ObjectElement || rootElement is MutableElement.ArrayElement) {
            "the root element must be an object element or an array element; $rootElement is passed"
        }

        currentPath = ElementPath(listOf())
    }

    fun applyCommand(command: Command): CommandResult {
        return when (command) {
            is Command.Path -> doPath()
            Command.Print -> doPrint()
            Command.Up -> doUp()
            is Command.Go -> doGo(command)
            is Command.FindPaths -> doFindPaths(command)
            is Command.Delete -> doDelete(command)
            is Command.Replace -> doReplace(command)
            is Command.Insert -> doInsert(command)
            is Command.Copy -> doCopy(command)
            is Command.Paste -> doPaste()
        }
    }

    private fun findInvalidPath(
        path: ElementPath,
        checkElement: (MutableElement?) -> Boolean = { it != null }
    ): CommandResult.InvalidPath {
        for (i in currentPath.segments.size until path.segments.size) {
            if (!checkElement(path.getElement(rootElement, i)))
                return CommandResult.InvalidPath(i - currentPath.segments.size)
        }
        return CommandResult.InvalidPath(path.segments.size - 1)
    }

    private fun objectElementOrArrayElement(element: MutableElement?): Boolean {
        return element is MutableElement.ObjectElement || element is MutableElement.ArrayElement
    }


    /** Always returns [CommandResult.ResultValue] containing the current path starting from the root element, for
     * example, `foo/bar/[0]/baz/[1]`. */
    private fun doPath(): CommandResult.ResultValue {
        return CommandResult.ResultValue(currentPath.toString())
    }

    /**
     * Returns [CommandResult.ResultValue] containing the string representation of the current element.
     */
    private fun doPrint(): CommandResult.ResultValue {
        /** Use [Json.encodeToString] */
        return CommandResult.ResultValue(Json.encodeToString(currentPath.getElement(rootElement)?.toJsonElement()))
    }

    /**
     * Go to the parent element, one level up. Returns [CommandResult.Failure] if no parent element exists,
     * [CommandResult.Success] otherwise.
     */
    private fun doUp(): CommandResult {
        if (currentPath.segments.isEmpty())
            return CommandResult.Failure
        currentPath = ElementPath(currentPath.segments.dropLast(1))
        return CommandResult.Success
    }

    /**
     * Changes the current path to an existing object or array element. If the relative path in [Command.Go.path]
     * points to a non-existing element or an element that is not an object or array, [CommandResult.InvalidPath] is
     * returned. Returns [CommandResult.Success] if the path is changed.
     */
    private fun doGo(go: Command.Go): CommandResult {

        val newPath = currentPath.plus(go.path)
        if (objectElementOrArrayElement(newPath.getElement(rootElement))) {
            currentPath = newPath
            return CommandResult.Success
        }

        return findInvalidPath(newPath) { element -> objectElementOrArrayElement(element) }
    }

    /**
     * Returns [CommandResult.ResultValues] with matching paths. A paths is considered matched if ends with the same
     * segments as the requested [Command.FindPaths.path].
     *
     * In the bonus task, the latter may contain `*` segments, which match any single segment.
     *
     * Example: in an object `{ "a": { "b": { "c": 123 } } }`, there are three paths, `a`, `a/b`, and `a/b/c`.
     * Only the second one matches a request for `a/b` (as `a/b/c` doesn't end with `a/b`).
     */

    private fun doFindPaths(command: Command.FindPaths): CommandResult.ResultValues {
        val paths: MutableList<ElementPath> = mutableListOf()
        val prefix: MutableList<String> = mutableListOf()

        fun findPaths(element: MutableElement): MutableList<ElementPath> {

            paths.add(ElementPath(prefix.toMutableList()))

            fun goToNextElement(string: String, element: MutableElement) {
                prefix.add(string)
                findPaths(element)
                prefix.removeLast()
            }

            when (element) {
                is MutableElement.ObjectElement -> {
                    element.items.forEach { goToNextElement(it.key, it.value) }
                }
                is MutableElement.ArrayElement -> {
                    element.items.forEachIndexed { index, it -> goToNextElement("[${index}]", it) }
                }
            }
            return paths
        }

        val currentElement = currentPath.getElement(rootElement)!!
        val resultPaths = findPaths(currentElement).map { it.toString() }.distinct()
        return CommandResult.ResultValues(resultPaths.filter { it.endsWith(command.path.toString()) })
    }


    /**
     * Deletes the element found by the relative path in [Command.Delete.path], removing the element from its parent.
     * If the parent is an array element, the successor elements move one index lower.
     *
     * In the bonus task, the path may contain star segments, so more than one element may be removed.
     *
     * If the path points to a non-existing element, [CommandResult.InvalidPath] is returned. [CommandResult.Success]
     * is returned otherwise.
     */
    private fun doDelete(command: Command.Delete): CommandResult {
        if (currentPath.plus(command.path).getElement(rootElement) != null) {
            val parent = ElementPath(currentPath.segments.plus(command.path.segments).dropLast(1))
                .getElement(rootElement)
            when (parent) {
                is MutableElement.ObjectElement -> parent.items.remove(command.path.segments.last())
                is MutableElement.ArrayElement ->
                    parent.items.remove(parent.items[command.path.segments.last().drop(1).dropLast(1).toInt()])
            }
            return CommandResult.Success
        }
        return findInvalidPath(currentPath.plus(command.path))
    }

    /**
     * Replaces the element by the specified [Command.Replace.path] with a new element provided by
     * [Command.Replace.newElement].
     *
     * In the bonus task, the path may contain `*` segments, so more than one element may be replaced.
     */
    private fun doReplace(replace: Command.Replace): CommandResult {
        val elementPath = currentPath.plus(replace.path)
        val toReplace = elementPath.getElement(rootElement)
        if (toReplace != null) {
            when (val parent = ElementPath(elementPath.segments.dropLast(1)).getElement(rootElement)) {
                is MutableElement.ArrayElement -> {
                    val index = replace.path.segments.last().substring(1, 2).toInt()
                    parent.items[index] = replace.newElement.invoke()
                    return CommandResult.Success
                }
                is MutableElement.ObjectElement -> {
                    val index = replace.path.segments.last()
                    parent.items[index] = replace.newElement.invoke()
                    return CommandResult.Success
                }
            }
        }
        return findInvalidPath(replace.path)
    }

    private fun String.isArrayIndex() = this[0] == '[' && this.last() == ']'
    private fun String.getArrayIndex() = this.drop(1).dropLast(1).toIntOrNull()

    /**
     * Inserts the element provided by [Command.Insert.newElement] at the designated relative [Command.Insert.path] from
     * the current position.
     *
     * If any intermediate elements in the path do not exist, they are created (they may be object elements as well as
     * array elements if the next segment is an index).
     *
     * If the insertion point position is already taken, [CommandResult.InvalidPath] is returned and no change is made.
     *
     * In an [MutableElement.ArrayElement], as special index segment [ElementPath.appendSegment] specifies that the new
     * element should be added after the last element.
     *
     * If an *intermediate* segment is located in an [MutableElement.ArrayElement] and the index is taken by another
     * element, no change should be made, and the result should be [CommandResult.InvalidPath]. The append segment or
     * the index after the last element are fine to use in intermediate array elements.
     *
     * If the *last* segment is an index in an [MutableElement.ArrayElement], then the new element should be added at
     * the specified index, moving the successor elements one index forward.
     */
    private fun doInsert(insert: Command.Insert): CommandResult {
        var currentElement = currentPath.getElement(rootElement) ?: return CommandResult.Failure
        var parentElement: MutableElement? = null
        var elementToAdd: MutableElement? = null
        var segmentToAdd: String? = null
        for (segmentIndex in insert.path.segments.indices) {
            val currentSegment = insert.path.segments[segmentIndex]
            when (currentElement) {
                is MutableElement.ArrayElement -> {
                    if (!currentSegment.isArrayIndex())
                        return CommandResult.InvalidPath(segmentIndex)
                    val indexToAdd = currentSegment.getArrayIndex() ?: currentElement.items.size
                    if (indexToAdd > currentElement.items.size)
                        return CommandResult.InvalidPath(segmentIndex)
                    when {
                        segmentIndex == insert.path.segments.size - 1 -> {
                            currentElement.items.add(indexToAdd, insert.newElement.invoke())
                        }
                        insert.path.segments[segmentIndex+1].isArrayIndex() -> {
                            if (indexToAdd < currentElement.items.size) {
                                if (currentElement.items[indexToAdd] !is MutableElement.ArrayElement)
                                    return CommandResult.InvalidPath(segmentIndex+1)
                                currentElement = currentElement.items[indexToAdd]
                            } else {
                                if (parentElement == null) {
                                    parentElement = currentElement
                                    elementToAdd = MutableElement.ArrayElement(mutableListOf())
                                    segmentToAdd = indexToAdd.toString()
                                    currentElement = elementToAdd
                                } else {
                                    currentElement.items.add(MutableElement.ArrayElement(mutableListOf()))
                                    currentElement = currentElement.items[indexToAdd]
                                }
                            }
                        }
                        else -> {
                            if (indexToAdd < currentElement.items.size) {
                                if (currentElement.items[indexToAdd] !is MutableElement.ObjectElement)
                                    return CommandResult.InvalidPath(segmentIndex+1)
                                currentElement = currentElement.items[indexToAdd]
                            } else {
                                if (parentElement == null) {
                                    parentElement = currentElement
                                    elementToAdd = MutableElement.ObjectElement(mutableMapOf())
                                    segmentToAdd = indexToAdd.toString()
                                    currentElement = elementToAdd
                                } else {
                                    currentElement.items.add(MutableElement.ObjectElement(mutableMapOf()))
                                    currentElement = currentElement.items[indexToAdd]
                                }
                            }
                        }
                    }
                }
                is MutableElement.ObjectElement -> {
                    if (currentSegment.isArrayIndex())
                        return CommandResult.InvalidPath(segmentIndex)
                    when {
                        segmentIndex == insert.path.segments.size - 1 -> {
                            currentElement.items[currentSegment] = insert.newElement.invoke()
                        }
                        insert.path.segments[segmentIndex+1].isArrayIndex() -> {
                            if (currentElement.items.containsKey(currentSegment)) {
                                if (currentElement.items[currentSegment] !is MutableElement.ArrayElement)
                                    return CommandResult.InvalidPath(segmentIndex + 1)
                                currentElement = currentElement.items[currentSegment]
                                    ?: MutableElement.ArrayElement(mutableListOf())
                            } else {
                                if (parentElement == null) {
                                    parentElement = currentElement
                                    elementToAdd = MutableElement.ArrayElement(mutableListOf())
                                    segmentToAdd = currentSegment
                                    currentElement = elementToAdd
                                } else {
                                    currentElement.items[currentSegment] = MutableElement.ArrayElement(mutableListOf())
                                    currentElement = currentElement.items[currentSegment]
                                        ?: MutableElement.ArrayElement(mutableListOf())
                                }
                            }
                        }
                        else -> {
                            if (currentElement.items.containsKey(currentSegment)) {
                                if (currentElement.items[currentSegment] !is MutableElement.ObjectElement)
                                    return CommandResult.InvalidPath(segmentIndex + 1)
                                currentElement = currentElement.items[currentSegment]
                                    ?: MutableElement.ArrayElement(mutableListOf())
                            } else {
                                if (parentElement == null) {
                                    parentElement = currentElement
                                    elementToAdd = MutableElement.ObjectElement(mutableMapOf())
                                    segmentToAdd = currentSegment
                                    currentElement = elementToAdd
                                } else {
                                    currentElement.items[currentSegment] = MutableElement.ObjectElement(mutableMapOf())
                                    currentElement = currentElement.items[currentSegment]
                                        ?: MutableElement.ArrayElement(mutableListOf())
                                }
                            }
                        }
                    }
                }
            }
        }
        if (parentElement != null && elementToAdd != null && segmentToAdd != null) {
            if (parentElement is MutableElement.ObjectElement)
                parentElement.items[segmentToAdd] = elementToAdd
            if (parentElement is MutableElement.ArrayElement)
                parentElement.items.add(segmentToAdd.toInt(), elementToAdd)
        }
        return CommandResult.Success
    }

    /**
     * Remembers the path matching [Command.Copy.path] as well as the element, for future insertion into a different
     * element in [doPaste]. In the bonus task, the requested path may match more than one element, all of which need
     * to be copied to the corresponding relative paths at the destination.
     * Assume that no mutating operation is performed between copy and paste.
     */

    private var rememberedElement: MutableElement? = null
    private var rememberedPath: ElementPath? = null

    private fun doCopy(command: Command.Copy): CommandResult {
        rememberedElement = currentPath.plus(command.path).getElement(rootElement)?.copy()
        rememberedPath = command.path
        if (rememberedElement != null && rememberedPath != null) {
            return CommandResult.Success
        }
        return findInvalidPath(currentPath.plus(command.path))
    }

    /**
     * Inserts the element remembered with [doCopy] at the corresponding relative path in the current position.
     */
    private fun doPaste(): CommandResult {
        val tmpPath = rememberedPath
        val tmpElement = rememberedElement
        if (tmpElement == null || tmpPath == null) {
            return CommandResult.Failure
        }
        val result = doInsert(Command.Insert(tmpPath) { tmpElement })
        if (result is CommandResult.InvalidPath) {
            return CommandResult.InvalidPath(result.atIndex)
        }
        return CommandResult.Success
    }
}

class ElementPath(val segments: List<String>) {
    override fun toString() = segments.joinToString(segmentSeparator)

    fun hasIndexAt(index: Int) = indexAtOrNull(index) != null
    fun indexAtOrNull(index: Int) = toIndexOrNull(segments[index])
    fun hasStarSegmentAt(index: Int) = isStarSegment(segments[index])
    fun hasAppendSegmentAt(index: Int) = isAppendSegment(segments[index])

    fun getElement(rootElement: MutableElement, index: Int = segments.size - 1): MutableElement? {
        var element: MutableElement? = rootElement
        for (i in 0..index) {
            if (element == null) return null
            element = element.getChild(segments[i])
        }
        return element
    }

    fun plus(path: ElementPath) = ElementPath(segments.plus(path.segments))

    companion object {
        const val segmentSeparator = "/"
        const val appendSegment = "[_]"
        const val starSegment = "*"

        fun toIndexOrNull(segment: String): Int? =
            segment.removeSurrounding("[", "]").takeIf { it != segment }?.toIntOrNull()

        fun isAppendSegment(segment: String) = segment == appendSegment
        fun isStarSegment(segment: String) = segment == starSegment

        fun parse(string: String) = ElementPath(string.split(segmentSeparator))
    }
}

sealed class CommandResult {
    object Success : CommandResult()
    object Failure : CommandResult()
    class ResultValue(val value: String) : CommandResult()
    class ResultValues(val values: List<String>) : CommandResult()
    class InvalidPath(val atIndex: Int) : CommandResult()
}

sealed class Command {
    object Path : Command()
    object Up : Command()
    object Print : Command()
    class Go(val path: ElementPath) : Command()
    class Replace(val path: ElementPath, val newElement: () -> MutableElement) : Command()
    class Insert(val path: ElementPath, val newElement: () -> MutableElement) : Command()
    class Copy(val path: ElementPath) : Command()
    object Paste : Command()
    class FindPaths(val path: ElementPath) : Command()
    class Delete(val path: ElementPath) : Command()
}

sealed class MutableElement {
    class PrimitiveNumber(var number: Number) : MutableElement()
    class PrimitiveString(var string: String) : MutableElement()
    class PrimitiveBoolean(var value: Boolean) : MutableElement()
    object PrimitiveNull : MutableElement()

    class ArrayElement(val items: MutableList<MutableElement>) : MutableElement()
    class ObjectElement(val items: MutableMap<String, MutableElement>) : MutableElement()

    fun toJsonElement(): JsonElement = when (this) {
        is PrimitiveNumber -> JsonPrimitive(number)
        is PrimitiveString -> JsonPrimitive(string)
        is PrimitiveBoolean -> JsonPrimitive(value)
        PrimitiveNull -> JsonNull
        is ArrayElement -> JsonArray(items.map { it.toJsonElement() })
        is ObjectElement -> JsonObject(items.mapValues { it.value.toJsonElement() })
    }

    fun copy() : MutableElement? {
        when (this) {
            is PrimitiveBoolean -> return PrimitiveBoolean(value)
            PrimitiveNull -> return null
            is PrimitiveNumber -> return PrimitiveNumber(number)
            is PrimitiveString -> return PrimitiveString(string)

            is ArrayElement -> return ArrayElement(items.toMutableList())
            is ObjectElement -> return ObjectElement(items.toMutableMap())
        }
    }

    fun getChild(child: String): MutableElement? =
        when (this) {
            is ArrayElement -> {
                val index = child.drop(1).dropLast(1).toString().toIntOrNull()
                if (index == null)
                    null
                else
                    items.getOrNull(index)
            }
            is ObjectElement -> items[child]
            else -> null
        }

    companion object {
        fun fromJsonElement(json: JsonElement): MutableElement = when (json) {
            is JsonPrimitive -> when {
                json.isString -> PrimitiveString(json.content)
                json is JsonNull -> PrimitiveNull
                json.content == "true" || json.content == "false" -> PrimitiveBoolean(json.content.toBoolean())
                json.content.toIntOrNull() != null -> PrimitiveNumber(json.content.toInt())
                json.content.toDoubleOrNull() != null -> PrimitiveNumber(json.content.toDouble())
                else -> PrimitiveString(json.content)
            }
            is JsonArray -> ArrayElement(json.mapTo(mutableListOf()) { fromJsonElement(it) })
            is JsonObject -> ObjectElement(json.keys.associateTo(mutableMapOf()) { key ->
                key to fromJsonElement(json.getValue(key))
            })
        }
    }
}