package main

import (
	"bufio"
	"fmt"
	"log"
	"math"
	"os"
	"regexp"
	"strconv"
	"strings"
)

func processLines(path string, f func(string) (stop bool, err error)) error {
	file, err := os.Open(path)
	if err != nil {
		return err
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		stop, err := f(scanner.Text())
		if stop || err != nil {
			return err
		}
	}

	if err := scanner.Err(); err != nil {
		return err
	}

	return nil
}

var memExtractor = regexp.MustCompile(`^mem\[([0-9]+)] = ([0-9]+)`)

func powInt(x, y int) int {
	return int(math.Pow(float64(x), float64(y)))
}

func applyMask(mask string, v int) (int, error) {
	for i := 0; i < len(mask); i++ {
		m := mask[len(mask)-i-1]
		switch m {
		case 'X':
		case '1':
			v |= powInt(2, i)
		case '0':
			v &^= powInt(2, i)
		default:
			return 0, fmt.Errorf("could not apply mask %v in position %d", m, i)
		}
	}

	return v, nil
}

func sprintBinArray(a []int, z int) string {
	r := " "

	for _, v := range a {
		f := fmt.Sprintf("%%%db(%%d)", z)
		r += fmt.Sprintf(f, v, v)
	}

	return fmt.Sprintf("[%s]", r)
}

func applyFloatingMask(mask string, v int) ([]int, error) {
	base2 := 2
	results := []int{v}

	// fmt.Printf("Applying %s to %8b\n", mask, v)

	for i := 0; i < len(mask); i++ {
		m := mask[len(mask)-i-1]
		// fmt.Printf("Apply %s at %d ", string(m), i)

		resultLen := len(results)
		for j := 0; j < resultLen; j++ {
			switch m {
			case 'X':
				// Add a 0 mask to end of results
				results = append(results, results[j]|powInt(base2, i))
				// Apply a 1 mask in place
				results[j] &^= powInt(base2, i)
			case '1':
				results[j] |= powInt(base2, i)
			case '0':
				// results[j] &^= powInt(base2, i)
			default:
				return results, fmt.Errorf("could not apply mask %v in position %d", m, i)
			}
		}

		// fmt.Printf("Result %s\n", sprintBinArray(results, 8))
	}

	return results, nil
}

func parseMemValues(line string) (loc int, val int, err error) {
	match := memExtractor.FindStringSubmatch(line)
	if match == nil {
		err = fmt.Errorf("could not extract mem val in %s", line)

		return
	}

	loc, err = strconv.Atoi(match[1])
	if err != nil {
		err = fmt.Errorf("error parsing int from %s: %w", line, err)

		return
	}

	val, err = strconv.Atoi(match[2])
	if err != nil {
		err = fmt.Errorf("error parsing int from %s: %w", line, err)

		return
	}

	return
}

func part1() {
	var mask string

	mem := map[int]int{}

	if err := processLines("input.txt", func(line string) (bool, error) {
		if strings.HasPrefix(line, "mask") {
			mask = strings.TrimPrefix(line, "mask = ")
		} else {
			loc, val, err := parseMemValues(line)
			if err != nil {
				return false, err
			}
			val, err = applyMask(mask, val)
			if err != nil {
				return false, fmt.Errorf("error applying mask to value on line %s: %w", line, err)
			}
			// fmt.Printf("mem[%d] = %d\n", loc, val)
			mem[loc] = val
		}

		return false, nil
	}); err != nil {
		log.Fatal(err)
	}

	sum := 0
	for _, v := range mem {
		sum += v
	}

	fmt.Printf("Total value in memory %d\n", sum)
}

func part2() {
	var mask string

	mem := map[int]int{}

	if err := processLines("input.txt", func(line string) (bool, error) {
		if strings.HasPrefix(line, "mask") {
			mask = strings.TrimPrefix(line, "mask = ")
		} else {
			loc, val, err := parseMemValues(line)
			if err != nil {
				return false, err
			}
			locs, err := applyFloatingMask(mask, loc)
			if err != nil {
				return false, fmt.Errorf("error applying mask to value on line %s: %w", line, err)
			}
			for _, l := range locs {
				mem[l] = val
			}
		}

		return false, nil
	}); err != nil {
		log.Fatal(err)
	}

	sum := 0
	for _, v := range mem {
		sum += v
	}

	fmt.Printf("Total value in memory %d\n", sum)
}

func main() {
	part1()
	part2()
}