import Foundation

enum Direction : Hashable {
    case n, s, e, w
    func turnRight() -> Direction {
        let map: [Direction: Direction] = [.n: .e, .s: .w, .e: .s, .w: .n]
        return map[self]!
    }
    func asChar() -> Character {
        let map: [Direction: Character] = [.n: "^", .s: "v", .e: ">", .w: "<"]
        return map[self] ?? "?"
    }
}

struct GuardPos : Hashable {
    let x: Int
    let y: Int
    let dir: Direction
    func turnRight() -> GuardPos {
        return GuardPos(x: x, y: y, dir: dir.turnRight())
    }
    func move() -> GuardPos {
        switch dir {
        case Direction.n:
            return GuardPos(x: x-1, y: y, dir: dir)
        case Direction.s:
            return GuardPos(x: x+1, y: y, dir: dir)
        case Direction.e:
            return GuardPos(x: x, y: y+1, dir: dir)
        case Direction.w:
            return GuardPos(x: x, y: y-1, dir: dir)
        }
    }
}

struct Map {
    var obst: [[Bool]]
    var patrol: GuardPos
    var visited: Set<GuardPos>
    var done = false
    var looped = false

    func print() {
        for (i, row) in obst.enumerated() {
            Swift.print(String(row.enumerated().map { j, c in
                (i == patrol.x && j == patrol.y) ?
                    patrol.dir.asChar() : (c ? "#" : ".")
            }))
        }
    }

    mutating func step() {
        let next = patrol.move()
        if next.x < 0 || next.x >= obst.count
           || next.y < 0 || next.y >= obst[0].count {
            done = true
        } else if obst[next.x][next.y] {
            patrol = patrol.turnRight()
        } else if visited.contains(next) {
            looped = true
            done = true
        } else {
            visited.insert(next)
            patrol = next
        }
    }

    func addObst(x: Int, y: Int) -> Map? {
        if obst[x][y] {
            return nil
        }
        var newObst = obst
        newObst[x][y] = true
        return Map(
            obst: newObst, patrol: patrol,
            visited: visited, done: done, looped: looped
        )
    }
}

func readInput(_ filePath: String) throws -> Map {
    let content = try String(contentsOfFile: filePath, encoding: .ascii)
    let rows = content.split(separator: "\n").map(Array.init)
    let (patrolX, patrolY) = rows.map {
        $0.enumerated().compactMap { j, cell in cell == "^" ? j : nil }
    }.enumerated().compactMap { $1.count > 0 ? ($0, $1[0]) : nil }[0]
    let initPatrol = GuardPos(x: patrolX, y: patrolY, dir: Direction.n)
    return Map(
        obst: rows.map { $0.map { $0 == "#" } },
        patrol: initPatrol,
        visited: [initPatrol]
    )
}

func solve(_ map: Map) -> Int {
    var initMap = map
    while !initMap.done {
        initMap.step()
    }
    let visiteds = Set(initMap.visited.map {
        pos in GuardPos(x: pos.x, y: pos.y, dir: Direction.n)
    })
    print("\(visiteds.count) cells to try blocking")

    var count = 0
    for pos in visiteds {
        if var newMap = map.addObst(x: pos.x, y: pos.y) {
            while !newMap.done {
                newMap.step()
            }
            if newMap.looped {
                count += 1
                print("Found \(pos.x), \(pos.y)...")
            }
        }
    }
    return count
}

let map = try readInput(CommandLine.arguments[1])
print(solve(map))