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

// Compared to day 12, for day 23 it makes more sense to abstract
// over operands in this way; they could be registers or literals.
#[derive(Copy, Clone, Debug)]
enum Arg<'a> {
    Reg(&'a str),
    Lit(isize),
}

// It's convenient to use the `From' trait for parsing operands
impl<'a> From<&'a str> for Arg<'a> {
    fn from(s: &'a str) -> Self {
        if s.chars().all(|c| c.is_ascii_digit() || c == '-') { // yes I know
            Self::Lit(s.parse().unwrap())
        } else {
            Self::Reg(s)
        }
    }
}

#[derive(Copy, Clone, Debug)]
enum Oper<'a> {
    Inc(Arg<'a>),
    Dec(Arg<'a>),
    Tgl(Arg<'a>),
    Cpy(Arg<'a>, Arg<'a>),
    Jnz(Arg<'a>, Arg<'a>),
}

fn parse(lines: Vec<&str>) -> Vec<Oper> {
    let mut insts = Vec::new();

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

        if words[0] == "inc" {
            insts.push(Oper::Inc(Arg::from(words[1])));
        } else if words[0] == "dec" {
            insts.push(Oper::Dec(Arg::from(words[1])));
        } else if words[0] == "tgl" {
            insts.push(Oper::Tgl(Arg::from(words[1])));
        } else if words[0] == "cpy" {
            insts.push(Oper::Cpy(Arg::from(words[1]), Arg::from(words[2])));
        } else if words[0] == "jnz" {
            insts.push(Oper::Jnz(Arg::from(words[1]), Arg::from(words[2])));
        }
    }

    insts
}

fn execute(program: &Vec<Oper>, init: Vec<(&str, isize)>) -> isize {
    // Make a local mutable copy of the program
    let mut insts = program.clone();

    // `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() {
        match insts[ip] {
            Oper::Inc(arg) => {
                if let Arg::Reg(dst) = arg {
                    *state.entry(dst).or_insert(0) += 1;
                }
            },
            Oper::Dec(arg) => {
                if let Arg::Reg(dst) = arg {
                    *state.entry(dst).or_insert(0) -= 1;
                }
            },
            Oper::Tgl(arg) => {
                let off = match arg {
                    Arg::Reg(reg) => state.get(&reg).copied().unwrap_or(0),
                    Arg::Lit(lit) => lit,
                };
                // Calculate index of target instruction
                let t = ((ip as isize) + off) as usize;
                if t < insts.len() {
                    insts[t] = match insts[t] {
                        Oper::Inc(a)      => Oper::Dec(a),
                        Oper::Dec(a)      => Oper::Inc(a),
                        Oper::Tgl(a)      => Oper::Inc(a),
                        Oper::Cpy(a1, a2) => Oper::Jnz(a1, a2),
                        Oper::Jnz(a1, a2) => Oper::Cpy(a1, a2),
                    };
                }
            },
            Oper::Cpy(arg1, arg2) => {
                // If destination is a literal (due to `Tgl'), ignore this `Cpy'
                if let Arg::Reg(dst) = arg2 {
                    let val = match arg1 {
                        Arg::Reg(src) => state.get(&src).copied().unwrap_or(0),
                        Arg::Lit(lit) => lit,
                    };
                    *state.entry(dst).or_insert(0) = val;
                }
            },
            Oper::Jnz(arg1, arg2) => {
                let val = match arg1 {
                    Arg::Reg(reg) => state.get(&reg).copied().unwrap_or(0),
                    Arg::Lit(lit) => lit,
                };
                // Should the branch be taken? If yes:
                if val != 0 {
                    let off = match arg2 {
                        Arg::Reg(reg) => state.get(&reg).copied().unwrap_or(0),
                        Arg::Lit(lit) => lit,
                    };
                    // Subtract 1 to compensate for `ip += 1;' below
                    ip = ((ip as isize) + off - 1) as usize;
                }
            },
        };

        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 14065 for my input
    println!("Part 1 solution: {}", execute(&insts, vec![("a", 7)]));

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

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

    #[test]
    fn part1_example1() {
        let lines = vec![
            "cpy 2 a", "tgl a", "tgl a", "tgl a", "cpy 1 a", "dec a", "dec a"
        ];
        let insts = parse(lines);
        assert_eq!(execute(&insts, vec![]), 3);
    }
}