use std::fs;
use std::collections::HashMap;

enum Oper<'a> {
    IncReg,
    DecReg,
    CpyReg(&'a str),
    CpyLit(isize),
    JnzReg(&'a str, isize),
    JnzLit(isize,   isize),
}

fn parse(lines: Vec<&str>) -> Vec<(Oper, &str)> {
    // We emit tuples of type `(Oper, &str)', where the second field
    // is the destination register. Storing the destination this way,
    // instead of inside the `Oper' enum, helps keep `execute()' DRY.
    let mut insts = Vec::new();

    for line in lines {
        let words: Vec<&str> = line.split_ascii_whitespace().collect();

        if words[0] == "inc" {
            // words[1] is register name
            insts.push((Oper::IncReg, words[1]));
        } else if words[0] == "dec" {
            // words[1] is register name
            insts.push((Oper::DecReg, words[1]));
        } else if words[0] == "cpy" {
            // words[1] is source, can be literal number or register name
            // words[2] is destination register name
            if words[1].chars().all(|c| c.is_ascii_digit() || c == '-') { // yes I know
                let val = words[1].parse().unwrap();
                insts.push((Oper::CpyLit(val), words[2]));
            } else {
                insts.push((Oper::CpyReg(words[1]), words[2]));
            }
        } else if words[0] == "jnz" {
            // words[1] is value, can be literal number or register name
            // words[2] is offset to jump by if branch is taken
            let off = words[2].parse().unwrap();
            if words[1].chars().all(|c| c.is_ascii_digit() || c == '-') { // I know
                let val = words[1].parse().unwrap();
                insts.push((Oper::JnzLit(val, off), "_"));
            } else {
                insts.push((Oper::JnzReg(words[1], off), "_"));
            }
        }
    }

    insts
}

fn execute(insts: &Vec<(Oper, &str)>, init: Vec<(&str, isize)>) -> isize {
    // `init' contains initial register values
    let mut state: HashMap<&str, isize> = init.into_iter().collect();

    let mut ip = 0; // instruction pointer
    while ip < insts.len() {
        // Get destination register now, because we'll mess with `ip' later
        let dst = insts[ip].1;

        // Too early for mutable borrow of `state', so write to `tmp' instead
        let mut tmp = state.get(dst).copied().unwrap_or(0);
        match insts[ip].0 {
            Oper::IncReg => {
                tmp += 1;
            },
            Oper::DecReg => {
                tmp -= 1;
            },
            Oper::CpyLit(val) => {
                tmp = val;
            },
            Oper::CpyReg(src) => {
                let val = state.get(&src).copied().unwrap_or(0);
                tmp = val;
            },
            Oper::JnzLit(val, off) => {
                if val != 0 {
                    // Subtract 1 to compensate for `ip += 1;' below
                    ip = ((ip as isize) + off - 1) as usize;
                }
            },
            Oper::JnzReg(src, off) => {
                let val = state.get(&src).copied().unwrap_or(0);
                if val != 0 {
                    // Subtract 1 to compensate for `ip += 1;' below
                    ip = ((ip as isize) + off - 1) as usize;
                }
            },
        };

        // "_" is a dummy destination register used by the "jnz" instructions
        if dst != "_" {
            *state.entry(dst).or_insert(0) = tmp;
        }

        ip += 1;
    }

    state.get("a").copied().unwrap()
}

fn main() {
    // Read assembly from input text file and parse it
    let input = fs::read_to_string("input.txt").unwrap();
    let lines = input.lines().collect();
    let insts = parse(lines);

    // Part 1 gives 318117 for my input
    println!("Part 1 solution: {}", execute(&insts, vec![]));

    // Part 2 gives 9227771 for my input
    println!("Part 2 solution: {}", execute(&insts, vec![("c", 1)]));
}



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

    #[test]
    fn part1_example1() {
        let lines = vec!["cpy 41 a", "inc a", "inc a", "dec a", "jnz a 2", "dec a"];
        let insts = parse(lines);
        assert_eq!(execute(&insts, vec![]), 42);
    }
}