package main

import (
	"fmt"

	h "git.bullercodeworks.com/brian/adventofcode/helpers"
)

func main() {
	inp := h.StdinToStringSlice()
	part1(inp)
	fmt.Println()
	part2(inp)
}

type node struct {
	left, right string
}

func build(input []string) (string, map[string]node) {
	inst := input[0]
	input = input[2:]
	m := make(map[string]node)
	for i := range input {
		var name, l, r string
		fmt.Sscanf(input[i], "%s = (%3s, %3s)", &name, &l, &r)
		m[name] = node{left: l, right: r}
	}
	return inst, m
}

func part1(input []string) {
	inst, m := build(input)
	curr := "AAA"
	instIdx := 0
	for curr != "ZZZ" {
		switch inst[instIdx%len(inst)] {
		case 'L':
			curr = m[curr].left
		case 'R':
			curr = m[curr].right
		}
		instIdx++
	}
	fmt.Println("# Part 1")
	fmt.Println(instIdx)
}

func part2(input []string) {
	inst, m := build(input)
	// Find all starting nodes
	var starts []string
	for k := range m {
		if k[2] == 'A' {
			starts = append(starts, k)
		}
	}
	var dists []int
	// Find the distances from 'starts' to the end nodes
	for _, v := range starts {
		cnt := 0
		distMap := make(map[string]int)
		// find the path to a 'Z'
		for {
			if v[2] == 'Z' {
				// Found it
				key := fmt.Sprintf("%s-%d", v, cnt%len(inst))
				if distMap[key] == 1 {
					break
				}
				distMap[key] = 1
				dists = append(dists, cnt)
			}
			switch inst[cnt%len(inst)] {
			case 'L':
				v = m[v].left
			case 'R':
				v = m[v].right
			}
			cnt++
		}
	}

	cnt := 1
	for _, n := range dists {
		cnt = h.Lcm(cnt, n)
	}

	fmt.Println("# Part 2")
	fmt.Println(cnt)
}