use std::{
    collections::{HashMap, HashSet},
    process::exit,
    thread,
    time::Duration,
};

use advent_of_code_2024::{make_main, SResult};
use regex::Regex;
use threadpool::ThreadPool;

make_main!();

struct State<'a> {
    registers: [usize; 3],
    pointer: usize,
    program: &'a [usize],
}

impl<'a> State<'a> {
    fn new(program: &'a [usize], registers: [usize; 3]) -> Self {
        Self {
            registers,
            pointer: 0,
            program,
        }
    }

    fn run_to_end(&mut self) -> Vec<usize> {
        let mut results = vec![];
        let mut i = 0;
        while self.pointer < self.program.len() {
            if let Some(val) = self.opcode() {
                if val != self.program[i] {
                    break;
                }
                results.push(val);
                i += 1;
                if i == self.program.len() {
                    break;
                }
            }
        }
        results
    }

    // Processes opcode and returns the program pointer
    fn opcode(&mut self) -> Option<usize> {
        let opcode = self.program[self.pointer];
        let operand = self.program[self.pointer + 1];
        match opcode {
            0 => {
                self.registers[0] = self.registers[0] / (usize::pow(2, self.combo(operand) as u32));
                self.pointer += 2
            }
            1 => {
                self.registers[1] = self.registers[1] ^ operand;
                self.pointer += 2
            }
            2 => {
                self.registers[1] = self.combo(operand) % 8;
                self.pointer += 2
            }
            3 => {
                if self.registers[0] == 0 {
                    self.pointer += 2
                } else {
                    self.pointer = operand;
                }
            }
            4 => {
                self.registers[1] = self.registers[1] ^ self.registers[2];
                self.pointer += 2
            }
            5 => {
                self.pointer += 2;
                return Some(self.combo(operand) % 8);
            }
            6 => {
                self.registers[1] = self.registers[0] / (usize::pow(2, self.combo(operand) as u32));
                self.pointer += 2
            }
            7 => {
                self.registers[2] = self.registers[0] / (usize::pow(2, self.combo(operand) as u32));
                self.pointer += 2
            }
            _ => unimplemented!(),
        };
        None
    }
    fn combo(&self, val: usize) -> usize {
        match val {
            0..4 => val,
            4 => self.registers[0],
            5 => self.registers[1],
            6 => self.registers[2],
            _ => unimplemented!(),
        }
    }
}

fn solve(lines: Vec<String>) -> SResult<(usize, usize)> {
    let mut registers: HashMap<String, usize> = HashMap::default();

    let regex = Regex::new(r"Register (\w+): (\d+)").unwrap();
    let mut i = 0;
    while i < lines.len() && lines[i].trim() != "" {
        let caps = regex.captures(&lines[i]).unwrap();
        i += 1;
        registers.insert(caps[1].into(), caps[2].parse().unwrap());
    }
    i += 1;
    let program: Vec<usize> = lines[i]
        .strip_prefix("Program: ")
        .unwrap()
        .split(",")
        .map(|s| s.parse().unwrap())
        .collect();
    let values =
        State::new(&program, [registers["A"], registers["B"], registers["C"]]).run_to_end();
    for (i, n) in values.iter().enumerate() {
        if i > 0 {
            print!(",");
        }
        print!("{}", n);
    }
    println!();

    let pool = ThreadPool::new(20);
    //let mut n = 3536000000;
    let mut n = 0;
    let mut prev_best = 0;
    let program: &'static mut [usize] = program.leak();
    loop {
        let values = State::new(&program, [n, 0, 0]).run_to_end();
        if values.len() > prev_best {
            println!("{} => {:?}", n, values);
            prev_best = values.len();
        }
        while &values == program {
            println!("finished at {}", n);
            //exit(n as i32);
        }
        n += 1;
        if n % 1000000 == 0 {
            println!("n = {}", n);
        }
    }
    Ok((values.len(), n))
}

#[cfg(test)]
mod tests {
    use advent_of_code_2024::input;

    use super::*;
    #[test]
    fn d17i1() {
        let strings: Vec<String> = input!("d17i1.txt");
        let got = solve(strings).unwrap();
        assert_eq!(got, (4, 117440));
    }
    #[test]
    fn d17i2() {
        return;
        let strings: Vec<String> = input!("d17i2.txt");
        let got = solve(strings).unwrap();
        assert_eq!(got, (4, 117440));
    }
}