import java.io.File

typealias Coordinate = Pair<Int, Int>

fun main() {
    val input = File("input")

    var numbers: MutableMap<Coordinate, Int> = mutableMapOf()
    var symbols: MutableMap<Coordinate, Char> = mutableMapOf()

    input.useLines {
        it.forEachIndexed { row, line ->
            line.forEachIndexed { col, char ->
                when (char) {
                    in '0'..'9' -> {
                        // In this case, we're either starting a new
                        // number, or in the middle of an existing
                        // number.

                        // We need to mark all previous coordinates as
                        // well as this one as the full multi-digit
                        // number, so we need to check the previous
                        // columns.
                        var length = 1
                        while (numbers.get(Coordinate(row, col - length)) != null) {
                            length += 1
                        }

                        // Now we check if the column already has a
                        // number; if not, we're starting a new one,
                        // otherwise mark this column as well as any
                        // previous ones as the new number.
                        val number = numbers.get(Coordinate(row, col - length + 1))
                        if (number == null) {
                            numbers.put(Coordinate(row, col), char.digitToInt())
                        } else {
                            for (coordinate in col downTo col - length + 1) {
                                numbers.put(
                                        Coordinate(row, coordinate),
                                        number * 10 + char.digitToInt()
                                )
                            }
                        }
                    }
                    // Dots are ignored
                    '.' -> Unit
                    // Collect the coordinates of all symbols, too
                    else -> symbols.put(Coordinate(row, col), char)
                }
            }
        }
    }

    val gears: List<Int> =
            symbols
                    .map { symbol ->
                        if (symbol.value == '*') {
                            var adjacent = 0
                            var last: Int? = null
                            var ratio = 1

                            for (coord in nextTo(symbol.key)) {
                                val number = numbers.get(coord)

                                if (number != last && number != null) {
                                    last = number
                                    adjacent += 1
                                    ratio *= number
                                }
                            }

                            if (adjacent == 2) {
                                ratio
                            } else {
                                null
                            }
                        } else {
                            null
                        }
                    }
                    .filterNotNull()

    println("Part 1: ${gears.sum()}")

    val parts =
            symbols.flatMap { symbol ->
                nextTo(symbol.key)
                        .map {
                            val number = numbers.remove(it)

                            if (number == null) {
                                number
                            } else {
                                var prev: Coordinate = it
                                do {
                                    prev = Coordinate(prev.first, prev.second - 1)
                                } while (numbers.remove(prev) != null)

                                var next: Coordinate = it
                                do {
                                    next = Coordinate(next.first, next.second + 1)
                                } while (numbers.remove(next) != null)

                                number
                            }
                        }
                        .filterNotNull()
            }

    println("Part 2: ${parts.sum()}")
}

fun nextTo(coordinate: Coordinate): List<Coordinate> {
    return (coordinate.first - 1..coordinate.first + 1).flatMap { row ->
        (coordinate.second - 1..coordinate.second + 1).map { col -> Coordinate(row, col) }
    }
}