use crate::prelude::*;

mod operation;

use operation::*;

pub fn solution() -> Solution {
  Solution::new(Day::new(14, 2020), part_1, part_2)
    .with_expected(9967721333886_i64, 4355897790573_i64)
}

fn part_1(input: &str) -> Result<String> {
  let operations = input.lines().map(Into::into).collect::<Vec<Operation>>();

  let mut current_mask = "".to_string();
  let mut memory = HashMap::new();

  for operation in operations {
    match operation {
      Operation::SetMask(mask) => current_mask = mask.clone(),
      Operation::SetMemory(address, mut value) => {
        for (index, character) in current_mask.chars().rev().enumerate() {
          match character {
            '0' => value &= !(1 << index),
            '1' => value |= 1 << index,
            _ => (),
          }
        }

        *memory.entry(address).or_insert(value) = value;
      }
    }
  }

  Ok(memory.values().sum::<i64>().to_string())
}

fn part_2(input: &str) -> Result<String> {
  let operations = input.lines().map(Into::into).collect::<Vec<Operation>>();

  let mut current_mask = "".to_string();
  let mut memory = HashMap::new();

  for operation in operations {
    match operation {
      Operation::SetMask(mask) => current_mask = mask.clone(),
      Operation::SetMemory(mut address, value) => {
        for (index, character) in current_mask.chars().rev().enumerate() {
          if character == '1' {
            address |= 1 << index;
          }
        }

        let mut new_mask = String::new();
        for (index, character) in format!("{:036b}", address).char_indices() {
          if current_mask.chars().nth(index) == Some('X') {
            new_mask += "X";
          } else {
            new_mask += &character.to_string();
          }
        }

        let mut floating_addresses = HashSet::new();
        floating_addresses.insert(new_mask);
        while floating_addresses
          .iter()
          .any(|address| address.contains('X'))
        {
          for address in floating_addresses.clone() {
            for new_address in combine(address.clone()) {
              floating_addresses.insert(new_address);
            }

            floating_addresses.remove(&address);
          }
        }

        let floating_addresses = floating_addresses
          .iter()
          .map(|address| i64::from_str_radix(address, 2).unwrap())
          .collect::<Vec<i64>>();
        for mut address in floating_addresses {
          for (index, character) in current_mask.chars().rev().enumerate() {
            if character == '1' {
              address |= 1 << index;
            }
          }

          *memory.entry(address).or_insert(value) = value;
        }
      }
    }
  }

  Ok(memory.values().sum::<i64>().to_string())
}

fn combine(input: String) -> Vec<String> {
  let mut result = vec![];

  for (index, character) in input.char_indices() {
    if character == 'X' {
      let zero = format!("{}{}{}", &input[0..index], '0', &input[index + 1..]);
      result.push(zero);
      let one = format!("{}{}{}", &input[0..index], '1', &input[index + 1..]);
      result.push(one);
    }
  }

  result
}