package day21

import (
	"fmt"

	"gitea.paas.celticinfo.fr/oabrivard/aoc2023/utils"
)

type Coord struct {
	row int
	col int
}

func BFS1(matrix [][]byte, root Coord, maxSteps int) int {
	height := len(matrix)
	width := len(matrix[0])
	queue := utils.NewQueue[*Coord]()

	matrix[root.row][root.col] = 'O'
	queue.Enqueue(&root)

	moves := 0
	for moves < maxSteps {
		//fmt.Println(moves + 1)
		newQueue := utils.NewQueue[*Coord]()

		for queue.HasElement() {
			coord := queue.Dequeue()
			matrix[coord.row][coord.col] = '.'

			if coord.row > 0 && matrix[coord.row-1][coord.col] != '#' {
				newCoord := Coord{coord.row - 1, coord.col}
				newQueue.Enqueue(&newCoord)
				matrix[newCoord.row][newCoord.col] = 'O'
			}

			if coord.row < height-1 && matrix[coord.row+1][coord.col] != '#' {
				newCoord := Coord{coord.row + 1, coord.col}
				newQueue.Enqueue(&newCoord)
				matrix[newCoord.row][newCoord.col] = 'O'
			}

			if coord.col > 0 && matrix[coord.row][coord.col-1] != '#' {
				newCoord := Coord{coord.row, coord.col - 1}
				newQueue.Enqueue(&newCoord)
				matrix[newCoord.row][newCoord.col] = 'O'
			}

			if coord.col < width-1 && matrix[coord.row][coord.col+1] != '#' {
				newCoord := Coord{coord.row, coord.col + 1}
				newQueue.Enqueue(&newCoord)
				matrix[newCoord.row][newCoord.col] = 'O'
			}
		}

		queue = newQueue
		moves++
	}

	// count 'O' in matrix
	count := 0
	for r := range matrix {
		for _, b := range matrix[r] {
			if b == 'O' {
				count++
			}
		}
	}

	return count
}

func redrawFromCache(matrix [][]byte, cacheQueue utils.Queue[*Coord]) {
	for _, coord := range cacheQueue.GetData() {
		if coord != nil {
			matrix[coord.row][coord.col] = 'O'
		}
	}
}

func BFS2(matrix [][]byte, root Coord, maxSteps int) int {
	height := len(matrix)
	width := len(matrix[0])
	queue := utils.NewQueue[*Coord]()

	cache := make([][]utils.Queue[*Coord], height)
	for r, row := range matrix {
		cache[r] = make([]utils.Queue[*Coord], width)
		for c := range row {
			cache[r][c] = utils.NewQueue[*Coord]()
		}
	}

	matrix[root.row][root.col] = 'O'
	queue.Enqueue(&root)

	moves := 0
	for moves < maxSteps {
		fmt.Println(moves + 1)
		newQueue := utils.NewQueue[*Coord]()

		for queue.HasElement() {
			coord := queue.Dequeue()

			if cache[coord.row][coord.col].HasElement() {

				matrix[coord.row][coord.col] = '.'
				redrawFromCache(matrix, cache[coord.row][coord.col])

			} else {
				cacheQueue := utils.NewQueue[*Coord]()
				matrix[coord.row][coord.col] = '.'

				if coord.row > 0 && matrix[coord.row-1][coord.col] != '#' {
					newCoord := Coord{coord.row - 1, coord.col}
					newQueue.Enqueue(&newCoord)
					matrix[newCoord.row][newCoord.col] = 'O'
					cacheQueue.Enqueue(&newCoord)
				}

				if coord.row < height-1 && matrix[coord.row+1][coord.col] != '#' {
					newCoord := Coord{coord.row + 1, coord.col}
					newQueue.Enqueue(&newCoord)
					matrix[newCoord.row][newCoord.col] = 'O'
					cacheQueue.Enqueue(&newCoord)
				}

				if coord.col > 0 && matrix[coord.row][coord.col-1] != '#' {
					newCoord := Coord{coord.row, coord.col - 1}
					newQueue.Enqueue(&newCoord)
					matrix[newCoord.row][newCoord.col] = 'O'
					cacheQueue.Enqueue(&newCoord)
				}

				if coord.col < width-1 && matrix[coord.row][coord.col+1] != '#' {
					newCoord := Coord{coord.row, coord.col + 1}
					newQueue.Enqueue(&newCoord)
					matrix[newCoord.row][newCoord.col] = 'O'
					cacheQueue.Enqueue(&newCoord)
				}

				cache[coord.row][coord.col] = cacheQueue
			}
		}

		queue = newQueue
		moves++
	}

	// count 'O' in matrix
	count := 0
	for r := range matrix {
		for _, b := range matrix[r] {
			if b == 'O' {
				count++
			}
		}
	}

	return count
}

func BFS3(matrix [][]byte, root Coord, maxSteps int) int {
	height := len(matrix)
	width := len(matrix[0])
	queue := utils.NewQueue[*Coord]()

	matrix[root.row][root.col] = 'O'
	queue.Enqueue(&root)

	moves := 0
	for moves < maxSteps {
		//fmt.Println(moves + 1)
		newQueue := utils.NewQueue[*Coord]()
		coordInQueue := map[Coord]bool{}

		for queue.HasElement() {
			coord := queue.Dequeue()
			matrix[coord.row][coord.col] = '.'

			if coord.row > 0 && matrix[coord.row-1][coord.col] != '#' {
				newCoord := Coord{coord.row - 1, coord.col}
				if !coordInQueue[newCoord] {
					newQueue.Enqueue(&newCoord)
					matrix[newCoord.row][newCoord.col] = 'O'
					coordInQueue[newCoord] = true
				}
			}

			if coord.row < height-1 && matrix[coord.row+1][coord.col] != '#' {
				newCoord := Coord{coord.row + 1, coord.col}
				if !coordInQueue[newCoord] {
					newQueue.Enqueue(&newCoord)
					matrix[newCoord.row][newCoord.col] = 'O'
					coordInQueue[newCoord] = true
				}
			}

			if coord.col > 0 && matrix[coord.row][coord.col-1] != '#' {
				newCoord := Coord{coord.row, coord.col - 1}
				if !coordInQueue[newCoord] {
					newQueue.Enqueue(&newCoord)
					matrix[newCoord.row][newCoord.col] = 'O'
					coordInQueue[newCoord] = true
				}
			}

			if coord.col < width-1 && matrix[coord.row][coord.col+1] != '#' {
				newCoord := Coord{coord.row, coord.col + 1}
				if !coordInQueue[newCoord] {
					newQueue.Enqueue(&newCoord)
					matrix[newCoord.row][newCoord.col] = 'O'
					coordInQueue[newCoord] = true
				}
			}
		}

		queue = newQueue
		moves++
	}

	// count 'O' in matrix
	count := 0
	for r := range matrix {
		for _, b := range matrix[r] {
			if b == 'O' {
				count++
			}
		}
	}

	return count
}

func Part1(lines []string, maxSteps int) int {
	matrix := make([][]byte, len(lines))

	var root Coord
	for r, line := range lines {
		matrix[r] = make([]byte, len(line))
		for c, b := range line {
			matrix[r][c] = byte(b)
			if b == 'S' {
				root.row = r
				root.col = c
			}
		}
	}

	result := BFS3(matrix, root, maxSteps)

	return result
}