package main

import (
	"AoC2020/helper"
	"fmt"
	"os"
	"strconv"
	"strings"
)

var Mem = make(map[int]int64)
var Mask string

func main() {
	args := os.Args[1:]
	input, err := helper.GetInput(args[0])
	if err != nil {
		fmt.Println(err)
	}
	execute(input, run)
	Mem = make(map[int]int64)
	execute(input, runPart2)
}

func execute(input []string, f func(a string)) {
	for _, val := range input {
		f(val)
	}
	count := 0
	for  _, value := range Mem {
		count = count + int(value)
	}
	fmt.Println(count)
}

func run(row string) {
	command := strings.Split(row," = ")
	if command[0] == "mask" {
		Mask = command[1]
	} else {
		location, _ := strconv.Atoi(command[0][4:len(command[0])-1])
		number, _ := strconv.Atoi(command[1])
		save(number, Mask, location)
	}
}

func save(number int, mask string, location int) {
	bitNumber := strconv.FormatInt(int64(number), 2)
	bitNumber = addPrecedingZeroes(bitNumber)
	bitNumber = applyMask(bitNumber,mask)
	saveToM(bitNumber, location)
}

func applyMask(number string, mask string) string{
	var newNumber string
	for i, val := range mask {
		switch val {
		case 'X': newNumber = newNumber + string(number[i])
		case '1': newNumber = newNumber + "1"
		case '0': newNumber = newNumber + "0"
		}
	}
	return newNumber
}

func runPart2(row string) {
	command := strings.Split(row," = ")
	if command[0] == "mask" {
		Mask = command[1]
	} else {
		location, _ := strconv.Atoi(command[0][4:len(command[0])-1])
		number, _ := strconv.Atoi(command[1])
		savePart2(number, Mask, location)
	}
}

func savePart2(number int, mask string, location int) {
	bitNumber := strconv.FormatInt(int64(location), 2)
	bitNumber = addPrecedingZeroes(bitNumber)
	bitNumbers := applyMaskPart2(bitNumber,mask)
	for _,val := range bitNumbers {
		Mem[int(getIntFromBitNumber(val))] = int64(number)
	}
}

func saveToM(bitNumber string, location int) {
	Mem[location] = getIntFromBitNumber(bitNumber)
}

func getIntFromBitNumber(bitNumber string) int64 {
	if i, err := strconv.ParseInt(bitNumber, 2, 64); err != nil {
		fmt.Println(err)
	} else {
		return i
	}
	return -1
}

func applyMaskPart2(number string, mask string) []string{
	var newNumber string
	for i, val := range mask {
		switch val {
		case 'X': newNumber = newNumber + "X"
		case '1': newNumber = newNumber + "1"
		case '0': newNumber = newNumber + string(number[i])
		}
	}
	return getVariations(0,newNumber)
}

func getVariations(pointer int ,number string) []string{
	var variations []string
	for i := pointer; i < len(number); i++ {
		if number[i] == 'X' {
			number1 := number + ""
			number1 = replaceAtIndex(number1,'0',i)
			number2 := number + ""
			number2 = replaceAtIndex(number1,'1',i)
			variations = append(variations, getVariations(i, number1) ...)
			variations = append(variations, getVariations(i, number2) ...)
		}
	}
	if strings.Index(number, "X") == -1 {
		variations = append(variations, number)
	}
	return variations
}

func addPrecedingZeroes(number string) string {
	return strings.Repeat("0", 36-len(number)) + number
}

func replaceAtIndex(in string, r rune, i int) string {
	out := []rune(in)
	out[i] = r
	return string(out)
}