use std::io::{self, BufRead};
use std::env;

// https://stackoverflow.com/q/37888042/remove-single-trailing-newline-from-string-without-cloning
fn trim_newline(s: &mut String) {
    if s.ends_with('\n') {
        s.pop();
        if s.ends_with('\r') {
            s.pop();
        }
    }
}

fn read_input() -> Vec<usize> {
    let stdin = io::stdin();

    // Far from neat, but gets the job done...
    stdin.lock().split(b',').map(|v| {
        let mut s = String::from_utf8(v.unwrap()).unwrap();
        trim_newline(&mut s);
        let u:u32 = s.parse().unwrap();
        u as usize
    }).collect()
}

fn restore_gravity_assist(window:&mut [usize], noun:usize, verb:usize) {
    window[0] = noun;
    window[1] = verb;
}

fn add(ip:usize, memory:&mut Vec<usize>) {
    let result = memory[ip + 3];
    let left = memory[ip + 1];
    let right = memory[ip + 2];
    memory[result] = memory[left] + memory[right];
}

fn mul(ip:usize, memory:&mut Vec<usize>) {
    let result = memory[ip + 3];
    let left = memory[ip + 1];
    let right = memory[ip + 2];
    memory[result] = memory[left] * memory[right];
}

fn intcode(mut memory:Vec<usize>) -> usize {
    let mut ip = 0;
    while ip < memory.len() {
        let op = memory[ip];
        match op {
            1 => { add(ip, &mut memory); ip += 4 },
            2 => { mul(ip, &mut memory); ip += 4 },
            99 => break,
            op => panic!("Invalid operation {}", op),
        }
    }
    memory[0]
}

fn first_puzzle(mut memory:Vec<usize>) -> usize {
    println!("No argument provided, hence attemping to solve first puzzle.");

    restore_gravity_assist(&mut memory[1..3], 12, 2);
    intcode(memory)
}

fn second_puzzle(desired:usize, mut memory:Vec<usize>) -> Option<usize> {
    println!("Desired output {} specified. Assuming second puzzle to search for it.", desired);
    for noun in 0..100 {
        for verb in 0..100 {
            restore_gravity_assist(&mut memory[1..3], noun, verb);
            if intcode(memory.clone()) == desired {
                return Some(100 * noun + verb);
            }
        }
    }

    None
}

fn main() {
    let memory = read_input();

    let needle = env::args().nth(1);
    let solution = match needle {
        None    => Some(first_puzzle(memory)),
        Some(n) => second_puzzle(n.parse().unwrap_or(0), memory),
    };

    match solution {
        Some(s) => println!("Solution: {}.", s),
        _       => println!("No solution found."),
    }
}