use std::{collections::HashMap, error::Error};

use itertools::Itertools;

pub fn process_part1(input: &str) -> (String, HashMap<&str, u32>) {
    let (registers, program) = input.split_once("\n\n").unwrap();
    let mut registers = parse_registers(registers);
    let (_, instructions) = program.split_once(": ").unwrap();
    let mut out = Vec::new();
    let instructions = instructions
        .lines()
        .map(|line| {
            line.split(",")
                .collect_vec()
                .chunks(2)
                .map(|chunk| (chunk[0], chunk[1]))
                .collect_vec()
        })
        .collect_vec()
        .concat();

    let mut instruction_pointer: u32 = 0;

    while let Some((opcode, operand)) = instructions.get(instruction_pointer as usize) {
        let opcode = OpCodes::try_from(*opcode).unwrap();
        if let Some((output, skip)) = opcode.exec(operand, &mut registers) {
            if skip {
                instruction_pointer = 0;
            } else {
                out.push(output);
                instruction_pointer += 1;
            }
        } else {
            instruction_pointer += 1;
        }
    }
    (out.into_iter().join(","), registers)
}

pub fn process_part2(input: &str) -> u32 {
    let (registers, program) = input.split_once("\n\n").unwrap();
    let mut registers = parse_registers(registers);
    let (_, instructions) = program.split_once(": ").unwrap();
    let instructions = instructions
        .lines()
        .map(|line| {
            line.split(",")
                .collect_vec()
                .chunks(2)
                .map(|chunk| (chunk[0], chunk[1]))
                .collect_vec()
        })
        .collect_vec()
        .concat();
    let orig = instructions
        .iter()
        .map(|(opcode, operand)| {
            [
                opcode.parse::<u32>().unwrap(),
                operand.parse::<u32>().unwrap(),
            ]
        })
        .collect_vec()
        .concat();
    loop {
        let out = exec_program(instructions.clone(), &mut registers);
        if out == orig {
            break;
        }
    }

    *registers.get("A").unwrap()
}

fn exec_program(instructions: Vec<(&str, &str)>, registers: &mut HashMap<&str, u32>) -> Vec<u32> {
    let mut out = Vec::new();
    let mut instruction_pointer: u32 = 0;

    while let Some((opcode, operand)) = instructions.get(instruction_pointer as usize) {
        let opcode = OpCodes::try_from(*opcode).unwrap();
        if let Some((output, skip)) = opcode.exec(operand, registers) {
            if skip {
                instruction_pointer = 0;
            } else {
                out.push(output);
                instruction_pointer += 1;
            }
        } else {
            instruction_pointer += 1;
        }
    }
    out
}

fn parse_registers(input: &str) -> HashMap<&str, u32> {
    let mut registers = HashMap::new();
    input.lines().for_each(|line| {
        let (register, value) = line.split_once(": ").unwrap();
        let (_, register) = register.split_once(" ").unwrap();
        let value = value.parse().unwrap();
        registers.insert(register, value);
    });
    registers
}

#[derive(Debug)]
enum OpCodes {
    Adv,
    Bxl,
    Bst,
    Jnz,
    Bxc,
    Out,
    Bdv,
    Cdv,
}

