use aoc2019::*;
use rpds::list;
use rpds::list::List;
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 23: Category Six")]
/// Executes an Intcode program on a network of computers
///
/// See https://adventofcode.com/2019/day/23 for details.
struct Opt {
    #[structopt(short = "n", long = "network-size")]
    network_size: usize,
    #[structopt(long = "nat")]
    send_nat: 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>();

    let network = Network::new(program, opt.network_size);

    if opt.send_nat {
        println!("{}", network.first_repeated_nat_packet().y);
    } else {
        println!("{}", network.first_nat_packet().y);
    }
}

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 Network {
    computers: Vector<IntcodeProgram>,
    nat: Option<Packet>,
    is_potentially_idle: bool,
    is_idle: bool,
}

impl Network {
    fn new(program: IntcodeProgram, network_size: usize) -> Network {
        Network {
            computers: (0..network_size)
                .map(|ip| program.with_input(list![ip.into(), Intcode::from(-1)]))
                .collect(),
            nat: None,
            is_potentially_idle: false,
            is_idle: false,
        }
    }

    fn first_nat_packet(&self) -> Packet {
        iter::successors(Some(self.clone()), |network| Some(network.next()))
            .filter_map(|network| network.nat)
            .next()
            .unwrap()
    }

    fn first_repeated_nat_packet(&self) -> Packet {
        self.sent_nat_packets()
            .zip(self.sent_nat_packets().skip(1))
            // .inspect(|(nat, _p2)| eprintln!("{}", nat.y))
            .find(|(p1, p2)| p1.y == p2.y)
            .unwrap()
            .0
    }

    fn sent_nat_packets<'a>(&'a self) -> impl Iterator<Item = Packet> + 'a {
        iter::successors(Some(self.clone()), |network| {
            Some(network.with_nat_packet_sent().run_to_network_idle())
        })
        .filter_map(|network| network.nat)
    }

    fn run_to_network_idle(&self) -> Network {
        iter::successors(Some(self.clone()), |network| Some(network.next()))
            .find(|network| network.is_idle)
            .unwrap()
    }

    fn with_nat_packet_sent(&self) -> Network {
        if self.nat.is_some() {
            Network {
                computers: self
                    .computers
                    .iter()
                    .enumerate()
                    .map(|(ip, computer)| {
                        computer
                            .with_additional_input(
                                self.nat
                                    .iter()
                                    .filter(|packet| packet.dest == Intcode::from(ip))
                                    .flat_map(|packet| vec![packet.x.clone(), packet.y.clone()])
                                    .collect(),
                            )
                            .run_to_termination_or_input()
                    })
                    .collect(),
                nat: self.nat.clone(),
                is_potentially_idle: false,
                is_idle: false,
            }
        } else {
            self.clone()
        }
    }

    fn next(&self) -> Network {
        Network {
            computers: self
                .computers
                .iter()
                .enumerate()
                .map(|(ip, computer)| {
                    computer
                        .with_cleared_output()
                        .with_additional_input(if self.empty_queues() {
                            list![Intcode::from(-1)]
                        } else {
                            list![]
                        })
                        .with_additional_input(
                            self.pending_packets()
                                .filter(|packet| packet.dest == Intcode::from(ip))
                                .flat_map(|packet| vec![packet.x, packet.y])
                                .collect(),
                        )
                        .run_to_termination_or_input()
                })
                .collect(),
            nat: self
                .pending_packets()
                .filter(|packet| packet.is_nat())
                .map(|packet| packet.with_dest(0.into()))
                .last()
                .or_else(|| self.nat.clone()),
            is_potentially_idle: self.empty_queues(),
            is_idle: self.is_potentially_idle && self.empty_queues(),
        }
    }

    fn pending_packets<'a>(&'a self) -> impl Iterator<Item = Packet> + 'a {
        self.computers.iter().flat_map(|computer| {
            computer
                .output
                .iter()
                .cloned()
                .collect::<Vec<Intcode>>()
                .chunks_exact(3)
                .map(|packet| Packet {
                    dest: packet[0].clone(),
                    x: packet[1].clone(),
                    y: packet[2].clone(),
                })
                .collect::<Vec<Packet>>()
        })
    }

    fn pending_input<'a>(&'a self) -> impl Iterator<Item = Intcode> + 'a {
        self.computers
            .iter()
            .flat_map(|computer| computer.input.iter().cloned().collect::<Vec<Intcode>>())
    }

    fn empty_queues(&self) -> bool {
        self.pending_packets().count() == 0 && self.pending_input().count() == 0
    }
}

#[derive(Debug, Clone)]
struct Packet {
    dest: Intcode,
    x: Intcode,
    y: Intcode,
}

impl Packet {
    fn is_nat(&self) -> bool {
        self.dest == 255.into()
    }

    fn with_dest(&self, dest: Intcode) -> Packet {
        Packet {
            dest,
            ..self.clone()
        }
    }
}