use itertools::Itertools;
use regex::Regex;
use std::fmt;

struct Operator {
    name: char,
    f: Box<dyn Fn(u64, u64) -> u64>,
}

impl fmt::Display for Operator {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.name)
    }
}

#[allow(dead_code)]
fn op_chain_to_string(chain: &Vec<&Operator>) -> String {
    format!("[{}]", chain.iter().map(|o| o.to_string()).join(", "))
}

#[derive(Debug)]
struct Equation {
    result: u64,
    values: Vec<u64>,
}

impl Equation {
    fn parse(s: &str) -> Self {
        let re = Regex::new("^([0-9]+): ((?:[0-9]+ ?)+)$").unwrap();
        let caps = re.captures(s).unwrap();
        let (_full, [result, values]) = caps.extract();
        Equation {
            result: result.parse().unwrap(),
            values: values.split(' ').map(|n| n.parse().unwrap()).collect(),
        }
    }

    fn is_solution(&self, chain: &Vec<&Operator>) -> bool {
        let mut chain_it = chain.iter();
        let mut self_it = self.values.iter();
        let mut acc: u64 = *self_it.next().unwrap();
        for value in self_it {
            let Some(op) = chain_it.next() else {
                panic!("Impossible")
            };
            acc = (op.f)(acc, *value);
            if acc > self.result {
                return false;
            }
        }
        // // TODO why does this not work?
        // self.values.iter()
        //     .reduce(|v1, v2| {
        //         let Some(op) = chain_it.next() else { panic!("Impossible") };
        //         (op.f)(v1, v2)
        //     });
        acc == self.result
    }

    fn find_solution<'a>(&self, ops: &'a [Operator]) -> Option<Vec<&'a Operator>> {
        // println!("Finding solution for {self:?}...");
        let op_chains = itertools::repeat_n(ops, self.values.len() - 1).multi_cartesian_product();
        for chain in op_chains {
            // print!("  Trying chain {}... ", op_chain_to_string(&chain));
            if self.is_solution(&chain) {
                // println!("YES!");
                return Some(chain.clone());
            }
            // println!("nope");
        }
        None
    }

    fn has_solution(&self, ops: &[Operator]) -> bool {
        self.find_solution(ops).is_some()
    }
}

fn sum_solutions(input: &str, ops: &[Operator]) -> u64 {
    let equations: Vec<Equation> = input.lines().map(Equation::parse).collect();
    // println!("Equations: {equations:?}");
    equations
        .iter()
        .filter(|e| e.has_solution(ops))
        .map(|e| e.result)
        .sum()
}

fn p1(input: &str) -> String {
    let ops: [Operator; 2] = [
        Operator {
            name: '+',
            f: Box::new(|x, y| x + y),
        },
        Operator {
            name: '*',
            f: Box::new(|x, y| x * y),
        },
    ];

    sum_solutions(input, &ops).to_string()
}

fn p2(input: &str) -> String {
    let ops: [Operator; 3] = [
        Operator {
            name: '+',
            f: Box::new(|x, y| x + y),
        },
        Operator {
            name: '*',
            f: Box::new(|x, y| x * y),
        },
        Operator {
            name: '|',
            f: Box::new(|x, y| {
                format!("{}{}", x, y).parse().unwrap()
            }),
        },
    ];

    sum_solutions(input, &ops).to_string()
}

fn main() {
    aoc2024::run_day("7", p1, p2);
}