impl OpCodes {
    fn exec(&self, operand: &str, registers: &mut HashMap<&str, u32>) -> Option<(u32, bool)> {
        let a = *registers.get("A").unwrap();
        let b = *registers.get("B").unwrap();
        let c = *registers.get("C").unwrap();
        let operand: u32 = operand.parse().unwrap();
        let combo = if operand <= 3 {
            operand
        } else if operand == 4 {
            a
        } else if operand == 5 {
            b
        } else if operand == 6 {
            c
        } else {
            0
        };
        match self {
            OpCodes::Adv => {
                //println!(
                //    "a = a / 2 ** combo = {a} / 2 ** {combo} = {}",
                //    a / 2_u32.pow(combo)
                //);
                registers.insert("A", a / 2_u32.pow(combo));
                None
            }
            OpCodes::Bxl => {
                //println!("b = b xor operand = {b} xor {operand} = {}", b ^ operand);
                registers.insert("B", b ^ operand);
                None
            }
            OpCodes::Bst => {
                //println!("b = combo mod 8 = {combo} mod 8 = {}", combo % 8);
                registers.insert("B", combo % 8);
                None
            }
            OpCodes::Jnz => {
                if a != 0 {
                    //println!("jump {operand}");
                    Some((operand, true))
                } else {
                    //println!("don't jump");
                    None
                }
            }
            OpCodes::Bxc => {
                //println!("b = b xor c = {b} xor {c} = {}", b ^ c);
                registers.insert("B", b ^ c);
                None
            }
            OpCodes::Out => {
                //println!("out {combo} % 8 = {}", combo % 8);
                Some((combo % 8, false))
            }
            OpCodes::Bdv => {
                //println!(
                //    "b = a / 2 ** combo = {a} / 2 ** {combo} = {}",
                //    a / 2_u32.pow(combo)
                //);
                registers.insert("B", a / 2_u32.pow(combo));
                None
            }
            OpCodes::Cdv => {
                //println!(
                //    "c = a / 2 ** combo = {a} / 2 ** {combo} = {}",
                //    a / 2_u32.pow(combo)
                //);
                registers.insert("C", a / 2_u32.pow(combo));
                None
            }
        }
    }
}

impl TryFrom<&str> for OpCodes {
    type Error = Box<dyn Error>;

    fn try_from(value: &str) -> std::result::Result<OpCodes, Box<dyn Error>> {
        match value {
            "0" => Ok(Self::Adv),
            "1" => Ok(Self::Bxl),
            "2" => Ok(Self::Bst),
            "3" => Ok(Self::Jnz),
            "4" => Ok(Self::Bxc),
            "5" => Ok(Self::Out),
            "6" => Ok(Self::Bdv),
            "7" => Ok(Self::Cdv),
            _ => Err(Box::from(format!("{value} is not a valid OpCode"))),
        }
    }
}

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

    const INPUT_MAIN: &str = "Register A: 729
Register B: 0
Register C: 0

Program: 0,1,5,4,3,0";

    const INPUT_SMALL_1: &str = "Register A: 0
Register B: 0
Register C: 9

Program: 2,6";

    const INPUT_SMALL_2: &str = "Register A: 10
Register B: 0
Register C: 0

Program: 5,0,5,1,5,4";

    const INPUT_SMALL_3: &str = "Register A: 2024
Register B: 0
Register C: 0

Program: 0,1,5,4,3,0";

    const INPUT_SMALL_4: &str = "Register A: 0
Register B: 29
Register C: 0

Program: 1,7";

    const INPUT_SMALL_5: &str = "Register A: 0
Register B: 2024
Register C: 43690

Program: 4,0";

    const INPUT_COPY: &str = "Register A: 2024
Register B: 0
Register C: 0

Program: 0,3,5,4,3,0";

    #[test]
    fn part1_main() {
        let result = process_part1(INPUT_MAIN);
        assert_eq!(result.0, "4,6,3,5,6,3,5,2,1,0".to_string());
    }

    #[test]
    fn part1_small_1() {
        let result = process_part1(INPUT_SMALL_1);
        assert_eq!(result.1.get("B").unwrap(), &1);
    }

    #[test]
    fn part1_small_2() {
        let result = process_part1(INPUT_SMALL_2);
        assert_eq!(result.0, "0,1,2".to_string());
    }

    #[test]
    fn part1_small_3() {
        let result = process_part1(INPUT_SMALL_3);
        assert_eq!(result.1.get("A").unwrap(), &0);
        assert_eq!(result.0, "4,2,5,6,7,7,7,7,3,1,0".to_string());
    }

    #[test]
    fn part1_small_4() {
        let result = process_part1(INPUT_SMALL_4);
        assert_eq!(result.1.get("B").unwrap(), &26);
    }

    #[test]
    fn part1_small_5() {
        let result = process_part1(INPUT_SMALL_5);
        assert_eq!(result.1.get("B").unwrap(), &44354);
    }

    #[test]
    fn part2() {
        let result = process_part2(INPUT_COPY);
        assert_eq!(result, 117440);
    }
}