use std::str::FromStr;
use anyhow::Result;

#[derive(Debug, Clone, PartialEq)]
enum ISet {
    Acc,
    Jmp,
    Nop,
}

#[derive(Debug, Clone)]
struct Instruction {
    execd: bool,
    instr: ISet,
    arg: i64
}

#[derive(Debug, thiserror::Error)]
enum ParseError {
    #[error("Expected {0}")]
    Expected(&'static str)
}


impl FromStr for Instruction {
    type Err = ParseError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let mut tokens = s.lines()
            .next()
            .ok_or(ParseError::Expected("one line"))?
            .split_whitespace();
        match tokens.next().ok_or(ParseError::Expected("instruction"))? {
            "acc" => {
                let arg = tokens
                    .next()
                    .ok_or(ParseError::Expected("acc argument"))?
                    .parse()
                    .unwrap();
                Ok(Instruction {
                    execd: false,
                    instr: ISet::Acc,
                    arg,
                    
                })
            },
            "nop" => {
                let arg = tokens
                    .next()
                    .ok_or(ParseError::Expected("acc argument"))?
                    .parse()
                    .unwrap();
                Ok(Instruction {
                    execd: false,
                    instr: ISet::Nop,
                    arg, // For part 1 at least NOP args are useless
                })
            },
            "jmp" => {
                let arg = tokens
                    .next()
                    .ok_or(ParseError::Expected("acc argument"))?
                    .parse()
                    .unwrap();
                Ok(Instruction {
                    execd: false,
                    instr: ISet::Jmp,
                    arg,
                })
            },
            _ => Err(ParseError::Expected("acc, nop or jmp instruction"))
        }
    }
}

impl Instruction {
    fn exec(&mut self, pc: &mut usize, acc: &mut i64) {
        self.execd = true;
        match self.instr {
            ISet::Acc => {
                *acc += self.arg;
                *pc += 1;
            }
            ISet::Jmp => {
                // Expect this to panic hard if you underflow :3
                *pc = ((*pc as i64) + self.arg) as usize;
            }
            ISet::Nop => {
                *pc += 1;
            }
        }
    }
}

#[derive(Debug, Clone)]
struct Program {
    prog: Vec<Instruction>,
    pc: usize,
    acc: i64,
}

impl Program {
    fn load(s: &str) -> Result<Program, ParseError> {
        let loaded_program: Vec<Instruction> = s
            .lines()
            .map(|l| Instruction::from_str(l))
            .collect::<Result<Vec<Instruction>, ParseError>>()?;
        Ok(
            Self {
                prog: loaded_program,
                pc: 0,
                acc: 0,
            }
        )
    }

    fn reset(&mut self) {
        self.pc = 0;
        self.acc = 0;
        for mut i in self.prog.iter_mut() {
            i.execd = false;
        }
    }

    /// Returns whether execution finished :
    /// Normally (true) - executing one instruction after the last one
    /// Or abnormalaly (false) - looping forever or executing random 
    /// instructions
    fn exec_until_loop(&mut self) -> bool {
        let pc_stop = self.prog.len();
        loop {
            // Terminate on infinite loop
            if self.prog[self.pc].execd {
                println!("Infinite loop !");
                return false;
            }
            
            self.prog[self.pc].exec(&mut self.pc, &mut self.acc);

            // Terminate on success
            if self.pc == pc_stop {
                println!("Good pc !");
                return true;
            }
            // Terminate on corrupted pc
            else if !(0..=pc_stop).contains(&(self.pc)) {
                println!("Corrupted pc !");
                return false;
            }
        }
    }

    fn find_corrupted_instruction(&mut self) -> Option<i64> {
        // Test jmps -> nops
        self.reset();
        for j in (0..self.prog.len())
            .filter(|&i| self.prog[i].instr == ISet::Jmp)
        {
            println!("Switch jmp-nop at pos {}", j);
            let mut pclone = self.clone();
            pclone.prog[j].instr = ISet::Nop;
            if pclone.exec_until_loop() {
                return Some(pclone.acc);
            }
        }
        // Test nops -> jmps
        for n in (0..self.prog.len())
            .filter(|&i| self.prog[i].instr == ISet::Nop)
        {
            println!("Switch nop-jmp at pos {}", n);
            let mut pclone = self.clone();
            pclone.prog[n].instr = ISet::Jmp;
            if pclone.exec_until_loop() {
                return Some(pclone.acc);
            }
        }
        None
    }
}

fn main() {
    let input = include_str!("../data/input");
    let mut prog = Program::load(input).unwrap();
    prog.exec_until_loop();
    println!("Acc before program gets into second loop : {}", prog.acc);
    println!("Acc after the program is set to terminate : {}", prog.find_corrupted_instruction().unwrap());
}

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

    #[test]
    fn test_nop() {
        let input = r"nop +0";
        let mut prog = Program::load(input).unwrap();
        prog.exec_until_loop();
        assert_eq!(prog.acc, 0)
    }

    #[test]
    fn test_acc() {
        let input = r"acc +5";
        let mut prog = Program::load(input).unwrap();
        prog.exec_until_loop();
        assert_eq!(prog.acc, 5)
    }

    #[test]
    fn test_jmp() {
        let input = r"jmp +2
acc +1
acc +1";
        let mut prog = Program::load(input).unwrap();
        prog.exec_until_loop();
        assert_eq!(prog.acc, 1)
    }

    #[test]
    fn test_short_loop() {
        let input = r"acc +1
jmp +2
acc +1
acc +1
jmp -4";
        let mut prog = Program::load(input).unwrap();
        prog.exec_until_loop();
        assert_eq!(prog.acc, 2)
    }

    #[test]
    fn test_exec() {
        let input = r"nop +0
acc +1
jmp +4
acc +3
jmp -3
acc -99
acc +1
jmp -4
acc +6";
        let mut prog = Program::load(input).unwrap();
        prog.exec_until_loop();
        assert_eq!(prog.acc, 5)
    }
}