use std;
use std::convert::TryInto;

fn get_modes(x: i32) -> (i32, i32, i32, i32) {
    (
        x as i32 % 100,
        (x as i32 / 100) % 10,
        (x as i32 / 1000) % 10,
        x as i32 / 10000,
    )
}

fn get_param(input: Vec<i32>, mode: i32, idx: usize) -> i32 {
    match mode {
        0 => {
            let new_idx: usize = input[idx as usize].try_into().unwrap();
            input[new_idx]
        }
        1 => input[idx],
        _ => panic!("invalid input"),
    }
}

fn set_param(input: &mut Vec<i32>, mode: i32, idx: i32, value: i32) -> () {
    match mode {
        0 => {
            let new_idx: usize = input[idx as usize].try_into().unwrap();
            input[new_idx] = value;
        }
        _ => {
            panic!("invalid mode for output");
        }
    }
}

pub fn compute_intcode_extended(p: &mut Vec<i32>, inp: Vec<i32>) -> (Vec<i32>, Vec<i32>) {
    let mut idx = 0;
    let mut output: Vec<i32> = Vec::new();
    let mut input = inp.into_iter();

    while idx < p.len() - 1 {
        let (opcode, m1, m2, m3) = get_modes(p[idx]);
        match opcode {
            1 => {
                let left = get_param(p.to_vec(), m1, idx + 1);
                let right = get_param(p.to_vec(), m2, idx + 2);
                let res = left + right;
                set_param(p, m3, idx as i32 + 3, res);
                idx += 4;
            }
            2 => {
                let left = get_param(p.to_vec(), m1, idx + 1);
                let right = get_param(p.to_vec(), m2, idx + 2);
                let res = left * right;
                set_param(p, m3, idx as i32 + 3, res);
                idx += 4;
            }
            3 => {
                let new_idx: usize = p[idx as usize + 1].try_into().unwrap();
                p[new_idx] = input.next().expect("no next input!");
                idx += 2;
            }
            4 => {
                let value = get_param(p.to_vec(), m1, idx + 1);
                output.push(value);
                idx += 2;
            }
            5 => {
                let left = get_param(p.to_vec(), m1, idx + 1);
                let right = get_param(p.to_vec(), m2, idx + 2);
                if left != 0 {
                    idx = right as usize;
                } else {
                    idx += 3;
                }
            }
            6 => {
                let left = get_param(p.to_vec(), m1, idx + 1);
                let right = get_param(p.to_vec(), m2, idx + 2);
                if left == 0 {
                    idx = right as usize;
                } else {
                    idx += 3;
                }
            }
            7 => {
                let left = get_param(p.to_vec(), m1, idx + 1);
                let right = get_param(p.to_vec(), m2, idx + 2);
                set_param(p, m3, idx as i32 + 3, (left < right) as i32);
                idx += 4;
            }
            8 => {
                let left = get_param(p.to_vec(), m1, idx + 1);
                let right = get_param(p.to_vec(), m2, idx + 2);
                set_param(p, m3, idx as i32 + 3, (left == right) as i32);
                idx += 4;
            }
            99 => {
                break;
            }
            _ => panic!("invalid opcode: {}", opcode),
        }
    }
    (p.to_vec(), output.to_vec())
}

pub fn main() -> std::io::Result<()> {
    let input: Vec<i32> = std::include_str!("../day5-input.txt")
        .trim()
        .split(",")
        .map(|x| x.parse().ok().expect(&format!("{} not a number", x)))
        .collect();

    let (_, output) = compute_intcode_extended(&mut input.clone(), vec![1]);
    println!("part 1 = {}", output.last().unwrap());

    let (_, output) = compute_intcode_extended(&mut input.clone(), vec![5]);
    println!("part 2 = {}", output[0]);
    Ok(())
}

#[test]
fn test_halt() {
    let (result, _) = compute_intcode_extended(&mut vec![99], vec![1]);
    assert_eq!(result, vec![99]);
}

#[test]
fn test_write_to_memory() {
    let (result, _) = compute_intcode_extended(&mut vec![3, 3, 99, 0], vec![20]);
    assert_eq!(20, result[3]);
}

