package day08

import (
	 _ "fmt"
	_ "adventofcode2023/utils"
	"regexp"
	"strings"
	"math"
)

type Node struct {
	left string
	right string
}

type Network map[string]Node

func Part1(input string) int {
	instructions, network := parseInput(input)
	return getSteps("AAA", instructions, network, true)
}

func Part2(input string) int {
	instructions, nexts, network := parseInput2(input)
	ans := 1
	for _, next := range nexts {
		ans = lcm(ans, getSteps(next, instructions, network, false))
	}
	return ans
}

// gcd calculates the greatest common divisor using Euclid's algorithm
func gcd(a, b int) int {
	for b != 0 {
		a, b = b, a%b
	}
	return a
}

// lcm calculates the least common multiple using the formula: LCM(a, b) = |a * b| / GCD(a, b)
func lcm(a, b int) int {
	if a == 0 || b == 0 {
		return 0
	}
	return int(math.Abs(float64(a*b)) / float64(gcd(a, b)))
}

func parseInput(input string) ([]rune, Network) {
	network := make(Network)
	lines := strings.Split(input, "\n")
	instructions := []rune(lines[0])
	re := regexp.MustCompile(`(\w+) = \((\w+), (\w+)\)`)
	
	for _, line := range lines[2:] {
		matches := re.FindAllStringSubmatch(line, -1)
		network[matches[0][1]] = Node{matches[0][2], matches[0][3]}
	}
	return instructions, network
}

func parseInput2(input string) ([]rune, []string, Network) {
	network := make(Network)
	lines := strings.Split(input, "\n")
	start := []string{}
	instructions := []rune(lines[0])
	re := regexp.MustCompile(`(\w+) = \((\w+), (\w+)\)`)
	
	for _, line := range lines[2:] {
		matches := re.FindAllStringSubmatch(line, -1)
		network[matches[0][1]] = Node{matches[0][2], matches[0][3]}
		if matches[0][1][2] == 'A' {
			start = append(start, matches[0][1])
		}
	}
	return instructions, start, network
}

func getSteps(next string, instructions []rune, network Network, partOne bool) int{
	steps := 0
	for {
		for _, ins := range instructions {
			steps++
			if ins == 'L' {
				next = network[next].left
			} else {
				next = network[next].right
			}
			if partOne {
				if next == "ZZZ" {
					return steps
				}
			} else {
				if next[2] == 'Z' {
					return steps
				}
			}
		}
	}
}