use aoc_runner_derive::{aoc, aoc_generator};
use std::collections::HashSet;

#[derive(Debug, Ord, PartialOrd, Eq, PartialEq, Clone)]
pub enum Inst {
    Acc { value: i64 },
    Jmp { offset: i64 },
    Nop { unused: i64 },
}
#[derive(Debug, Ord, PartialOrd, Eq, PartialEq)]
pub struct CPU {
    ip: usize,
    program: Vec<Inst>,
    acc: i64,
}

impl Inst {
    fn flip(&mut self) {
        if let Some(new_inst) = match &self {
            Inst::Jmp { offset } => Some(Inst::Nop { unused: *offset }),
            Inst::Nop { unused } => Some(Inst::Jmp { offset: *unused }),
            _ => None,
        } {
            *self = new_inst;
        };
    }
}

impl From<&str> for Inst {
    fn from(inst: &str) -> Self {
        let mut inst = inst.split_whitespace();
        let (inst, arg) = (
            inst.next().unwrap(),
            inst.next().map(|v| v.parse::<i64>().unwrap()),
        );
        match inst {
            "jmp" => Inst::Jmp {
                offset: arg.unwrap(),
            },
            "acc" => Inst::Acc {
                value: arg.unwrap(),
            },
            "nop" => Inst::Nop {
                unused: arg.unwrap(),
            },
            other => panic!("Invalid instruction: {}", other),
        }
    }
}

impl CPU {
    fn new(pgm: &[Inst]) -> Self {
        Self {
            ip: 0,
            program: pgm.to_vec(),
            acc: 0,
        }
    }

    fn flip(&mut self, idx: usize) {
        self.program[idx].flip();
    }

    fn step(&mut self) -> bool {
        // println!("{}: {:?} [Acc: {}]",self.ip,inst, self.acc);
        match self.program.get(self.ip) {
            Some(Inst::Acc { value }) => {
                self.acc += value;
            }
            Some(Inst::Jmp { offset }) => {
                self.ip = ((self.ip as i64) + (offset - 1)) as usize;
            }
            Some(Inst::Nop { .. }) => (),
            None => {
                return false;
            }
        }
        self.ip += 1;
        return true;
    }
}

#[aoc_generator(day8)]
#[inline(always)]
pub fn input_generator(input: &str) -> Vec<Inst> {
    input.trim().lines().map(|l| Inst::from(l.trim())).collect()
}

#[aoc(day8, part1)]
#[inline(always)]
pub fn solve_part1(input: &[Inst]) -> usize {
    let mut seen = HashSet::new();
    let mut cpu = CPU::new(input);
    loop {
        seen.insert(cpu.ip);
        cpu.step();
        if seen.contains(&cpu.ip) {
            break;
        }
    }
    return cpu.acc as usize;
}

#[aoc(day8, part2)]
#[inline(always)]
pub fn solve_part2(input: &[Inst]) -> usize {
    for (idx, inst) in input.iter().enumerate() {
        match inst {
            Inst::Nop { .. } | Inst::Jmp { .. } => {
                let mut seen = HashSet::new();
                let mut cpu = CPU::new(input);
                cpu.flip(idx);
                loop {
                    seen.insert(cpu.ip);
                    if !cpu.step() {
                        return cpu.acc as usize;
                    };
                    if seen.contains(&cpu.ip) {
                        break;
                    }
                }
            }
            _ => (),
        }
    }
    panic!("Failed to solve part 2");
}