package main

import (
	"bufio"
	"fmt"
	"os"
	"strconv"
	"strings"
	"time"
)

func main() {
	start := time.Now()
	partOne()
	duration := time.Since(start)
	partTwo()
	duration2 := time.Since(start)
	fmt.Printf("p1: %s, p2: %s\n", duration, duration2-duration)
}

type bitset struct {
	bits []string
}

func newBitset(i int) bitset {
	b := strings.Split(fmt.Sprintf("%036b", i), "")
	bs := bitset{
		bits: strings.Split("000000000000000000000000000000000000", ""),
	}
	for i, bt := range b {
		if bt != "0" {
			bs.bits[i] = bt
		}
	}
	return bs
}

func (b bitset) ToInt() int {
	i, _ := strconv.ParseInt(strings.Join(b.bits, ""), 2, 64)
	return int(i)
}

func parseMask(s string) string {
	return strings.Split(s, " = ")[1]
}

func parseMemoryToAddrAndBitset(s string) (int, bitset) {
	a := strings.Split(s, " = ")
	addr, _ := strconv.Atoi(strings.TrimFunc(a[0], func(r rune) bool {
		if r == 'm' || r == 'e' || r == '[' || r == ']' {
			return true
		}
		return false
	}))
	i, _ := strconv.Atoi(a[1])
	return addr, newBitset(i)
}

func parseMemoryToBitsetAndValue(s string) (bitset, int) {
	a := strings.Split(s, " = ")
	addr, _ := strconv.Atoi(strings.TrimFunc(a[0], func(r rune) bool {
		if r == 'm' || r == 'e' || r == '[' || r == ']' {
			return true
		}
		return false
	}))
	i, _ := strconv.Atoi(a[1])
	return newBitset(addr), i
}

func applyMaskToValue(m string, b bitset) int {
	for i, r := range m {
		if r == '0' {
			b.bits[i] = "0"
		}
		if r == '1' {
			b.bits[i] = "1"
		}
	}
	return b.ToInt()
}

func applyMaskToAddr(prefix []string, m string, attr bitset) []bitset {
	if len(m) == 0 {
		return []bitset{{bits: prefix}}
	}
	res := []bitset{}
	for _, r := range m {
		if r == '0' || r == 'X' {
			res = append(res, applyMaskToAddr(append(prefix[:], "0"), m[1:], attr)...)
		}
		if r == '1' || r == 'X' {
			res = append(res, applyMaskToAddr(append(prefix[:], "1"), m[1:], attr)...)
		}
	}
	return res
}

func partOne() {
	f, _ := os.Open("input")
	reader := bufio.NewReader(f)
	scanner := bufio.NewScanner(reader)

	mem := map[int]int{}
	var currentMask string

	for scanner.Scan() {
		line := scanner.Text()
		if strings.HasPrefix(line, "mask") {
			currentMask = parseMask(line)
		} else {
			addr, value := parseMemoryToAddrAndBitset(line)
			mem[addr] = applyMaskToValue(currentMask, value)
		}
	}
	sum := 0
	for _, v := range mem {
		sum = sum + v
	}
	fmt.Println(sum)
}

func partTwo() {
	f, _ := os.Open("testinput")
	reader := bufio.NewReader(f)
	scanner := bufio.NewScanner(reader)

	mem := map[int]int{}
	var currentMask string

	for scanner.Scan() {
		line := scanner.Text()
		if strings.HasPrefix(line, "mask") {
			currentMask = parseMask(line)
		} else {
			addr, value := parseMemoryToBitsetAndValue(line)
			for _, b := range applyMaskToAddr([]string{}, currentMask, addr) {
				mem[b.ToInt()] = value
			}
		}
	}

	sum := 0
	for _, v := range mem {
		sum = sum + v
	}
	fmt.Println(sum)
}