use nom::{
    branch::alt,
    character::complete::{char, line_ending, space1, u32},
    combinator::{map, value},
    multi::{many1, many1_count, separated_list1},
    sequence::separated_pair,
    IResult,
};
use std::{collections::BTreeMap, fs};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let input = fs::read_to_string("inputs/day_12.txt")?;
    let small = SpringField::parser(&input).unwrap().1;
    dbg!(&small.possibilities_sum());

    let large = small.expand();
    dbg!(&large.possibilities_sum());

    Ok(())
}

#[derive(Debug)]
struct SpringField(Vec<SpringRow>);

#[derive(Debug, Clone)]
struct SpringRow {
    springs: Vec<Spring>,
    check: Vec<u32>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Spring {
    Good,
    Bad(usize),
    Unknown,
}

impl SpringField {
    fn parser(input: &str) -> IResult<&str, Self> {
        map(separated_list1(line_ending, SpringRow::parser), SpringField)(input)
    }

    fn possibilities_sum(&self) -> usize {
        self.0.iter().map(|r| r.possibilities_count()).sum()
    }

    fn expand(&self) -> SpringField {
        SpringField(self.0.iter().map(|r| r.expand()).collect())
    }
}

impl SpringRow {
    fn parser(input: &str) -> IResult<&str, Self> {
        map(
            separated_pair(
                many1(Spring::parser),
                space1,
                separated_list1(char(','), u32),
            ),
            |(springs, check)| SpringRow { springs, check },
        )(input)
    }

    fn expand(&self) -> SpringRow {
        let mut expanded = SpringRow {
            springs: Vec::new(),
            check: Vec::new(),
        };

        for _ in 0..5 {
            expanded.springs.append(&mut self.springs.clone());
            expanded.springs.push(Spring::Unknown);
            expanded.check.append(&mut self.check.clone());
        }

        // should only have the extra Unknown between, not at the very end.
        expanded.springs.pop();

        expanded
    }

    fn possibilities_count(&self) -> usize {
        let optimized = self.optimize_tail();
        optimized.possibilities_count_inner(0, 0, &mut BTreeMap::new())
    }

    fn optimize_tail(&self) -> SpringRow {
        let mut optimized = self.clone();
        while optimized.springs.len() > 0 && optimized.check.len() > 0 {
            match optimized.springs[optimized.springs.len() - 1] {
                Spring::Good => {
                    optimized.springs.pop();
                }
                Spring::Bad(s) if s > optimized.check[optimized.check.len() - 1] as usize => {
                    panic!("Unsolvable row");
                }
                Spring::Bad(s) if s <= optimized.check[optimized.check.len() - 1] as usize => {
                    optimized.trim_bad_suffix();
                }
                Spring::Bad(_) => unreachable!(),
                Spring::Unknown => {
                    break;
                }
            }
        }

        optimized
    }

    fn trim_bad_suffix(&mut self) {
        let last_check = self.check.pop().expect("No trailing check to pop") as usize;
        let mut check_i = 0;
        while check_i < last_check {
            match self
                .springs
                .pop()
                .expect("Ran out of springs while optimizing suffix!")
            {
                Spring::Unknown => {
                    check_i += 1;
                }
                Spring::Bad(inc) => {
                    check_i += inc;
                }
                Spring::Good => {
                    panic!("Found a good spring in the middle of what must be a bad spring range");
                }
            }
        }
        let is_boundary = self.springs.len() == 0;
        let has_extra_good_or_unknown_to_trim =
            !is_boundary && !matches!(self.springs.last(), Some(Spring::Bad(_)));

        let valid_suffix =
            check_i == last_check && (is_boundary || has_extra_good_or_unknown_to_trim);
        if !valid_suffix {
            panic!("The suffix had another invalid bad section immediately!");
        }

        if has_extra_good_or_unknown_to_trim {
            self.springs.pop();
        }
    }

