package main

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

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

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

func do(input []string) (int, int) {
	// Parse input
	grid := make([][]int, len(input))
	for j, row := range input {
		gridRow := make([]int, len(row))
		for i, c := range row {
			gridRow[i] = int(c) - int('0')
		}
		grid[j] = gridRow
	}
	// For every grid cells check if its value is lower than its lowest neighbor
	// If it is, add the value to the risk + 1 and count the basin size
	risk := 0
	basins := []int{}
	for y, row := range grid {
		for x, cell := range row {
			if cell < lowestNeighbor(grid, x, y) {
				risk += cell + 1
				basins = append(basins, basinSize(grid, x, y))
			}
		}
	}
	sort.Ints(basins)
	return risk, basins[len(basins)-1] * basins[len(basins)-2] * basins[len(basins)-3]
}

func cellValue(grid [][]int, x, y int) int {
	if x < 0 || x >= len(grid[0]) || y < 0 || y >= len(grid) {
		// This cell is out of bounds
		return 9
	} else {
		return grid[y][x]
	}
}

func lowestNeighbor(grid [][]int, x, y int) int {
	lowest := math.MaxInt
	if n := cellValue(grid, x-1, y); n < lowest {
		lowest = n
	}
	if n := cellValue(grid, x+1, y); n < lowest {
		lowest = n
	}
	if n := cellValue(grid, x, y-1); n < lowest {
		lowest = n
	}
	if n := cellValue(grid, x, y+1); n < lowest {
		lowest = n
	}
	return lowest
}

func basinSize(grid [][]int, x, y int) int {
	if cellValue(grid, x, y) == 9 {
		return 0
	}
	// Mark cell as visited
	grid[y][x] = 9
	// Recurse on neighbors
	return 1 + basinSize(grid, x-1, y) + basinSize(grid, x+1, y) + basinSize(grid, x, y-1) + basinSize(grid, x, y+1)
}