package day21

import (
	"strconv"
	"strings"
)

type coordinate struct {
	row, col int
}

var numberPad = map[string]coordinate{
	"7": {0, 0}, "8": {0, 1}, "9": {0, 2},
	"4": {1, 0}, "5": {1, 1}, "6": {1, 2},
	"1": {2, 0}, "2": {2, 1}, "3": {2, 2},
	"0": {3, 1}, "A": {3, 2},
}

var robotPad = map[string]coordinate{
	"^": {0, 1}, "A": {0, 2},
	"<": {1, 0}, "v": {1, 1}, ">": {1, 2},
}

type cacheKey struct {
	sequence   string
	iterations int
}

var memo = make(map[cacheKey]int)

func Part1(input string) int {
	codes := parseInput(input)

	numpadGraph := createGraph(numberPad, coordinate{3, 0})
	dirpadGraph := createGraph(robotPad, coordinate{0, 0})

	total := 0
	for _, code := range codes {
		codeInt, _ := strconv.Atoi(code[:len(code)-1])
		total += codeInt * getLength(code, 3, true, numpadGraph, dirpadGraph)
	}
	return total
}

func Part2(input string) int {
	codes := parseInput(input)

	numpadGraph := createGraph(numberPad, coordinate{3, 0})
	dirpadGraph := createGraph(robotPad, coordinate{0, 0})

	total := 0
	for _, code := range codes {
		codeInt, _ := strconv.Atoi(code[:len(code)-1])
		total += codeInt * getLength(code, 26, true, numpadGraph, dirpadGraph)
	}
	return total
}

func parseInput(input string) []string {
	parts := strings.Split(input, "\n")
	codes := []string{}
	codes = append(codes, parts...)
	return codes
}

func createGraph(keypad map[string]coordinate, invalidCoords coordinate) map[string]string {
	graph := make(map[string]string)

	for a, coordA := range keypad {
		for b, coordB := range keypad {
			x1, y1 := coordA.row, coordA.col
			x2, y2 := coordB.row, coordB.col

			path := strings.Repeat("<", max(0, y1-y2)) +
				strings.Repeat("v", max(0, x2-x1)) +
				strings.Repeat("^", max(0, x1-x2)) +
				strings.Repeat(">", max(0, y2-y1))

			if invalidCoords == (coordinate{x1, y2}) || invalidCoords == (coordinate{x2, y1}) {
				path = reverseString(path)
			}
			graph[a+"->"+b] = path + "A"
		}
	}

	return graph
}

func reverseString(s string) string {
	runes := []rune(s)
	for i, j := 0, len(runes)-1; i < j; i, j = i+1, j-1 {
		runes[i], runes[j] = runes[j], runes[i]
	}
	return string(runes)
}

func getLength(sequence string, iterations int, firstIter bool, numpadGraph map[string]string, dirpadGraph map[string]string) int {

	if iterations == 0 {
		return len(sequence)
	}

	key := cacheKey{sequence: sequence, iterations: iterations}
	if cachedResult, exists := memo[key]; exists {
		return cachedResult
	}

	graph := numpadGraph
	if !firstIter {
		graph = dirpadGraph
	}

	totalLength := 0
	prev := "A"
	for _, char := range sequence {
		graphKey := prev + "->" + string(char)
		if subSequence, exists := graph[graphKey]; exists {
			totalLength += getLength(subSequence, iterations-1, false, numpadGraph, dirpadGraph)
		}
		prev = string(char)
	}

	memo[key] = totalLength
	return totalLength
}