use aoc2019::*;
use rpds::list;
use rpds::list::List;
use rpds::rbt_set;
use rpds::set::red_black_tree_set::RedBlackTreeSet;
use rpds::vector;
use rpds::vector::Vector;
use std::io;
use std::io::prelude::*;
use std::iter;
use std::process;
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
#[structopt(name = "Day 15: Oxygen System")]
/// Executes an Intcode robot that's searching a map. Prints the
/// time taken for oxygen to propagate to the whole area.
///
/// See https://adventofcode.com/2019/day/15 for details.
struct Opt {
    /// Run in 'find' mode, find the oxygen tank but don't see how long it takes the oxygen to propagate
    #[structopt(short = "f")]
    find: bool,
}

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

    let program: IntcodeProgram = 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::<IntcodeProgram>();

    if opt.find {
        println!("{}", shortest_distance(program));
    } else {
        println!("{}", max_depth_from_oxygen(program));
    }
}

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);
        }
    }
}

fn shortest_distance(program: IntcodeProgram) -> usize {
    iter::successors(
        Some((
            rbt_set![(0, 0)],
            vector![((0, 0), program.run_to_termination_or_input())],
        )),
        |(visited, programs)| {
            Some(next_depth_states(
                visited.clone(),
                next_program_states(programs, visited),
            ))
        },
    )
    .enumerate()
    .find(|(_i, (_visited, programs))| {
        programs
            .iter()
            .any(|(_location, program)| program.output == vector![2.into()])
    })
    .unwrap()
    .0
}

fn max_depth_from_oxygen(program: IntcodeProgram) -> usize {
    max_depth(
        iter::successors(
            Some((
                rbt_set![(0, 0)],
                vector![((0, 0), program.run_to_termination_or_input())],
            )),
            |(visited, programs)| {
                Some(next_depth_states(
                    visited.clone(),
                    next_program_states(programs, visited),
                ))
            },
        )
        .find(|(_visited, programs)| {
            programs
                .iter()
                .any(|(_location, program)| program.output == vector![2.into()])
        })
        .unwrap()
        .1
        .iter()
        .find(|(_location, program)| program.output == vector![2.into()])
        .cloned()
        .unwrap()
        .1,
    )
}

fn max_depth(program: IntcodeProgram) -> usize {
    iter::successors(
        Some((
            rbt_set![(0, 0)],
            vector![((0, 0), program.run_to_termination_or_input())],
        )),
        |(visited, programs)| {
            Some(next_depth_states(
                visited.clone(),
                next_program_states(programs, visited),
            ))
        },
    )
    .take_while(|(_visited, programs)| programs.len() > 0)
    .count()
        - 1
}

fn inputs() -> [((i32, i32), Intcode); 4] {
    [
        ((0, 1), Intcode::from(1)),
        ((0, -1), Intcode::from(2)),
        ((1, 0), Intcode::from(3)),
        ((-1, 0), Intcode::from(4)),
    ]
}

fn next_program_states(
    programs: &Vector<((i32, i32), IntcodeProgram)>,
    visited: &RedBlackTreeSet<(i32, i32)>,
) -> Vector<((i32, i32), IntcodeProgram)> {
    let inputs = inputs();
    programs
        .iter()
        .flat_map(|(location, program)| {
            inputs.iter().map(move |(vec, input)| {
                (
                    (location.0 + vec.0, location.1 + vec.1),
                    program
                        .with_cleared_output()
                        .with_input(list![input.clone()])
                        .run_to_termination_or_input(),
                )
            })
        })
        .filter(|(location, program)| {
            !visited.contains(location) && program.output != vector![0.into()]
        })
        .collect()
}

fn next_depth_states(
    previous_visited: RedBlackTreeSet<(i32, i32)>,
    programs: Vector<((i32, i32), IntcodeProgram)>,
) -> (
    RedBlackTreeSet<(i32, i32)>,
    Vector<((i32, i32), IntcodeProgram)>,
) {
    (
        programs
            .iter()
            .fold(previous_visited, |acc, (next, _)| acc.insert(*next)),
        dedup_locations(programs),
    )
}

fn dedup_locations(
    programs: Vector<((i32, i32), IntcodeProgram)>,
) -> Vector<((i32, i32), IntcodeProgram)> {
    programs.iter().fold(Vector::new(), |acc, next| {
        if acc.iter().any(|(loc, _)| *loc == next.0) {
            acc
        } else {
            acc.push_back(next.clone())
        }
    })
}