use im::vector::Vector;
use std::fmt;
use std::io;
use std::io::prelude::*;
use std::iter;
use std::iter::FromIterator;
use std::iter::IntoIterator;
use std::process;
use structopt::StructOpt;

type Intcode = u32;

#[derive(Debug, StructOpt)]
#[structopt(name = "Day 2: 1202 Program Alarm")]
/// Executes an Intcode program
///
/// The program is read from stdin as a series of comma-separated
/// values. Newlines are ignored. When the program halts, the value at
/// position 0 is returned.
///
/// If an output is provided, all possible inputs are tried to find
/// the input that results in the desired output. In this case, the
/// inputs are returned in the format (noun, verb).
///
///See https://adventofcode.com/2019/day/2 for details.
struct Opt {
    #[structopt(short = "n", long = "noun")]
    noun: Option<Intcode>,
    #[structopt(short = "v", long = "verb")]
    verb: Option<Intcode>,
    #[structopt(short = "o", long = "output")]
    output: Option<Intcode>,
}

fn main() {
    let stdin = io::stdin();
    let opt = Opt::from_args();

    let program: Program = stdin
        .lock()
        .split(b',')
        .map(|x| exit_on_failed_assertion(x, "Error reading input"))
        .map(|x| exit_on_failed_assertion(String::from_utf8(x), "Input was not valid UTF-8"))
        .map(|x| exit_on_failed_assertion(x.trim().parse::<Intcode>(), "Invalid number"))
        .collect();

    match (opt.noun, opt.verb, opt.output) {
        (Some(noun), Some(verb), _) => {
            let result = exit_on_failed_assertion(
                program.set_input(noun, verb).execute(),
                "Program errored",
            );
            println!("{}", result);
        }
        (_, _, Some(output)) => {
            let (noun, verb) =
                exit_on_failed_assertion(program.find_input(output), "Program errored");
            println!("({}, {})", noun, verb);
        }
        (None, None, None) => {
            let result = exit_on_failed_assertion(program.execute(), "Program errored");
            println!("{}", result);
        }
        _ => {
            eprintln!("Either a noun and verb or an expected output must be provided");
            process::exit(1);
        }
    }
}

fn exit_on_failed_assertion<A, E: std::error::Error>(data: Result<A, E>, message: &str) -> A {
    match data {
        Ok(data) => data,
        Err(e) => {
            eprintln!("{}: {}", message, e);
            process::exit(1);
        }
    }
}

#[derive(Debug, Clone)]
struct Program {
    program_counter: usize,
    error: bool,
    halted: bool,
    codes: Vector<Intcode>,
}

impl FromIterator<Intcode> for Program {
    fn from_iter<I: IntoIterator<Item = Intcode>>(iter: I) -> Self {
        Program {
            program_counter: 0,
            error: false,
            halted: false,
            codes: iter.into_iter().collect(),
        }
    }
}

impl Program {
    fn set_input(&self, noun: Intcode, verb: Intcode) -> Program {
        Program {
            codes: self.codes.update(1, noun).update(2, verb),
            ..self.clone()
        }
    }

    fn find_input(&self, output: Intcode) -> Result<(Intcode, Intcode), ProgramError> {
        (0..99)
            .flat_map(|noun| (0..99).map(move |verb| (noun, verb)))
            .map(|(noun, verb)| (noun, verb, self.set_input(noun, verb).execute()))
            .find(|(_noun, _verb, out)| *out == Ok(output))
            .map(|(noun, verb, _out)| Ok((noun, verb)))
            .unwrap_or(Err(ProgramError))
    }

    fn execute(&self) -> Result<Intcode, ProgramError> {
        self.run_to_termination().into_result()
    }

    fn into_result(&self) -> Result<Intcode, ProgramError> {
        if self.error {
            Err(ProgramError)
        } else {
            Ok(self.codes.head().unwrap().clone())
        }
    }

    fn run_to_termination(&self) -> Program {
        iter::successors(Some(self.clone()), |p| Some(Program::next(&p)))
            .find(|p| p.halted)
            .unwrap() // successors doesn't terminate, so this will never be none.
    }

    fn next(&self) -> Program {
        //eprintln!("{:?}", self);
        match self.codes.get(self.program_counter) {
            Some(1) => self.add(),
            Some(2) => self.multiply(),
            Some(99) => self.halt(),
            Some(_) => self.error(),
            None => self.error(),
        }
    }

    fn add(&self) -> Program {
        match (self.get_deref(1), self.get_deref(2), self.get(3)) {
            (Some(in1), Some(in2), Some(out)) => Program {
                program_counter: self.program_counter + 4,
                codes: self.codes.update(out as usize, in1 + in2),
                ..self.clone()
            },
            _ => self.error(),
        }
    }

    fn multiply(&self) -> Program {
        match (self.get_deref(1), self.get_deref(2), self.get(3)) {
            (Some(in1), Some(in2), Some(out)) => Program {
                program_counter: self.program_counter + 4,
                codes: self.codes.update(out as usize, in1 * in2),
                ..self.clone()
            },
            _ => self.error(),
        }
    }

    fn halt(&self) -> Program {
        Program {
            halted: true,
            ..self.clone()
        }
    }

    fn error(&self) -> Program {
        Program {
            halted: true,
            error: true,
            ..self.clone()
        }
    }

    fn get(&self, pointer_offset: usize) -> Option<Intcode> {
        self.codes
            .get(self.program_counter + pointer_offset)
            .cloned()
    }

    fn get_deref(&self, pointer_offset: usize) -> Option<Intcode> {
        self.get(pointer_offset)
            .and_then(|r| self.codes.get(r as usize))
            .cloned()
    }
}

#[derive(Debug, PartialEq)]
struct ProgramError;

impl fmt::Display for ProgramError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Unknown error")
    }
}
impl std::error::Error for ProgramError {}