package main

import (
	"log"
	"math"
	"os"
	"strings"
)

func main() {
	inputBytes, _ := os.ReadFile("input.txt")
	input := strings.Split(string(inputBytes), "\n")

	log.Println("Puzzle 1:", puzzle1(input))
	log.Println("Puzzle 2:", puzzle2(input))
}

func puzzle1(input []string) int {
	// Read input into grid
	grid := readInput(input)
	totalFlashes := 0
	// 100 steps
	for step := 1; step <= 100; step++ {
		// Increase every energy level (every cell) by 1
		for i := range grid {
			for j := range grid[i] {
				grid[i][j]++
			}
		}
		// Flash cells
		for i := range grid {
			for j := range grid[i] {
				if grid[i][j] > 9 {
					totalFlashes += flash(grid, j, i)
				}
			}
		}
	}
	return totalFlashes
}

func puzzle2(input []string) int {
	// Read input into grid
	grid := readInput(input)
	totalFlashes := 0
forsteps:
	for step := 1; step < math.MaxInt; step++ {
		// Increase every energy level (every cell) by 1
		for i := range grid {
			for j := range grid[i] {
				grid[i][j]++
			}
		}
		// Flash cells
		for i := range grid {
			for j := range grid[i] {
				if grid[i][j] > 9 {
					totalFlashes += flash(grid, j, i)
				}
			}
		}
		// Check if all cells are now 0
		for i := range grid {
			for j := range grid[i] {
				if grid[j][i] != 0 {
					continue forsteps
				}
			}
		}
		// Return step
		return step
	}
	return totalFlashes
}

func readInput(input []string) [][]int {
	// Read input into 10x10 grid
	grid := make([][]int, 10)
	for i := range grid {
		grid[i] = make([]int, 10)
	}
	for i, line := range input {
		for j, c := range line {
			grid[i][j] = int(c - '0')
		}
	}
	return grid
}

func flash(grid [][]int, x, y int) int {
	grid[y][x] = 0
	flashes := 1
	// Iterate over all adjacent cells
	ax, ay := adjacent(x, y)
	for i := range ax {
		if grid[ay[i]][ax[i]] == 0 {
			// Already flashed
			continue
		}
		// Increase energy level by 1
		grid[ay[i]][ax[i]]++
		if grid[ay[i]][ax[i]] > 9 {
			// Flash cell aswell
			flashes += flash(grid, ax[i], ay[i])
		}
	}
	return flashes
}

func adjacent(x, y int) (xs, ys []int) {
	for i := 0; i < 10; i++ {
		for j := 0; j < 10; j++ {
			if i >= 0 && i <= 9 && i >= x-1 && i <= x+1 && j >= 0 && j <= 9 && j >= y-1 && j <= y+1 && !(i == x && j == y) {
				xs = append(xs, i)
				ys = append(ys, j)
			}
		}
	}
	return
}