package day10

import (
	"strings"
)

var directions = [][2]int{
	{-1, 0},
	{1, 0},
	{0, -1},
	{0, 1},
}

func Part1(input string) int {
	grid := ParseInput(input)
	return CalculateTotalTrailheadScore(grid)
}

func Part2(input string) int {
	grid := ParseInput(input)
	return CalculateTotalTrailheadRating(grid)
}

func ParseInput(input string) [][]int {
	lines := strings.Split(strings.TrimSpace(input), "\n")
	grid := make([][]int, len(lines))
	for i, line := range lines {
		grid[i] = make([]int, len(line))
		for j, char := range line {
			grid[i][j] = int(char - '0')
		}
	}
	return grid
}

func IsValidPosition(x, y, rows, cols int) bool {
	return x >= 0 && x < rows && y >= 0 && y < cols
}

func DFS(grid [][]int, x, y, prevHeight int, visited [][]bool, reached9 map[[2]int]bool) {
	rows := len(grid)
	cols := len(grid[0])
	if !IsValidPosition(x, y, rows, cols) || visited[x][y] || grid[x][y] != prevHeight+1 {
		return
	}

	if grid[x][y] == 9 {
		reached9[[2]int{x, y}] = true
		return
	}

	visited[x][y] = true
	for _, dir := range directions {
		nx, ny := x+dir[0], y+dir[1]
		DFS(grid, nx, ny, grid[x][y], visited, reached9)
	}
	visited[x][y] = false
}

func CalculateTrailheadScore(grid [][]int, x, y int) int {
	rows := len(grid)
	cols := len(grid[0])
	visited := make([][]bool, rows)
	for i := range visited {
		visited[i] = make([]bool, cols)
	}

	reached9 := make(map[[2]int]bool)
	DFS(grid, x, y, -1, visited, reached9)
	return len(reached9)
}

func DFSForRatings(grid [][]int, x, y, prevHeight int, visited [][]bool, trailCache map[[3]int]int) int {
	rows := len(grid)
	cols := len(grid[0])
	if !IsValidPosition(x, y, rows, cols) || visited[x][y] || grid[x][y] != prevHeight+1 {
		return 0
	}

	if grid[x][y] == 9 {
		return 1
	}

	cacheKey := [3]int{x, y, grid[x][y]}
	if count, exists := trailCache[cacheKey]; exists {
		return count
	}

	visited[x][y] = true
	totalTrails := 0
	for _, dir := range directions {
		nx, ny := x+dir[0], y+dir[1]
		totalTrails += DFSForRatings(grid, nx, ny, grid[x][y], visited, trailCache)
	}
	visited[x][y] = false

	trailCache[cacheKey] = totalTrails
	return totalTrails
}

func CalculateTrailheadRating(grid [][]int, x, y int) int {
	rows := len(grid)
	cols := len(grid[0])
	visited := make([][]bool, rows)
	for i := range visited {
		visited[i] = make([]bool, cols)
	}

	trailCache := make(map[[3]int]int)
	return DFSForRatings(grid, x, y, -1, visited, trailCache)
}

func CalculateTotalTrailheadScore(grid [][]int) int {
	rows := len(grid)
	cols := len(grid[0])
	totalScore := 0
	for i := 0; i < rows; i++ {
		for j := 0; j < cols; j++ {
			if grid[i][j] == 0 {
				score := CalculateTrailheadScore(grid, i, j)
				totalScore += score
			}
		}
	}
	return totalScore
}

func CalculateTotalTrailheadRating(grid [][]int) int {
	rows := len(grid)
	cols := len(grid[0])
	totalRating := 0
	for i := 0; i < rows; i++ {
		for j := 0; j < cols; j++ {
			if grid[i][j] == 0 {
				rating := CalculateTrailheadRating(grid, i, j)
				totalRating += rating
			}
		}
	}
	return totalRating
}