package main

import (
	"bufio"
	"fmt"
	"os"
	"strings"
)

func main() {
	fmt.Println(find(readData("example.txt")))
	fmt.Println(36, 81)
	fmt.Println(find(readData("data.txt")))
}

type Candidate struct {
	X, Y int
	Alt  byte
	Head *TrailHead
}

func (c Candidate) Identity() int {
	return (c.X << 16) | c.Y
}

type TrailHead struct {
	Reaches map[int]bool
}

func findNextCandidates(x, y int, targetAlt byte, head *TrailHead, topo [][]byte) (result []Candidate) {
	if x > 0 && topo[y][x-1] == targetAlt {
		result = append(result, Candidate{x - 1, y, targetAlt, head})
	}
	if x < len(topo[y])-1 && topo[y][x+1] == targetAlt {
		result = append(result, Candidate{x + 1, y, targetAlt, head})
	}
	if y > 0 && topo[y-1][x] == targetAlt {
		result = append(result, Candidate{x, y - 1, targetAlt, head})
	}
	if y < len(topo)-1 && topo[y+1][x] == targetAlt {
		result = append(result, Candidate{x, y + 1, targetAlt, head})
	}
	return
}

func find(topo [][]byte) (resultPart1, resultPart2 int) {
	candidates := make([]Candidate, 0)
	for y, line := range topo {
		for x, alt := range line {
			if alt == 0 {
				trailhead := &TrailHead{make(map[int]bool)}
				candidate := Candidate{x, y, alt, trailhead}
				candidates = append(candidates, candidate)
			}
		}
	}

	for len(candidates) > 0 {
		current := candidates[len(candidates)-1]
		candidates = candidates[:len(candidates)-1]

		nextCandidates := findNextCandidates(current.X, current.Y, current.Alt+1, current.Head, topo)
		if current.Alt == 8 {
			for _, summit := range nextCandidates {
				id := summit.Identity()
				_, alreadyIn := summit.Head.Reaches[id]
				if !alreadyIn {
					summit.Head.Reaches[id] = true
					resultPart1++
				}
			}
			resultPart2 += len(nextCandidates)
		} else {
			candidates = append(candidates, nextCandidates...)
		}
	}

	return
}

func readData(fileName string) (topo [][]byte) {
	fp, err := os.Open(fileName)
	if err != nil {
		panic(err)
	}

	scanner := bufio.NewScanner(fp)
	for scanner.Scan() {
		line := []byte(strings.TrimSpace(scanner.Text()))
		if len(line) == 0 {
			continue
		}

		for i := range line {
			line[i] = line[i] - '0'
		}

		topo = append(topo, line)
	}

	return
}