typealias Grid = List<List<Char>>

private val up: Char = '^'
private val down: Char = 'v'
private val left: Char = '<'
private val right: Char = '>'
private val obstacle: Char = '#'
private val directions: Set<Char> = setOf(up, down, left, right)

data class Position(
    var direction: Char,
    var row: Int,
    var col: Int,
    val visited: MutableSet<Pair<Int, Int>> = mutableSetOf(),
    val history: MutableSet<Triple<Char, Int, Int>> = mutableSetOf(),
) {
    override fun toString(): String = "Position(direction: $direction, position: ($row,$col))"
}

fun part1(input: String): Int {
    val grid: Grid = input.lines().map(String::toList)
    val position = findStartPosition(grid)
    while (!isEndPosition(position, grid)) {
        moveOrTurn(position, grid)
    }
    return position.visited.size
}

fun part2(input: String): Int {
    var loops = 0
    for (i in input.indices) {
        if (input[i] != '.') {
            continue
        }
        val sb = StringBuilder(input)
        sb.setCharAt(i, obstacle)
        val grid: Grid = sb.toString().lines().map(String::toList)
        val position = findStartPosition(grid)
        while (!isEndPosition(position, grid)) {
            moveOrTurn(position, grid)
            if (isLoop(position)) {
                loops++
                break
            }
        }
    }
    return loops
}

private fun findStartPosition(grid: Grid): Position {
    for (row in grid.indices) {
        for (col in grid[row].indices) {
            val c = grid[row][col]
            if (directions.contains(c)) {
                val pos = Position(c, row, col)
                pos.visited.add(Pair(row, col))
                return pos
            }
        }
    }
    throw IllegalStateException("No start position found")
}

private fun isEndPosition(position: Position, grid: Grid): Boolean {
    return position.row == 0 || position.col == 0
            || position.row == grid.size - 1 || position.col == grid.size - 1
}

private fun isLoop(position: Position): Boolean {
    return position.history.contains(Triple(position.direction, position.row, position.col))
}

private fun moveOrTurn(position: Position, grid: Grid) {
    position.history.add(Triple(position.direction, position.row, position.col))

    when (position.direction) {
        up ->
            if (grid[position.row - 1][position.col] == obstacle)
                position.direction = right
            else
                position.row--

        right ->
            if (grid[position.row][position.col + 1] == obstacle)
                position.direction = down
            else
                position.col++

        down ->
            if (grid[position.row + 1][position.col] == obstacle)
                position.direction = left
            else
                position.row++

        left ->
            if (grid[position.row][position.col - 1] == obstacle)
                position.direction = up
            else
                position.col--

        else -> throw IllegalStateException("Invalid direction, cannot move")
    }

    position.visited.add(Pair(position.row, position.col))
}