    fn possibilities_count_inner(
        &self,
        springs_offset: usize,
        check_offset: usize,
        cache: &mut BTreeMap<(usize, usize), usize>,
    ) -> usize {
        if let Some(cached) = cache.get(&(springs_offset, check_offset)) {
            return *cached;
        };

        let (springs_offset, check_offset) = match self.optimize_head(springs_offset, check_offset)
        {
            Some(optimized) => optimized,
            None => {
                return 0;
            }
        };

        let springs = &self.springs[springs_offset..];
        let check = &self.check[check_offset..];

        if check.len() == 1 && springs.iter().all(|s| !matches!(s, Spring::Bad(_))) {
            let mut contigous_unknowns = Vec::new();
            let mut next_contiguous_unknown = 0;
            let mut is_in_unknown = false;
            for spring in springs.iter() {
                match spring {
                    Spring::Unknown => {
                        next_contiguous_unknown += 1;
                        is_in_unknown = true;
                    }
                    Spring::Good => {
                        if is_in_unknown {
                            contigous_unknowns.push(next_contiguous_unknown);
                            next_contiguous_unknown = 0;
                            is_in_unknown = false;
                        }
                    }
                    Spring::Bad(_) => unreachable!(),
                }
            }
            if is_in_unknown {
                contigous_unknowns.push(next_contiguous_unknown);
            }

            return contigous_unknowns
                .iter()
                .map(|region| {
                    if *region >= check[0] as usize {
                        region - check[0] as usize + 1
                    } else {
                        0
                    }
                })
                .sum();
        }

        if check.len() == 0 {
            if springs.iter().any(|s| matches!(s, Spring::Bad(_))) {
                return 0;
            } else {
                return 1;
            }
        }

        let valid_prefix_possibilities_count =
            if let Some((without_prefix_springs_offset, without_prefix_check_offset)) =
                self.trim_bad_prefix(springs_offset, check_offset)
            {
                self.possibilities_count_inner(
                    without_prefix_springs_offset,
                    without_prefix_check_offset,
                    cache,
                )
            } else {
                0
            };

        let non_prefix_possibilities_count =
            self.possibilities_count_inner(springs_offset + 1, check_offset, cache);

        let total_possibilities = valid_prefix_possibilities_count + non_prefix_possibilities_count;

        cache.insert((springs_offset, check_offset), total_possibilities);

        total_possibilities
    }

    fn optimize_head(
        &self,
        mut springs_offset: usize,
        mut check_offset: usize,
    ) -> Option<(usize, usize)> {
        while springs_offset < self.springs.len() && check_offset < self.check.len() {
            match self.springs[springs_offset] {
                Spring::Good => {
                    springs_offset += 1;
                }
                Spring::Bad(s) if s > self.check[check_offset] as usize => {
                    return None;
                }
                Spring::Bad(s) if s <= self.check[check_offset] as usize => {
                    match self.trim_bad_prefix(springs_offset, check_offset) {
                        Some((new_springs, new_check)) => {
                            springs_offset = new_springs;
                            check_offset = new_check;
                        }
                        None => return None,
                    };
                }
                Spring::Bad(_) => unreachable!(),
                Spring::Unknown => {
                    break;
                }
            }
        }

        if springs_offset == self.springs.len() && check_offset != self.check.len() {
            None
        } else {
            Some((springs_offset, check_offset))
        }
    }

    fn trim_bad_prefix(
        &self,
        springs_offset: usize,
        check_offset: usize,
    ) -> Option<(usize, usize)> {
        let first_check = self.check[check_offset] as usize;
        let mut check_i = 0;
        let mut springs_i = springs_offset;
        while check_i < first_check {
            if springs_i >= self.springs.len() {
                return None;
            }

            match self.springs[springs_i] {
                Spring::Unknown => {
                    check_i += 1;
                }
                Spring::Bad(inc) => {
                    check_i += inc;
                }
                Spring::Good => {
                    return None;
                }
            }
            springs_i += 1;
        }

        let is_boundary = springs_i == self.springs.len();
        let has_extra_good_or_unknown_to_trim =
            !is_boundary && !matches!(self.springs[springs_i], Spring::Bad(_));
        let valid_prefix =
            check_i == first_check && (is_boundary || has_extra_good_or_unknown_to_trim);

        if valid_prefix {
            let new_springs_i = if has_extra_good_or_unknown_to_trim {
                springs_i + 1
            } else {
                springs_i
            };
            Some((new_springs_i, check_offset + 1))
        } else {
            None
        }
    }
}

impl Spring {
    fn parser(input: &str) -> IResult<&str, Self> {
        alt((
            value(Spring::Good, many1_count(char('.'))),
            map(many1_count(char('#')), Spring::Bad),
            value(Spring::Unknown, char('?')),
        ))(input)
    }
}