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::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)]
struct SpringRow {
    springs: Vec<Spring>,
    check: Vec<u32>,
}

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

#[derive(Debug, Clone)]
struct WorkingSpringRow<'a> {
    springs: &'a [Spring],
    check: &'a [u32],
}

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_alt()).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
    }

    /// fail if it's impossible with that prefix to find any matches. Result is (valid, failed-after-reading-x)
    fn validate_check_prefix(&self, bad_springs: &[usize]) -> (bool, usize) {
        let mut check_iter = self.check.iter();
        let mut current_check = check_iter.next();
        let mut bad_spring_iter = bad_springs.iter();
        let mut current_bad_spring = bad_spring_iter.next();
        let mut bad_spring_read_count = 0;
        let mut current_count = 0;
        let mut in_bad_lands = false;

        for (i, spring) in self.springs.iter().enumerate() {
            let (bad, increment) = match spring {
                Spring::Good => (false, 1),
                Spring::Bad(s) => (true, *s),
                Spring::Unknown => {
                    if current_bad_spring.is_none() {
                        // That's the end of the prefix. Time to exit early.
                        let too_many_bads = in_bad_lands
                            && current_check
                                .map_or(false, |expected| *expected < current_count as u32);
                        return (!too_many_bads, bad_spring_read_count);
                    }

                    let bad_check = current_bad_spring == Some(&i);
                    if bad_check {
                        current_bad_spring = bad_spring_iter.next();
                        bad_spring_read_count += 1;
                    }
                    (bad_check, 1)
                }
            };

            if bad {
                current_count += increment;
                in_bad_lands = true;
            } else {
                if in_bad_lands {
                    let valid =
                        current_check.map_or(false, |expected| *expected == current_count as u32);
                    if !valid {
                        return (false, bad_spring_read_count);
                    }

                    current_count = 0;
                    current_check = check_iter.next();
                    in_bad_lands = false;
                }
            }
        }

        (true, bad_spring_read_count)
    }

    fn validate_check(&self, bad_springs: &[usize]) -> bool {
        let mut check_iter = self.check.iter();
        let mut current_check = check_iter.next();
        let mut bad_spring_iter = bad_springs.iter();
        let mut current_bad_spring = bad_spring_iter.next();
        let mut current_count = 0;
        let mut in_bad_lands = false;

        for (i, spring) in self.springs.iter().enumerate() {
            let (bad, increment) = match spring {
                Spring::Good => (false, 1),
                Spring::Bad(s) => (true, *s),
                Spring::Unknown => {
                    let bad_check = current_bad_spring == Some(&i);
                    if bad_check {
                        current_bad_spring = bad_spring_iter.next();
                    }
                    (bad_check, 1)
                }
            };

            if bad {
                current_count += increment;
                in_bad_lands = true;
            } else {
                if in_bad_lands {
                    let valid =
                        current_check.map_or(false, |expected| *expected == current_count as u32);
                    if !valid {
                        return false;
                    }
                    current_count = 0;
                    current_check = check_iter.next();
                    in_bad_lands = false;
                }
            }
        }

        if in_bad_lands {
            let valid = current_check.map_or(false, |expected| *expected == current_count as u32);
            if !valid {
                return false;
            }
            current_check = check_iter.next();
        }

        current_check == None
    }

    fn possibilities_count(&self) -> usize {
        let unknown_indexes: Vec<usize> = self
            .springs
            .iter()
            .enumerate()
            .filter_map(|(i, s)| (s == &Spring::Unknown).then_some(i))
            .collect();

        let total_bad_count: usize = self.check.iter().sum::<u32>() as usize;
        let known_bad_count: usize = self
            .springs
            .iter()
            .map(|s| match s {
                Spring::Bad(c) => *c,
                _ => 0,
            })
            .sum();

        let expected_ones: usize = total_bad_count - known_bad_count;

        let possibilities = count_choices(
            &unknown_indexes,
            &mut Vec::with_capacity(expected_ones),
            expected_ones,
            |bad_springs_prefix: &[usize]| self.validate_check_prefix(&bad_springs_prefix),
            |bad_springs: &[usize]| self.validate_check(&bad_springs),
        );

        dbg!(possibilities)
    }

    fn possibilities_count_alt(&self) -> usize {
        let working_row = WorkingSpringRow {
            springs: &self.springs,
            check: &self.check,
        };
        dbg!(working_row.possibilities_count())
    }
}

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)
    }
}

