use std::collections::HashMap;

use anyhow::Context;

use crate::{Problem, Solution};

pub struct Day15;

impl Problem for Day15 {
    const DAY: u8 = 15;

    const INPUT: &'static str = include_str!("../input/day_15.txt");
}

impl Solution for Day15 {
    type Answer1 = usize;

    type Answer2 = usize;

    fn part_1(input: &str) -> anyhow::Result<Self::Answer1> {
        Ok(input.trim().split(',').map(hash).sum())
        // 509784 ==
    }

    fn part_2(input: &str) -> anyhow::Result<Self::Answer2> {
        Ok(input
            .trim()
            .split(',')
            .try_fold(Boxes::new(), process_step)?
            .into_iter()
            .map(calculate_power)
            .sum())

        // 339464 >
        // 230197 ==
    }
}

type Lenses = Vec<(String, usize)>;
type Boxes = HashMap<usize, Lenses>;

fn hash(s: &str) -> usize {
    s.trim()
        .chars()
        .fold(0, |acc, c| (acc + c as usize) * 17 % 256)
}

fn process_step(mut acc: Boxes, s: &str) -> anyhow::Result<Boxes> {
    let (label, power) = s.split_once(['-', '=']).context("Failed to parse input")?;

    let key = hash(label);
    let lenses = acc.entry(key).or_default();

    match lenses.iter().position(|(s, _)| s == label) {
        pos if !power.is_empty() => {
            lenses.push((label.to_owned(), power.parse()?));
            if let Some(index) = pos {
                lenses.swap_remove(index);
            }
        }
        Some(index) => {
            lenses.remove(index);
            if lenses.is_empty() {
                acc.remove(&key);
            }
        }
        _ => {}
    };

    Ok(acc)
}

fn calculate_power((id, lenses): (usize, Lenses)) -> usize {
    lenses
        .into_iter()
        .enumerate()
        .map(|(i, (_, len))| (id + 1) * (i + 1) * len)
        .sum::<usize>()
}

#[cfg(test)]
mod tests {
    use super::*;

    const INPUT: &str = "rn=1,cm-,qp=3,cm=2,qp-,pc=4,ot=9,ab=5,pc-,pc=6,ot=7";

    #[test]
    fn test_part_1() -> anyhow::Result<()> {
        Ok(assert_eq!(1320, Day15::part_1(INPUT)?))
    }

    #[test]
    fn test_part_2() -> anyhow::Result<()> {
        Ok(assert_eq!(145, Day15::part_2(INPUT)?))
    }
}