use aoc2019::*;
use rpds::list;
use rpds::list::List;
use rpds::vector::Vector;
use rpds::RedBlackTreeMap;
use std::cmp::Ordering;
use std::io;
use std::io::prelude::*;
use std::iter;
use std::process;
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
#[structopt(name = "Day 13: Care Package")]
/// Executes an Intcode game
///
/// The program is read from stdin as a series of comma-separated
/// values. Newlines are ignored.
///
/// See https://adventofcode.com/2019/day/13 for details.
struct Opt {
    #[structopt(short = "q", long = "quarters")]
    quarters: Option<Intcode>,
}

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

    match opt.quarters {
        Some(quarters) => {
            let result = score_on_won_game(program.with_mem_0(quarters));
            println!("{}", result);
        }
        None => {
            let result = exit_on_failed_assertion(program.execute(), "Program errored");
            println!("{}", count_blocks(&result));
        }
    }
}

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(Default, Clone)]
struct Screen {
    screen: RedBlackTreeMap<(Intcode, Intcode), Intcode>,
    previous_ball: (Intcode, Intcode),
    ball: (Intcode, Intcode),
    paddle: (Intcode, Intcode),
    score: Intcode,
}

impl Screen {
    fn render(output: &Vector<Intcode>) -> Screen {
        (0..output.len() / 3)
            .map(|i| i * 3)
            .map(|i| {
                (
                    output[i].clone(),
                    output[i + 1].clone(),
                    output[i + 2].clone(),
                )
            })
            .fold(Screen::default(), |acc, (x, y, tile)| {
                if x == Intcode::from(-1) && y == Intcode::from(0) {
                    Screen { score: tile, ..acc }
                } else if tile == Intcode::from(4) {
                    Screen {
                        ball: (x, y),
                        previous_ball: acc.ball,
                        ..acc
                    }
                } else if tile == Intcode::from(3) {
                    Screen {
                        paddle: (x.clone(), y.clone()),
                        screen: acc.screen.insert((x, y), tile),
                        ..acc
                    }
                } else if tile == Intcode::from(0) {
                    Screen {
                        screen: acc.screen.remove(&(x, y)),
                        ..acc
                    }
                } else {
                    Screen {
                        screen: acc.screen.insert((x, y), tile),
                        ..acc
                    }
                }
            })
    }

    fn paddle_required_direction(&self) -> Intcode {
        match self.paddle.0.cmp(&self.ball.0) {
            Ordering::Less => Intcode::from(1),
            Ordering::Equal => Intcode::from(0),
            Ordering::Greater => Intcode::from(-1),
        }
    }
}

fn count_blocks(output: &Vector<Intcode>) -> usize {
    Screen::render(output)
        .screen
        .values()
        .filter(|val| **val == Intcode::from(2))
        .count()
}

fn score_on_won_game(program: IntcodeProgram) -> Intcode {
    Screen::render(
        &iter::successors(Some(program.run_to_termination_or_input()), |program| {
            Some(next_game_state(program.clone()))
        })
        .take_while(|program| !program.halted)
        .find(|program| count_blocks(&program.output) == 0)
        .unwrap()
        .output,
    )
    .score
}

fn next_game_state(program: IntcodeProgram) -> IntcodeProgram {
    program_with_next_input(program.run_to_termination_or_input())
}

fn program_with_next_input(program: IntcodeProgram) -> IntcodeProgram {
    program.with_input(list![next_input(&program.output)])
}

fn next_input(output: &Vector<Intcode>) -> Intcode {
    Screen::render(output).paddle_required_direction()
}