package main

import (
	"fmt"
	"os"

	"git.annabunch.es/annabunches/adventofcode/2020/lib/util"
)

func occupied1(board [][]byte, i, j int) int {
	if i < 0 || j < 0 || i >= len(board) || j >= len(board[0]) {
		return 0
	}

	if board[i][j] == '#' {
		return 1
	}

	return 0
}

func occupied2(board [][]byte, i, j, iDir, jDir, level int) int {
	si := i + (iDir * level)
	sj := j + (jDir * level)

	if si < 0 || sj < 0 || si >= len(board) || sj >= len(board[0]) {
		return 0
	}

	if board[si][sj] == 'L' {
		return 0
	}

	if board[si][sj] == '#' {
		return 1
	}

	return occupied2(board, i, j, iDir, jDir, level+1)
}

func getNewValue(board [][]byte, i, j int, step string) byte {
	if board[i][j] == '.' {
		return '.'
	}

	occ := 0
	occLimit := 4

	switch step {
	case "1":
		occ = occupied1(board, i-1, j-1) +
			occupied1(board, i-1, j) +
			occupied1(board, i-1, j+1) +
			occupied1(board, i, j-1) +
			occupied1(board, i, j+1) +
			occupied1(board, i+1, j-1) +
			occupied1(board, i+1, j) +
			occupied1(board, i+1, j+1)
	case "2":
		occLimit = 5
		occ = occupied2(board, i, j, -1, -1, 1) +
			occupied2(board, i, j, -1, 0, 1) +
			occupied2(board, i, j, -1, +1, 1) +
			occupied2(board, i, j, 0, -1, 1) +
			occupied2(board, i, j, 0, +1, 1) +
			occupied2(board, i, j, 1, -1, 1) +
			occupied2(board, i, j, 1, 0, 1) +
			occupied2(board, i, j, 1, +1, 1)
	}

	if board[i][j] == 'L' && occ == 0 {
		return '#'
	}

	if board[i][j] == '#' && occ >= occLimit {
		return 'L'
	}

	return board[i][j]
}

func countOccupied(board [][]byte) int {
	count := 0

	for i := 0; i < len(board); i++ {
		for j := 0; j < len(board[0]); j++ {
			if board[i][j] == '#' {
				count++
			}
		}
	}

	return count
}

func main() {
	step := os.Args[1]
	board := util.InputParserBytes(os.Args[2])

	changed := true
	for changed {
		changed = false

		// initialize the next generation
		newBoard := make([][]byte, len(board))
		for i := 0; i < len(newBoard); i++ {
			newBoard[i] = make([]byte, len(board[0]))
			copy(newBoard[i], board[i])
		}

		for i := 0; i < len(board); i++ {
			for j := 0; j < len(board[0]); j++ {
				value := getNewValue(board, i, j, step)

				if board[i][j] != value {
					newBoard[i][j] = value
					changed = true
				}
			}
		}

		board = newBoard
	}

	fmt.Println(countOccupied(board))
}