package main

import (
	"fmt"
	"io/ioutil"
	"regexp"
	"strconv"
	"strings"
)

func main() {
	mem := make(map[int64]int64)                       // The values are 36-bit - using int64s
	memre, _ := regexp.Compile(`mem\[(\d+)\] = (\d+)`) // isn't matching all memory lines?
	maskre, _ := regexp.Compile(`mask = ([01X]+)`)
	d, _ := ioutil.ReadFile("input")
	lines := strings.Split(string(d), "\n")
	var mask string
	for _, l := range lines {
		memMatch := memre.FindStringSubmatch(l)
		maskMatch := maskre.FindStringSubmatch(l)
		if len(memMatch) != 0 {
			addr, _ := strconv.ParseInt(memMatch[1], 10, 64)
			value, _ := strconv.ParseInt(memMatch[2], 10, 64)

			// Applying the bitmask using strings because I'm doing this the stupid way
			valueString := strconv.FormatInt(value, 2) // value as a binary string
			v := ""
			// Pad out to 36 bits
			for i := 36 - len(valueString); i > 0; i-- {
				v = "0" + v
			}
			v = v + valueString
			applied := ""
			for i, _ := range v {
				if mask[i] != byte('X') {
					applied += string(mask[i])
				} else {
					applied += string(v[i])
				}
			}
			av, _ := strconv.ParseInt(applied, 2, 64) // applied
			mem[addr] = av
		} else if len(maskMatch) != 0 {
			mask = maskMatch[1]
		} else {
		}
	}
	var sum int64 = 0
	for _, v := range mem {
		sum += v
	}
	fmt.Println("Part 1:", sum)
	// Part 2
	mem = make(map[int64]int64) // reset memory
	mask = ""
	for _, l := range lines {
		memMatch := memre.FindStringSubmatch(l)
		maskMatch := maskre.FindStringSubmatch(l)
		if len(memMatch) != 0 {
			addr, _  := strconv.ParseInt(memMatch[1], 10, 64)
			value, _ := strconv.ParseInt(memMatch[2], 10, 64)
			addrString := strconv.FormatInt(addr, 2) // addr as a binary string
			v := ""
			// Pad out to 36 bits
			for i := 36 - len(addrString); i > 0; i-- {
				v = "0" + v
			}
			v = v + addrString
			applied := ""
			var addresses []string
			for i, _ := range v {
				if mask[i] != byte('0') {
					applied += string(mask[i])
				} else {
					applied += string(v[i])
				}
			}
			addresses = constructAddresses(applied)
			for _, a := range addresses {
				d, _ := strconv.ParseInt(a, 2, 64)
				mem[d] = value
			}
		} else if len(maskMatch) != 0 {
			mask = maskMatch[1]
		} else {
		}
	}	
	sum = 0
	for _, v := range mem {
		sum += v
	}
	fmt.Println("Part 2:", sum)

}

func constructAddresses(addr string) []string {
	var out []string
	x := false
	for _, c := range addr {
		if c == 'X' {
			x = true
		}
	}
	if x == true {
		ones := constructAddresses(strings.Replace(addr, "X", "1", 1))
		zeroes := constructAddresses(strings.Replace(addr, "X", "0", 1))
		for _, i := range ones {
			out = append(out, i)
		}
		for _, i := range zeroes {
		out = append(out, i)
		}
		return out
	} 
	return []string{addr}
}