enum class Team {
    RED,
    BLACK,
    EMPTY;

    override fun toString(): String = when (this) {
        RED -> "R"
        BLACK -> "B"
        EMPTY -> "O"
    }
}

class Board() {

    // Kotlin's way of doing constants
    companion object {
    	const val ROWS = 6;
    	const val COLS = 7;
    }

    // Init array for board
    var board = Array(7*6) { Team.EMPTY }

    // Overload operator so [row,col] works
    operator fun get(row: Int, col: Int): Team? {
        if (row < 0 || row > ROWS || col < 0 || col > COLS) {
            return null
        } else {
    		return board[row * COLS + col]
        }
    }

    // Overload operator so [row,col] = value works
    operator fun set(row: Int, col: Int, value: Team) {
		board[row * COLS + col] = value
    }

    // Overload toString to output pretty format
    override fun toString(): String {
        var boardStr = ""
        board.forEachIndexed { idx, space -> boardStr += "${space}" + if ((idx + 1) % 7 == 0) "\n" else " " }
        return boardStr
    }
}

fun findWin(board: Board): Team? {
    // Ranges are inclusive, so minus 1
    for (row in 0..Board.ROWS - 1) {
        for (col in 0..Board.COLS - 1) {
            val team = board[row, col]
            // If the space under test is empty, move on
            if (team == Team.EMPTY) {
                continue
            }

            if (team == board[row, col + 1] && team == board[row, col + 2] && team == board[row, col + 3]) {
                println("Horizontal Winner: ${team}")
                return team;
            }

            if (team == board[row + 1, col] && team == board[row + 2, col] && team == board[row + 3, col]) {
                println("Vertical Winner: ${team}")
                return team;
            }

            if ((team == board[row + 1, col + 1] && team == board[row + 2, col + 2] && team == board[row + 3, col + 3]) ||
                team == board[row + 1, col - 1] && team == board[row + 2, col - 2] && team == board[row + 3, col - 3]) {
                println("Diagonal Winner: ${team}")
                return team;
            }
        }
    }

    return null
}

fun main(args: Array<String>) {
    val board = Board()

    board[2, 6] = Team.RED
    board[3, 5] = Team.RED
    board[4, 4] = Team.RED
    board[5, 3] = Team.RED

    println(board)

    findWin(board)
}