#[path = "../lib.rs"]
mod lib;

fn parse_input() -> (usize, Vec<usize>) {
    let mut a = 0;
    let mut program = vec![];

    for line in lib::iter_input() {
        match line.split_once(": ") {
            Some(("Register A", v)) => {
                a = v.parse::<usize>().unwrap();
            }
            Some(("Register B", v)) => {
                assert!(v == "0")
            }
            Some(("Register C", v)) => {
                assert!(v == "0")
            }
            Some(("Program", v)) => {
                program = v.split(',').map(|c| c.parse::<usize>().unwrap()).collect();
            }
            _ => {}
        };
    }

    return (a, program);
}

fn decompile_combo(value: usize) -> String {
    return match value {
        0 | 1 | 2 | 3 => value.to_string(),
        4 => "A".to_string(),
        5 => "B".to_string(),
        6 => "C".to_string(),
        _ => unreachable!(),
    };
}

fn decompile(program: &Vec<usize>) {
    let mut i = 0;
    while i < program.len() {
        match program[i] {
            0 => {
                println!("A = A >> {}", decompile_combo(program[i + 1]));
            }
            1 => {
                println!("B ^= {}", program[i + 1]);
            }
            2 => {
                println!("B = {} % 8", decompile_combo(program[i + 1]));
            }
            3 => {
                println!("if A != 0: jump({})", program[i + 1]);
            }
            4 => {
                println!("B ^= C");
            }
            5 => {
                println!("output({} % 8)", decompile_combo(program[i + 1]))
            }
            6 => {
                println!("B = A >> {}", decompile_combo(program[i + 1]));
            }
            7 => {
                println!("C = A >> {}", decompile_combo(program[i + 1]));
            }
            _ => unreachable!(),
        };
        i += 2;
    }
    println!("");
}

fn get_combo(value: usize, reg: &[usize; 3]) -> usize {
    return match value {
        0 | 1 | 2 | 3 => value,
        4 => reg[0],
        5 => reg[1],
        6 => reg[2],
        _ => unreachable!(),
    };
}

fn exec(i: usize, reg: &mut [usize; 3], program: &Vec<usize>, output: &mut Vec<usize>) -> usize {
    match program[i] {
        0 => {
            reg[0] >>= get_combo(program[i + 1], reg);
        }
        1 => {
            reg[1] ^= program[i + 1];
        }
        2 => {
            reg[1] = get_combo(program[i + 1], reg) % 8;
        }
        3 => {
            if reg[0] != 0 {
                return program[i + 1];
            }
        }
        4 => {
            reg[1] ^= reg[2];
        }
        5 => {
            output.push(get_combo(program[i + 1], reg) % 8);
        }
        6 => {
            reg[1] = reg[0] >> get_combo(program[i + 1], reg);
        }
        7 => {
            reg[2] = reg[0] >> get_combo(program[i + 1], reg);
        }
        _ => unreachable!(),
    };
    return i + 2;
}

fn exec_cycle(program: &Vec<usize>, a: usize) -> usize {
    let mut reg = [a, 0, 0];
    let mut output = vec![];
    let mut i = 0;
    while output.len() == 0 {
        i = exec(i, &mut reg, program, &mut output);
    }
    return output[0];
}

fn reverse(program: &Vec<usize>, a: usize, i: usize) -> Option<usize> {
    // I noticed some properties about the input
    //
    // - it ends with a conditional jump to the start of the program
    // - that is the only jump
    // - a single value is output on every iteration
    // - that value can be derived from A, without knowing the initial B and C
    // - A is shifted by 3 bits in every iteration and not otherwise changed
    //
    // In other words: The program iterate through the 3-bit blocks in A and
    // outputs some values for each one. These values may also depend on the
    // higher bits, but not on the lower. This means that we can construct an
    // initial value for A by starting at the highest bits (that will generate
    // the last output).

    if i == program.len() {
        return Some(a);
    }

    for j in 0..8 {
        let next = (a << 3) | j;
        if exec_cycle(program, next) == program[program.len() - i - 1] {
            if let Some(result) = reverse(program, next, i + 1) {
                return Some(result);
            }
        }
    }

    return None;
}

fn part1(program: &Vec<usize>, a: usize) -> String {
    let mut reg = [a, 0, 0];
    let mut output = vec![];
    let mut i = 0;

    while i < program.len() {
        i = exec(i, &mut reg, &program, &mut output);
    }

    return output
        .iter()
        .map(|&x| x.to_string())
        .collect::<Vec<String>>()
        .join(",");
}

fn part2(program: &Vec<usize>) -> usize {
    return reverse(program, 0, 0).unwrap();
}

fn main() {
    let (a, program) = parse_input();

    decompile(&program);

    println!("part1: {}", part1(&program, a));
    println!("part2: {}", part2(&program));
}