package day07

import (
	"adventofcode2024/utils"
	"fmt"
	"strconv"
	"strings"
)

type Op int

const (
	Add Op = iota
	Mul
	Conc
)

func Part1(input string) int {
	val := 0
	lines := strings.Split(input, "\n")
	for _, line := range lines {
		parts := strings.Split(line, ":")
		result := utils.MustAtoi(parts[0])
		numStrings := strings.Fields(parts[1])
		var numbers []int
		for _, numStr := range numStrings {
			num := utils.MustAtoi(numStr)
			numbers = append(numbers, num)
		}
		bitmask := 1 << (len(numbers) - 1)
		for i := 0; i < bitmask; i++ {
			calc := 0
			switch isNthBitSet(i, 0) {
			case true:
				calc = numbers[0] + numbers[1]
			case false:
				calc = numbers[0] * numbers[1]
			}

			for n := 1; n < len(numbers)-1; n++ {
				switch isNthBitSet(i, n) {
				case true:
					calc = calc + numbers[n+1]
				case false:
					calc = calc * numbers[n+1]
				}
			}
			if calc == result {
				fmt.Printf("result: %v numbers:%v  count: %d\n", result, numbers, bitmask)
				val += result
				break
			}
		}
	}
	return val
}

func Part2(input string) int {
	val := 0
	lines := strings.Split(input, "\n")
	for _, line := range lines {
		parts := strings.Split(line, ":")
		result := utils.MustAtoi(parts[0])
		numStrings := strings.Fields(parts[1])
		var numbers []int
		for _, numStr := range numStrings {
			num := utils.MustAtoi(numStr)
			numbers = append(numbers, num)
		}
		operations := []Op{Add, Mul, Conc}
		results := [][]Op{}
		generatePermutations(operations, len(numbers)-1, []Op{}, &results)
		for _, op := range results {
			calc := 0
			switch op[0] {
			case Add:
				calc = numbers[0] + numbers[1]
			case Mul:
				calc = numbers[0] * numbers[1]
			case Conc:
				calc = concat(numbers[0], numbers[1])
			}

			for n := 1; n < len(numbers)-1; n++ {
				switch op[n] {
				case Add:
					calc = calc + numbers[n+1]
				case Mul:
					calc = calc * numbers[n+1]
				case Conc:
					calc = concat(calc, numbers[n+1])
				}
			}
			if calc == result {
				fmt.Printf("result: %v numbers:%v\n", result, numbers)
				val += result
				break
			}
		}
	}
	return val
}

func isNthBitSet(num, n int) bool {
	return num&(1<<n) != 0
}

func getOp(row, col int) Op {
	vals := []Op{Add, Mul, Conc}
	// Column 0 cycles as row / 3
	if col == 0 {
		return vals[row/3]
	}
	// Column 1 cycles as row % 3
	return vals[row%3]
}

// Recursive function to generate permutations
func generatePermutations(ops []Op, n int, prefix []Op, results *[][]Op) {
	if n == 0 {
		// Append a copy of the current prefix to the results
		*results = append(*results, append([]Op{}, prefix...))
		return
	}

	// Iterate over all operations and build the combinations
	for _, op := range ops {
		generatePermutations(ops, n-1, append(prefix, op), results)
	}
}

func concat(a, b int) int {
	aStr := strconv.Itoa(a)
	bStr := strconv.Itoa(b)

	concatStr := aStr + bStr

	result, _ := strconv.Atoi(concatStr)
	return result
}
