#!/usr/bin/env -S java --source 11
package io.github.hiljusti.adventofcode.solutions;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import java.util.stream.Stream;

class Main {
    public static void main(String[] args) throws IOException {
        final List<String> input = Files.readAllLines(Path.of("./input"));

        System.out.printf("Sum of all memory values: %d%n", solve1(input));

        System.out.printf("Sum of all memory values (version 2): %d%n", solve2(input));
    }

    static Long solve1(List<String> input) {
        final Map<String, Long> memory = new HashMap<>();
        String mask = "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX";

        for (String line: input) {
            if (line.length() == 0) { continue; }
            var parts = line.split(" = ");
            
            if (parts[0].startsWith("mask")) {
                mask = parts[1];
            }
            else {
                var address = parts[0].replace("mem[", "").replace("]", "");
                var value = String.format("%36s",
                        Long.toBinaryString(Long.parseLong(parts[1])))
                    .replace(' ', '0');

                char[] result = new char[36];
                for (int i = 0; i < 36; i++) {
                    var maskChar = mask.charAt(i);
                    result[i] = maskChar != 'X'? maskChar : value.charAt(i);
                }

                memory.put(address, Long.parseLong(new String(result), 2));
            }
        }

        return memory.values().stream().mapToLong(n -> n).sum();
    }
    
    static Long solve2(List<String> input) {
        final Map<String, Long> memory = new HashMap<>();
        String mask = "000000000000000000000000000000000000";

        for (String line: input) {
            if (line.length() == 0) { continue; }
            var parts = line.split(" = ");
            
            if (parts[0].startsWith("mask")) {
                mask = parts[1];
            }
            else {
                var value = Long.parseLong(parts[1]);
                
                var address = String.format(
                        "%36s",
                        Long.toBinaryString(Long.parseLong(parts[0].replace("mem[", "").replace("]", "")))
                    )
                    .replace(' ', '0');
                char[] result = new char[36];
                for (int i = 0; i < 36; i++) {
                    var maskChar = mask.charAt(i);
                    result[i] = maskChar != '0'? maskChar : address.charAt(i);
                }
                
                List<String> addresses = List.of(new String(result));
                for (int i = 0; i < 36; i++) {
                    if (mask.charAt(i) == 'X') {
                        final int ptr = i; // rebound for lambda context.
                        addresses = addresses.stream()
                            .flatMap(addr -> {
                                final char[] raw0 = addr.toCharArray();
                                raw0[ptr] = '0';
                                final String a = new String(raw0);

                                final char[] raw1 = addr.toCharArray();
                                raw1[ptr] = '1';
                                final String b = new String(raw1);

                                return Stream.of(a, b);
                            })
                            .collect(Collectors.toList());
                    }
                }
                
                addresses.forEach(addr -> memory.put(addr, value));
            }
        }

        return memory.values().stream().mapToLong(n -> n).sum();
    }
}