use aoc2019::*;
use rpds::list;
use rpds::list::List;
use rpds::RedBlackTreeMap;
use std::fmt;
use std::io;
use std::io::prelude::*;
use std::iter;
use std::process;
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
#[structopt(name = "Day 11: Space Police")]
/// Calculates how many blocks a painting robot would paint.
///
/// Takes the program to run on the robot in on stdin.
///
/// See https://adventofcode.com/2019/day/11 for details.
struct Opt {
    /// debug mode prints the size of the painted area on a black background
    #[structopt(short = "d", long = "debug")]
    debug: bool,
}

fn main() {
    let opt = Opt::from_args();
    let stdin = io::stdin();
    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>();

    let finished_robot = exit_on_failed_assertion(
        Robot::new(program, !opt.debug).execute(),
        "Robot encountered an error",
    );
    if opt.debug {
        println!("{}", finished_robot.canvas.size());
    } else {
        println!("{}", finished_robot);
    }
}

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 Robot {
    program: IntcodeProgram,
    position: (i32, i32),
    facing: Direction,
    canvas: RedBlackTreeMap<(i32, i32), bool>,
    background: bool,
}

impl fmt::Display for Robot {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        (self.min_white_y()..=self.max_white_y())
            .map(move |y| {
                (self.min_white_x()..=self.max_white_x())
                    .map(move |x| (x, y))
                    .map(|coord| self.canvas.get(&coord).cloned().unwrap_or(self.background))
                    .map(|c| write!(f, "{}", if c { '#' } else { ' ' }))
                    .collect::<fmt::Result>()
                    .and_then(|_| writeln!(f, ""))
            })
            .collect()
    }
}

#[derive(Debug, Clone)]
enum Direction {
    Up,
    Down,
    Left,
    Right,
}

impl Robot {
    fn new(program: IntcodeProgram, background: bool) -> Robot {
        Robot {
            program: program.run_to_termination_or_input(),
            position: (0, 0),
            facing: Direction::Up,
            canvas: RedBlackTreeMap::new(),
            background,
        }
    }

    fn execute(&self) -> Result<Robot, IntcodeProgramError> {
        iter::successors(Some(self.clone()), |robot| Some(robot.next()))
            .find(|robot| {
                robot.program.error.is_some()
                    || (robot.program.halted && robot.program.output.is_empty())
            })
            .unwrap() // infinite iterator won't terminate unless this is Some
            .as_result()
    }

    fn as_result(&self) -> Result<Robot, IntcodeProgramError> {
        match self.program.error {
            Some(ref error) => Err(error.clone()),
            None => Ok(self.clone()),
        }
    }

    fn next(&self) -> Robot {
        match (
            self.program.output.get(0).map(intcode_to_bool),
            self.program.output.get(1).map(intcode_to_bool),
        ) {
            (Some(paint), Some(rot)) => Robot {
                program: self
                    .program
                    .with_cleared_output()
                    .with_input(list![bool_to_intcode(
                        self.canvas
                            .get(&self.facing.rotate(rot).move_position(self.position))
                            .cloned()
                            .unwrap_or(self.background),
                    )])
                    .run_to_termination_or_input(),
                position: self.facing.rotate(rot).move_position(self.position),
                facing: self.facing.rotate(rot),
                canvas: self.canvas.insert(self.position, paint),
                background: self.background,
            },
            _ => Robot {
                program: self
                    .program
                    .with_input(list![bool_to_intcode(
                        self.canvas
                            .get(&self.position)
                            .cloned()
                            .unwrap_or(self.background),
                    )])
                    .run_to_termination_or_input(),
                ..self.clone()
            },
        }
    }

    fn min_white_x(&self) -> i32 {
        self.white_blocks().map(|(x, _y)| x).min().unwrap_or(0)
    }
    fn min_white_y(&self) -> i32 {
        self.white_blocks().map(|(_x, y)| y).min().unwrap_or(0)
    }
    fn max_white_x(&self) -> i32 {
        self.white_blocks().map(|(x, _y)| x).max().unwrap_or(0)
    }
    fn max_white_y(&self) -> i32 {
        self.white_blocks().map(|(_x, y)| y).max().unwrap_or(0)
    }

    fn white_blocks<'a>(&'a self) -> impl 'a + Iterator<Item = (i32, i32)> {
        self.canvas
            .iter()
            .filter(|(_, val)| **val)
            .map(|(coord, _)| coord)
            .cloned()
    }
}

impl Direction {
    fn rotate(&self, clockwise: bool) -> Direction {
        use Direction::*;

        if clockwise {
            match self {
                Up => Right,
                Right => Down,
                Down => Left,
                Left => Up,
            }
        } else {
            match self {
                Up => Left,
                Left => Down,
                Down => Right,
                Right => Up,
            }
        }
    }

    fn move_position(&self, position: (i32, i32)) -> (i32, i32) {
        use Direction::*;
        match self {
            Up => (position.0, position.1 + 1),
            Down => (position.0, position.1 - 1),
            Left => (position.0 - 1, position.1),
            Right => (position.0 + 1, position.1),
        }
    }
}