use std::collections::{HashMap, HashSet};

use crate::{Problem, Solution};

pub struct Day03;

impl Problem for Day03 {
    const DAY: u8 = 3;

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

impl Solution for Day03 {
    type Answer1 = usize;

    type Answer2 = usize;

    fn part_1(input: &str) -> anyhow::Result<Self::Answer1> {
        let grid = input
            .trim()
            .split('\n')
            .map(|l| l.trim().chars().collect::<Vec<_>>())
            .collect::<Vec<_>>();

        Ok(grid.iter().enumerate().fold(0, |mut acc, (row, line)| {
            let line_iter = &mut line.iter().enumerate().peekable();
            while line_iter.peek().is_some() {
                let (part, is_part) = line_iter
                    .skip_while(|(_, n)| !n.is_ascii_digit())
                    .map_while(|(i, n)| n.to_digit(10).map(|u| (i, u)))
                    .fold((0, false), |(part, is_part), (col, n)| {
                        let is_part = is_part
                            || [row.saturating_sub(1), row, row + 1]
                                .into_iter()
                                .flat_map(|r| [col.saturating_sub(1), col, col + 1].map(|c| (r, c)))
                                .filter(|(r, c)| *r < grid.len() && *c < line.len())
                                .filter(|(r, c)| *r != row || *c != col)
                                .flat_map(|(r, c)| grid.get(r).and_then(|line| line.get(c)))
                                .any(|n| *n != '.' && !n.is_ascii_digit());

                        (part * 10 + n as usize, is_part)
                    });

                if is_part {
                    acc += part;
                }
            }
            acc
        }))
    }

    fn part_2(input: &str) -> anyhow::Result<Self::Answer2> {
        let grid = input
            .trim()
            .split('\n')
            .map(|l| l.trim().chars().collect::<Vec<_>>())
            .collect::<Vec<_>>();

        Ok(grid
            .iter()
            .enumerate()
            .fold(
                &mut HashMap::<Symbol, HashSet<Part>>::new(),
                |acc, (row, line)| {
                    let line_iter = &mut line.iter().enumerate().peekable();
                    while line_iter.peek().is_some() {
                        let (part, symbols) = line_iter
                            .skip_while(|(_, n)| !n.is_ascii_digit())
                            .map_while(|(i, n)| n.to_digit(10).map(|u| (i, u)))
                            .fold(
                                (Part::default(), HashSet::<Symbol>::new()),
                                |(mut part, mut neighbors), (col, n)| {
                                    let new_neighbors = [row.saturating_sub(1), row, row + 1]
                                        .into_iter()
                                        .flat_map(|r| {
                                            [col.saturating_sub(1), col, col + 1].map(|c| (r, c))
                                        })
                                        .filter(|(r, c)| *r < grid.len() && *c < line.len())
                                        .filter(|(r, c)| *r != row || *c != col)
                                        .flat_map(|(r, c)| {
                                            grid.get(r)
                                                .and_then(|line| line.get(c))
                                                .map(|n| ((r, c), *n))
                                        })
                                        .filter(|(_, n)| *n != '.' && !n.is_ascii_digit())
                                        .map(Symbol::from);

                                    neighbors.extend(new_neighbors);
                                    part.push_value(n as usize);
                                    (part, neighbors)
                                },
                            );

                        symbols
                            .into_iter()
                            .filter(|s| s.value == '*')
                            .for_each(|s| {
                                acc.entry(s)
                                    .and_modify(|h| {
                                        h.insert(part);
                                    })
                                    .or_insert(HashSet::from([part]));
                            });
                    }
                    acc
                },
            )
            .values_mut()
            .filter(|h| h.len() == 2)
            .fold(0, |acc, h| {
                acc + h
                    .iter()
                    .copied()
                    .map(|p| p.value)
                    .reduce(|acc, v| acc * v)
                    .unwrap()
            }))
    }
}

#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash)]
struct Symbol {
    value: char,
    pos: (usize, usize),
}

impl From<((usize, usize), char)> for Symbol {
    fn from((pos, value): ((usize, usize), char)) -> Self {
        Symbol { value, pos }
    }
}

#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash)]
struct Part {
    value: usize,
    start: (usize, usize),
    end: (usize, usize),
}

impl Part {
    fn push_value(&mut self, n: usize) {
        self.value = self.value * 10 + n
    }
}

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

    const INPUT: &str = indoc::indoc! {"
        467..114..
        ...*......
        ..35..633.
        ......#...
        617*......
        .....+.58.
        ..592.....
        ......755.
        ...$.*....
        .664.598..
    "};

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

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