fn count_choices(
    possibilities: &[usize],
    prefix: &mut Vec<usize>,
    pulls: usize,
    prefix_filter: impl Fn(&[usize]) -> (bool, usize) + Copy,
    filter: impl Fn(&[usize]) -> bool + Copy,
) -> usize {
    let mut res = 0;
    if pulls == 0 {
        res = 1;
    } else if pulls == 1 {
        for suffix in possibilities {
            prefix.push(*suffix);
            if filter(&prefix) {
                res += 1;
            }
            prefix.pop();
        }
    } else {
        for (i, suffix) in possibilities.iter().enumerate() {
            prefix.push(*suffix);
            let (prefix_could_be_valid, prefix_filter_read_count) = prefix_filter(&prefix);
            if prefix_could_be_valid {
                res += count_choices(
                    &possibilities[i + 1..],
                    prefix,
                    pulls - 1,
                    prefix_filter,
                    filter,
                );
            }
            prefix.pop();

            if !prefix_could_be_valid && prefix_filter_read_count < prefix.len() {
                break;
            }
        }
    }
    res
}

impl<'a> WorkingSpringRow<'a> {
    fn possibilities_count(&self) -> usize {
        // this now definitely begins and ends with unknowns
        let optimized = match self.optimize() {
            Some(optimized) => optimized,
            None => {
                return 0;
            }
        };

        if optimized.check.len() == 1
            && optimized
                .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 optimized.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 >= optimized.check[0] as usize {
                        region - optimized.check[0] as usize + 1
                    } else {
                        0
                    }
                })
                .sum();
        }

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

        let valid_prefix_possibilities_count =
            if let Some(without_prefix) = optimized.trim_bad_prefix() {
                without_prefix.possibilities_count()
            } else {
                0
            };

        let non_prefix_possibilities_count = {
            let subset = WorkingSpringRow {
                springs: &optimized.springs[1..],
                check: &optimized.check,
            };
            subset.possibilities_count()
        };

        valid_prefix_possibilities_count + non_prefix_possibilities_count
    }

    fn trim_bad_prefix(&self) -> Option<WorkingSpringRow<'a>> {
        let first_check = self.check[0] as usize;
        let mut check_i = 0;
        let mut springs_i = 0;
        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(WorkingSpringRow {
                springs: &self.springs[new_springs_i..],
                check: &self.check[1..],
            })
        } else {
            None
        }
    }

    fn trim_bad_suffix(&self) -> Option<WorkingSpringRow<'a>> {
        let last_check = self.check[self.check.len() - 1] as usize;
        let mut check_i = 0;
        let mut new_springs_len = self.springs.len();
        while check_i < last_check {
            if new_springs_len == 0 {
                return None;
            }
            match self.springs[new_springs_len - 1] {
                Spring::Unknown => {
                    check_i += 1;
                }
                Spring::Bad(inc) => {
                    check_i += inc;
                }
                Spring::Good => {
                    return None;
                }
            }
            new_springs_len -= 1;
        }
        let is_boundary = new_springs_len == 0;
        let has_extra_good_or_unknown_to_trim =
            !is_boundary && !matches!(self.springs[new_springs_len - 1], Spring::Bad(_));
        let valid_suffix =
            check_i == last_check && (new_springs_len == 0 || has_extra_good_or_unknown_to_trim);
        if valid_suffix {
            let new_springs_len = if has_extra_good_or_unknown_to_trim {
                new_springs_len - 1
            } else {
                new_springs_len
            };
            Some(WorkingSpringRow {
                springs: &self.springs[..new_springs_len],
                check: &self.check[..self.check.len() - 1],
            })
        } else {
            None
        }
    }

    fn optimize(&'a self) -> Option<WorkingSpringRow<'a>> {
        let mut optimized = self.clone();
        while optimized.springs.len() > 0 && optimized.check.len() > 0 {
            match optimized.springs[0] {
                Spring::Good => {
                    optimized.springs = &optimized.springs[1..];
                }
                Spring::Bad(s) if s > optimized.check[0] as usize => {
                    return None;
                }
                Spring::Bad(s) if s <= optimized.check[0] as usize => {
                    let trimmed = match optimized.trim_bad_prefix() {
                        Some(opt) => opt,
                        None => return None,
                    };
                    optimized = trimmed;
                }
                Spring::Bad(_) => unreachable!(),
                Spring::Unknown => {
                    break;
                }
            }
        }

        while optimized.springs.len() > 0 && optimized.check.len() > 0 {
            match optimized.springs[optimized.springs.len() - 1] {
                Spring::Good => {
                    optimized.springs = &optimized.springs[..optimized.springs.len() - 1];
                }
                Spring::Bad(s) if s > optimized.check[optimized.check.len() - 1] as usize => {
                    return None;
                }
                Spring::Bad(s) if s <= optimized.check[optimized.check.len() - 1] as usize => {
                    let trimmed = match optimized.trim_bad_suffix() {
                        Some(opt) => opt,
                        None => return None,
                    };
                    optimized = trimmed;
                }
                Spring::Bad(_) => unreachable!(),
                Spring::Unknown => {
                    break;
                }
            }
        }

        if optimized.springs.len() == 0 && optimized.check.len() != 0 {
            None
        } else {
            Some(optimized)
        }
    }
}