use std::ops::{Add, AddAssign, Mul};

use advent_of_code::{
    numberlength,
    strings::{convert_to_array, parsenumber},
};
use advent_of_code_macros::include_aoc;
use num::{pow::Pow, Integer};

include_aoc!(2024, 07);

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
enum Operation {
    Add,
    Multiply,
    Concat,
}

impl Operation {
    fn op<T: Integer + Add + Mul + Pow<u32, Output = T> + std::ops::DivAssign + AddAssign + From<u32> + Copy>(self, a: T, b: T) -> T {
        match self {
            Operation::Add => a + b,
            Operation::Multiply => a * b,
            Operation::Concat => {
                let ta = a * T::from(10).pow(numberlength(b));
                ta + b
            },
        }
    }
}

fn get_data(line: &str) -> (u64, Vec<u64>) {
    let (result, numbers) = line.split_once(':').unwrap();

    (parsenumber(result.trim()), Vec::from(convert_to_array::<_, _, ' '>(numbers, parsenumber)))
}

fn test_equation(result: u64, numbers: Vec<u64>, ops: Vec<Operation>) -> bool {
    let mut numberiter = numbers.iter();
    let mut opiter = ops.iter();
    let mut calcresult = numberiter.next().unwrap().clone();
    let mut next_number;
    debug_assert!(numbers.len() - 1 == ops.len());
    loop {
        next_number = numberiter.next().cloned();
        if next_number.is_none() {
            return calcresult == result;
        }
        calcresult = opiter.next().unwrap().op(calcresult, next_number.unwrap())
    }
}

fn get_ops(len: usize, conf: usize) -> Vec<Operation> {
    let mut shift = 0;
    let mut ops = Vec::with_capacity(len);
    loop {
        match (conf & (0b1 << shift)) >> shift {
            0 => ops.push(Operation::Add),
            1 => ops.push(Operation::Multiply),
            _ => unreachable!(),
        };
        shift += 1;
        if ops.len() == len {
            break;
        }
    }
    ops
}

fn get_ops_with_concat(len: u32) -> Vec<Vec<Operation>> {
    let mut ops = Vec::with_capacity(3u64.saturating_pow(len) as usize);
    for i in [
        Operation::Add,
        Operation::Concat,
        Operation::Multiply,
    ] {
        if len > 1 {
            for mut oplist in get_ops_with_concat(len - 1) {
                oplist.insert(0, i);
                ops.push(oplist);
            }
        } else {
            ops.push(vec![i]);
        }
    }
    ops
}

#[inline]
fn pow2(i: usize) -> usize {
    2_u64.pow(i as u32) as usize - 1
}

fn solution1(data: Vec<(u64, Vec<u64>)>) -> u64 {
    let mut output = 0;
    for line in data {
        for i in 0..=pow2(line.1.len()) {
            if test_equation(line.0, line.1.clone(), get_ops(line.1.len() - 1, i)) {
                output += line.0;
                break;
            }
        }
    }
    output
}

fn solution2(data: Vec<(u64, Vec<u64>)>) -> u128 {
    let mut output = 0;
    for line in data {
        for ops in get_ops_with_concat((line.1.len() - 1) as u32) {
            if test_equation(line.0, line.1.clone(), ops) {
                output += line.0 as u128;
                break;
            }
        }
    }
    output
}

fn main() {
    let data = convert_to_array::<_, _, '\n'>(DATA, get_data);
    println!("Found Calibration Total: {}", solution1(data.clone()));
    println!("Calibration total with Concat: {}", solution2(data.clone()));
}

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

    #[test]
    fn test_test_equation() {
        for c in [
            (
                190,
                Vec::from([
                    10, 19,
                ]),
                Vec::from([Operation::Multiply]),
                true,
            ),
            (
                3267,
                Vec::from([
                    81, 40, 27,
                ]),
                Vec::from([
                    Operation::Add,
                    Operation::Multiply,
                ]),
                true,
            ),
            (
                3267,
                Vec::from([
                    81, 40, 27,
                ]),
                Vec::from([
                    Operation::Multiply,
                    Operation::Add,
                ]),
                true,
            ),
            (
                292,
                Vec::from([
                    11, 6, 16, 20,
                ]),
                Vec::from([
                    Operation::Add,
                    Operation::Multiply,
                    Operation::Add,
                ]),
                true,
            ),
            (
                292,
                Vec::from([
                    11, 6, 16, 20,
                ]),
                Vec::from([
                    Operation::Add,
                    Operation::Multiply,
                    Operation::Multiply,
                ]),
                false,
            ),
        ] {
            assert_eq!(test_equation(c.0, c.1.into(), c.2.into()), c.3)
        }
    }
}