package main

import (
	"fmt"
	"strings"

	h "git.bullercodeworks.com/brian/adventofcode/helpers"
)

func main() {
	fmt.Println("# Day 14")
	inp := h.StdinToStringSlice()
	vsn := h.Atoi(h.OptArgNumber(1, "1"))
	if vsn == 1 {
		part1(inp)
	} else {
		part2(inp)
	}
}

func part1(inp []string) {
	memory := make(map[int]uint)
	var mask map[int]uint
	for _, v := range inp {
		pts := strings.Split(v, " = ")
		if pts[0] == "mask" {
			mask = make(map[int]uint)
			parsemask := pts[1]
			for k := range parsemask {
				if parsemask[k] != 'X' {
					place := (len(parsemask) - k - 1)
					if parsemask[k] == '0' {
						mask[place] = 0
					} else {
						mask[place] = 1
					}
				}
			}
		} else {
			var reg int
			fmt.Sscanf(pts[0], "mem[%d]", &reg)
			val := uint(h.Atoi(pts[1]))
			store := applyMask(val, mask)
			memory[reg] = store
		}
	}

	var res uint
	for _, v := range memory {
		res = res + v
	}
	fmt.Println("## Part 1")
	fmt.Println("Sum:", res)
}

func applyMask(val uint, mask map[int]uint) uint {
	for k, v := range mask {
		if v == 0 {
			val = val &^ (1 << k)
		} else if v == 1 {
			val = val | (1 << k)
		}
	}
	return val
}

func part2(inp []string) {
	memory := make(map[uint]uint)
	var mask string
	for _, v := range inp {
		pts := strings.Split(v, " = ")
		if pts[0] == "mask" {
			mask = pts[1]
		} else {
			var reg int
			fmt.Sscanf(pts[0], "mem[%d]", &reg)
			val := uint(h.Atoi(pts[1]))
			stores := []uint{uint(reg)}
			for k := range mask {
				bit := len(mask) - k - 1
				var wrk []uint
				for i := range stores {
					switch mask[k] {
					case '0':
						wrk = append(wrk, stores[i])
					case '1':
						wrk = append(wrk, setBitTo(stores[i], bit, true))
					case 'X':
						wrk = append(wrk, floatBit(stores[i], bit)...)
					}
				}
				stores = append([]uint{}, wrk...)
			}
			for k := range stores {
				memory[stores[k]] = val
			}
		}
	}
	var sum uint
	for k := range memory {
		sum = sum + memory[k]
	}
	fmt.Println("## Part 2")
	fmt.Println("Sum:", sum)
}

func setBitTo(addr uint, pos int, val bool) uint {
	if val {
		return addr | (1 << pos)
	} else {
		return addr &^ (1 << pos)
	}
}
func floatBit(addr uint, pos int) []uint {
	var ret []uint
	ret = append(ret, addr|(1<<pos))
	ret = append(ret, addr&^(1<<pos))
	return ret
}