package two

import (
	"bytes"
	"strings"

	"github.com/onyx-and-iris/aoc2024/day-21/internal/pad"
)

func Solve(buf []byte) (int, error) {
	r := bytes.NewReader(buf)
	codes, err := parseLines(r)
	if err != nil {
		return 0, err
	}

	complexityChan := make(chan int)
	conc := len(codes)

	for _, code := range codes {
		go func() {
			numpad := pad.NewNumpad()
			dirpad := pad.NewDirpad()
			complexityChan <- numFromCode(code) * generateIntFromMoves(numpad, dirpad, code, 25)
		}()
	}

	var complexity int
	for range conc {
		complexity += <-complexityChan
	}

	return complexity, nil
}

type generator interface {
	Generate(code string) string
}

func generateIntFromMoves(n, d generator, code string, numRobots int) int {
	moves := n.Generate(code)
	return recurseForRobots(d, moves, numRobots, 1, newRobotCache())
}

func recurseForRobots(
	d generator,
	moves string,
	numRobots int,
	id int,
	memo robotCache,
) int {
	if memo.contains(moves) {
		if v, ok := memo.read(moves, id-1); ok {
			return v
		}
	}

	next := d.Generate(moves)
	if id == numRobots {
		return len(next)
	}

	var totalCount int
	for _, move := range strings.SplitAfter(next, "A") {
		count := recurseForRobots(d, move, numRobots, id+1, memo)
		totalCount += count
	}
	memo.insert(moves, id-1, totalCount)

	return totalCount
}