use crate::prelude::*;

pub fn solution() -> Solution {
  Solution::new(Day::new(16, 2020), part_1, part_2)
    .with_expected(26980, 3021381607403_i64)
}

fn solve(input: &str, return_part_1: bool) -> Result<String> {
  let mut parts = input.trim().split("\n\n");
  let mut ranges = vec![];

  for line in parts.next().unwrap().lines() {
    let name = &line[0..line.find(':').unwrap()];

    let mut first_range = line
      [line.find(':').unwrap() + 2..line.find(" or").unwrap()]
      .split('-')
      .map(str::parse::<usize>)
      .map(Result::unwrap);

    let first_range = first_range.next().unwrap()..=first_range.next().unwrap();

    let mut second_range = line[line.find(" or ").unwrap() + 4..]
      .split('-')
      .map(str::parse::<usize>)
      .map(Result::unwrap);

    let second_range =
      second_range.next().unwrap()..=second_range.next().unwrap();

    ranges.push((name, first_range, second_range));
  }

  let my_ticket = parts
    .next()
    .unwrap()
    .lines()
    .nth(1)
    .unwrap()
    .split(',')
    .map(str::parse::<_>)
    .map(Result::unwrap)
    .collect::<Vec<_>>();

  let mut tickets = vec![];

  // Skip the first line in the 3rd part, as it will say "nearby tickets:".
  for line in parts.next().unwrap().lines().skip(1) {
    tickets.push(
      line
        .split(',')
        .map(str::parse::<usize>)
        .map(Result::unwrap)
        .collect::<Vec<_>>(),
    );
  }

  let mut result_one = 0;
  let mut valid_tickets = vec![my_ticket.clone()];

  for ticket in &tickets {
    let mut ticket_is_valid = true;

    for field in ticket {
      let field_is_invalid = ranges.iter().all(|(_, range_1, range_2)| {
        !range_1.contains(field) && !range_2.contains(field)
      });

      if field_is_invalid {
        result_one += field;
        ticket_is_valid = false;
      }
    }

    if ticket_is_valid {
      valid_tickets.push(ticket.clone());
    }
  }

  if return_part_1 {
    return Ok(result_one.to_string());
  }

  let mut possibilities = HashMap::new();
  let ticket_length = my_ticket.len();

  for (name, range_1, range_2) in ranges {
    for index in 0..ticket_length {
      let mut fields = valid_tickets
        .iter()
        .map(|tickets| tickets.get(index))
        .map(Option::unwrap);

      if fields.all(|field| range_1.contains(field) || range_2.contains(field))
      {
        let possibility = possibilities.entry(name).or_insert(vec![]);
        possibility.push(index);
      }
    }
  }

  let mut possibilities = possibilities.into_iter().collect::<Vec<_>>();
  possibilities.sort_by(|a, b| a.1.len().cmp(&b.1.len()));
  let mut indices = HashMap::<usize, &str>::new();

  for (name, possibilities) in possibilities {
    'inner: for possibility in possibilities {
      if indices.contains_key(&possibility) {
        continue 'inner;
      }

      indices.insert(possibility, name);
    }
  }

  let result_two = my_ticket
    .iter()
    .enumerate()
    .filter_map(|(index, number)| {
      if let Some(name) = indices.get(&index) {
        if name.starts_with("departure") {
          Some(number)
        } else {
          None
        }
      } else {
        None
      }
    })
    .product::<usize>()
    .to_string();

  Ok(result_two)
}

fn part_1(input: &str) -> Result<String> {
  solve(input, true)
}

fn part_2(input: &str) -> Result<String> {
  solve(input, false)
}