import java.io.File
import kotlin.sequences.generateSequence
import kotlin.text.Regex

enum class Instruction {
    LEFT,
    RIGHT
}

fun main() {
    val input = File("input")
    val (instructions, map) = parseInput(input)

    println(gcd(48, 18))

    println(part1(instructions, map))
    println(part2(instructions, map))
}

fun part1(instructions: List<Instruction>, map: Map<String, Pair<String, String>>): Long {
    val sequence = generateSequence { instructions }.flatten()

    return traverse(sequence, map, "AAA") { it == "ZZZ" }
}

// Theory: Since the instructions repeat, we can walk through all
// paths individually, and then find the least common multiple
fun part2(instructions: List<Instruction>, map: Map<String, Pair<String, String>>): Long {
    val traversals =
            map.keys.filter { it.endsWith('A') }.map {
                val sequence = generateSequence { instructions }.flatten()
                traverse(sequence, map, it) { it.endsWith('Z') }
            }

    return traversals.reduce { i, distance -> lcm(i, distance) }
}

fun lcm(a: Long, b: Long): Long {
    return a * (b / gcd(a, b))
}

fun gcd(a: Long, b: Long): Long {
    tailrec fun gcdStep(a: Long, b: Long): Long {
        val rem = a % b
        if (rem == 0L) {
            return b
        } else {
            return gcdStep(b, rem)
        }
    }

    if (a > b) {
        return gcdStep(a, b)
    } else {
        return gcdStep(b, a)
    }
}

fun traverse(
        instructions: Sequence<Instruction>,
        map: Map<String, Pair<String, String>>,
        start: String,
        isEnd: (node: String) -> Boolean
): Long {
    var length = 0
    var current = start

    for (instruction in instructions) {
        val next = map.get(current)
        if (next == null) {
            throw Exception("Invalid map")
        }

        when (instruction) {
            Instruction.LEFT -> current = next.first
            Instruction.RIGHT -> current = next.second
        }

        length += 1

        if (isEnd(current)) {
            break
        }
    }

    return length.toLong()
}

fun parseInput(input: File): Pair<List<Instruction>, Map<String, Pair<String, String>>> {
    return input.useLines {
        val lines = it.iterator()

        val instructions =
                lines.next().map {
                    when (it) {
                        'R' -> Instruction.RIGHT
                        'L' -> Instruction.LEFT
                        else -> throw Error("Invalid direction")
                    }
                }

        lines.next() // Skip separator between instructions and nodes

        val nodeRegex = Regex("([A-Z]+) = \\(([A-Z]+), ([A-Z]+)\\)")

        val nodes =
                lines.asSequence()
                        .map {
                            val matches = nodeRegex.matchEntire(it)

                            if (matches == null) {
                                throw Exception("Invalid node: ${it}")
                            }

                            val (node, edge1, edge2) = matches.destructured
                            node to (edge1 to edge2)
                        }
                        .toMap()

        Pair(instructions, nodes)
    }
}