#[test]
fn test_write_to_output() {
    let (_, output) = compute_intcode_extended(&mut vec![4, 2, 99], vec![1]);
    assert_eq!(99, output[0]);
}

#[test]
fn test_add() {
    let (result, _) = compute_intcode_extended(&mut vec![1, 2, 2, 0, 99], vec![1]);
    assert_eq!(4, result[0]);
}

#[test]
fn test_mul() {
    let (result, _) = compute_intcode_extended(&mut vec![2, 2, 4, 0, 99], vec![1]);
    assert_eq!(396, result[0]);
}

#[test]
fn test_immediate_mode() {
    let (result, _) = compute_intcode_extended(&mut vec![1102, 2, 4, 0, 99], vec![1]);
    assert_eq!(8, result[0]);
}

#[test]
fn test_position_mode_equal() {
    let (_, output) =
        compute_intcode_extended(&mut vec![3, 9, 8, 9, 10, 9, 4, 9, 99, -1, 8], vec![8]);
    assert_eq!(1, output[0]);
}

#[test]
fn test_position_mode_not_equal() {
    let (_, output) =
        compute_intcode_extended(&mut vec![3, 9, 8, 9, 10, 9, 4, 9, 99, -1, 8], vec![10]);
    assert_eq!(0, output[0]);
}

#[test]
fn test_position_mode_less_than() {
    let (_, output) =
        compute_intcode_extended(&mut vec![3, 9, 7, 9, 10, 9, 4, 9, 99, -1, 8], vec![3]);
    assert_eq!(1, output[0]);
}

#[test]
fn test_position_mode_greater_than() {
    let (_, output) =
        compute_intcode_extended(&mut vec![3, 9, 7, 9, 10, 9, 4, 9, 99, -1, 8], vec![10]);
    assert_eq!(0, output[0]);
}

#[test]
fn test_immediate_mode_equal() {
    let (_, output) = compute_intcode_extended(&mut vec![3, 3, 1108, -1, 8, 3, 4, 3, 99], vec![8]);
    assert_eq!(1, output[0]);
}

#[test]
fn test_immediate_mode_not_equal() {
    let (_, output) = compute_intcode_extended(&mut vec![3, 3, 1108, -1, 8, 3, 4, 3, 99], vec![10]);
    assert_eq!(0, output[0]);
}

#[test]
fn test_immediate_mode_less_than() {
    let (_, output) = compute_intcode_extended(&mut vec![3, 3, 1107, -1, 8, 3, 4, 3, 99], vec![3]);
    assert_eq!(1, output[0]);
}

#[test]
fn test_immediate_mode_greater_than() {
    let (_, output) = compute_intcode_extended(&mut vec![3, 3, 1107, -1, 8, 3, 4, 3, 99], vec![10]);
    assert_eq!(0, output[0]);
}

#[test]
fn test_position_jump_zero() {
    let (_, output) = compute_intcode_extended(
        &mut vec![3, 12, 6, 12, 15, 1, 13, 14, 13, 4, 13, 99, -1, 0, 1, 9],
        vec![0],
    );
    assert_eq!(0, output[0]);
}

#[test]
fn test_position_jump_one() {
    let (_, output) = compute_intcode_extended(
        &mut vec![3, 12, 6, 12, 15, 1, 13, 14, 13, 4, 13, 99, -1, 0, 1, 9],
        vec![100],
    );
    assert_eq!(1, output[0]);
}

#[test]
fn test_immediate_jump_zero() {
    let (_, output) = compute_intcode_extended(
        &mut vec![3, 3, 1105, -1, 9, 1101, 0, 0, 12, 4, 12, 99, 1],
        vec![0],
    );
    assert_eq!(0, output[0]);
}

#[test]
fn test_immediate_jump_one() {
    let (_, output) = compute_intcode_extended(
        &mut vec![3, 3, 1105, -1, 9, 1101, 0, 0, 12, 4, 12, 99, 1],
        vec![100],
    );
    assert_eq!(1, output[0]);
}