use std::collections::HashSet;
use std::error::Error;
use std::fs::File;
use std::io::{self, BufRead};
use std::vec::Vec;

enum Instruction {
    Nop(i32),
    Acc(i32),
    Jmp(i32),
}

struct Node {
    inst: Instruction,
    jmp_srcs: Vec<usize>,
    nop_srcs: Vec<usize>,
}

fn line_to_node(line: String) -> Node {
    let mut parts = line.split(" ");
    let opcode = parts.nth(0).unwrap();
    let arg: i32 = parts.nth(0).unwrap().parse().unwrap();

    match opcode {
        "nop" => Node {
            inst: Instruction::Nop(arg),
            jmp_srcs: Vec::new(),
            nop_srcs: Vec::new(),
        },
        "acc" => Node {
            inst: Instruction::Acc(arg),
            jmp_srcs: Vec::new(),
            nop_srcs: Vec::new(),
        },
        "jmp" => Node {
            inst: Instruction::Jmp(arg),
            jmp_srcs: Vec::new(),
            nop_srcs: Vec::new(),
        },
        _ => panic!("Unsupported Instruction: {}", opcode),
    }
}

fn update_source_locations(programm: &mut Vec<Node>) {
    for pc in 0..programm.len() {
        let node = programm.get(pc).unwrap();
        match node.inst {
            Instruction::Nop(arg) => {
                if let Some(target) = programm.get_mut((pc as i32 + arg) as usize) {
                    target.nop_srcs.push(pc);
                }
            }
            Instruction::Jmp(arg) => {
                if let Some(target) = programm.get_mut((pc as i32 + arg) as usize) {
                    target.jmp_srcs.push(pc);
                }
            }
            _ => {}
        }
    }
}

fn backtrack_acc(
    programm: &Vec<Node>,
    mut visted: HashSet<usize>,
    pc: usize,
    mut acc: i32,
    changed: bool,
) -> Option<i32> {
    if visted.contains(&pc) {
        return None;
    }

    let node = programm.get(pc)?;
    visted.insert(pc);

    if let Instruction::Acc(arg) = node.inst {
        acc += arg;
    }

    if pc == 0 {
        return Some(acc);
    }

    let prev_node = programm.get(pc - 1).unwrap();
    match prev_node.inst {
        Instruction::Jmp(_) => {
            if !changed {
                if let Some(res) = backtrack_acc(programm, visted.clone(), pc - 1, acc, true) {
                    return Some(res);
                };
            }
        }
        _ => {
            if let Some(res) = backtrack_acc(programm, visted.clone(), pc - 1, acc, changed) {
                return Some(res);
            };
        }
    };

    for src in node.jmp_srcs.iter() {
        if let Some(res) = backtrack_acc(programm, visted.clone(), *src, acc, changed) {
            return Some(res);
        };
    }

    if !changed {
        for src in node.nop_srcs.iter() {
            if let Some(res) = backtrack_acc(programm, visted.clone(), *src, acc, true) {
                return Some(res);
            };
        }
    }

    return None;
}

fn main() -> Result<(), Box<dyn Error>> {
    let file = File::open("inputs/day8.txt")?;
    let lines = io::BufReader::new(file).lines().map(|l| l.unwrap());

    let mut programm: Vec<Node> = lines.map(line_to_node).collect();
    update_source_locations(&mut programm);

    println!(
        "{:?}",
        backtrack_acc(&programm, HashSet::new(), programm.len() - 1, 0, false)
    );

    Ok(())
}