use aoc_runner_derive::{aoc, aoc_generator};
use itertools::Itertools;
use misc::POW10;
use std::cmp::{max, min};
use std::ops::RangeInclusive;

#[aoc_generator(day2)]
fn parse(input: &str) -> Vec<RangeInclusive<u64>> {
    input
        .split(',')
        .map(|r| r.split_once('-').unwrap())
        .map(|(start, end)| RangeInclusive::new(start.parse().unwrap(), end.parse().unwrap()))
        .collect()
}

#[aoc(day2, part1, Brute)]
fn part1(input: &[RangeInclusive<u64>]) -> u64 {
    let mut invalid_sum = 0;

    for r in input {
        for product in r.clone() {
            let prod_s = product.to_string();
            if !prod_s.len().is_multiple_of(2) {
                continue;
            }
            let (left, right) = prod_s.split_at(prod_s.len() / 2);
            if left == right {
                invalid_sum += product
            }
        }
    }
    invalid_sum
}

#[aoc(day2, part2, Brute)]
fn part2(input: &[RangeInclusive<u64>]) -> u64 {
    let mut invalid_sum = 0;
    for r in input {
        for product in r.clone() {
            let prod_s = product.to_string();
            for n in 1..=prod_s.len() / 2 {
                if !prod_s.len().is_multiple_of(n) {
                    continue;
                }
                if prod_s
                    .chars()
                    .chunks(n)
                    .into_iter()
                    .map(|c| c.collect::<String>())
                    .all_equal()
                {
                    invalid_sum += product;
                    break;
                }
            }
        }
    }
    invalid_sum
}

#[aoc(day2, part1, ArithmeticBrute)]
fn part1_arithmetic_brute(input: &[RangeInclusive<u64>]) -> u64 {
    let mut invalid_sum = 0;
    for r in input {
        // println!("Range: {:?}", r);
        for product in r.clone() {
            let n_digits = (product.ilog10() + 1) as usize;
            if !n_digits.is_multiple_of(2) {
                continue;
            }

            let decade = POW10[n_digits / 2];
            let lhs = product / decade;
            let rhs = product % decade;
            // println!("{product} {n_digits}: D:{decade} L:{lhs} R:{rhs}");
            if lhs == rhs {
                invalid_sum += product
            }
        }
    }
    invalid_sum
}

#[aoc(day2, part2, ArithmeticBrute)]
fn part2_arithmetic_brute(input: &[RangeInclusive<u64>]) -> u64 {
    let mut invalid_sum = 0;
    for r in input {
        for product in r.clone() {
            let n_digits = (product.ilog10() + 1) as usize;

            for n in 1..=n_digits / 2 {
                if !n_digits.is_multiple_of(n) {
                    continue;
                }
                let repetitions = n_digits / n;
                let decade = POW10[n_digits - n]; // get the power of 10 to split the leftmost n digits
                let lhs = product / decade;
                let remainder = product % decade;

                // for each repetition we multiply by 10^(rep * n) to add the appropriate zeros
                let expected_remainder = (0..repetitions - 1).map(|rep| lhs * POW10[rep * n]).sum();

                if remainder == expected_remainder {
                    invalid_sum += product;
                    break;
                }
            }
        }
    }
    invalid_sum
}

#[aoc(day2, part1, Clever)]
fn part1_clever(input: &[RangeInclusive<u64>]) -> u64 {
    let mut invalid_sum = 0;
    for r in input {
        let n_digits = (r.start().ilog10() + 1) as usize;

        let mut lhs = if !n_digits.is_multiple_of(2) {
            // If the number of digits is odd, we start our guess at the first possibility in the next decade
            POW10[n_digits / 2]
        } else {
            let start_decade = POW10[n_digits / 2];
            r.start() / start_decade
        };

        // we will guess the next repeater, and check if it's in range, if not we are done.
        loop {
            let repeater = lhs * POW10[(lhs.ilog10() + 1) as usize] + lhs;
            if repeater < *r.start() {
                lhs += 1
            } else if r.contains(&repeater) {
                invalid_sum += repeater;
                lhs += 1
            } else {
                break;
            }
        }
    }
    invalid_sum
}

const fn n_digits(n: &u64) -> usize {
    (n.ilog10() + 1) as usize
}

fn generate_repeaters(r: &RangeInclusive<u64>, n: usize) -> Vec<u64> {
    let mut invalids = Vec::new();
    for r in split_by_decade(r) {
        let n_digits = n_digits(r.start());
        if !n_digits.is_multiple_of(n) || n_digits < 2 {
            continue;
        }

        let repetitions = n_digits / n;
        let decade = POW10[n_digits - n];
        for lhs in (r.start() / decade)..=(r.end() / decade) {
            let repeater = (0..repetitions)
                .map(|rep| lhs * POW10[rep * n])
                .sum::<u64>();
            if r.contains(&repeater) {
                invalids.push(repeater)
            }
        }
    }
    invalids
}

fn split_by_decade(r: &RangeInclusive<u64>) -> impl Iterator<Item = RangeInclusive<u64>> {
    let (start, end) = (*r.start(), *r.end());
    (n_digits(r.start())..n_digits(r.end()) + 1)
        .map(move |nd| RangeInclusive::new(max(start, POW10[nd - 1]), min(end, POW10[nd] - 1)))
}

#[aoc(day2, part2, Clever)]
fn part2_clever(input: &[RangeInclusive<u64>]) -> u64 {
    input
        .iter()
        .flat_map(|r| (1..=n_digits(r.end()) / 2).map(|nd| generate_repeaters(r, nd)))
        .flatten()
        .unique()
        .sum()
}

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

    const EXAMPLE: &str = "11-22,95-115,998-1012,1188511880-1188511890,222220-222224,1698522-1698528,446443-446449,38593856-38593862,565653-565659,824824821-824824827,2121212118-2121212124";

    #[test]
    fn part1_example() {
        assert_eq!(part1(&parse(EXAMPLE)), 1227775554);
    }

    #[test]
    fn part2_example() {
        assert_eq!(part2(&parse(EXAMPLE)), 4174379265);
    }

    #[test]
    fn part1_arithmetic_example() {
        assert_eq!(part1_arithmetic_brute(&parse(EXAMPLE)), 1227775554);
    }
    #[test]
    fn part2_arithmetic_example() {
        assert_eq!(part2_arithmetic_brute(&parse(EXAMPLE)), 4174379265);
    }
    #[test]
    fn part1_clever_example() {
        assert_eq!(part1_clever(&parse(EXAMPLE)), 1227775554);
    }
    #[test]
    fn part2_clever_example() {
        assert_eq!(part2_clever(&parse(EXAMPLE)), 4174379265);
    }
}