use nom::{
    branch::alt,
    bytes::complete::tag,
    character::complete::{i32 as nom_i32, line_ending},
    combinator::{map, value},
    multi::separated_list1,
    sequence::tuple,
    IResult,
};
use std::{cmp, fs};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let input = fs::read_to_string("inputs/day_22.txt")?;
    let instructions = parse_instructions(&input).unwrap().1;
    {
        let bounds_50 = Block {
            min_x: -50,
            max_x: 51,
            min_y: -50,
            max_y: 51,
            min_z: -50,
            max_z: 51,
        };
        let mut octtree_50 = OctTree::default();
        for instruction in &instructions {
            octtree_50.set_block(&bounds_50, &instruction.bounds, instruction.new_state);
        }
        dbg!(octtree_50.count_on_blocks(&bounds_50));
    }

    {
        let problem_boundary = Block {
            min_x: instructions
                .iter()
                .map(|i| i.bounds.min_x)
                .min()
                .unwrap_or(0),
            max_x: instructions
                .iter()
                .map(|i| i.bounds.max_x)
                .max()
                .unwrap_or(0),
            min_y: instructions
                .iter()
                .map(|i| i.bounds.min_y)
                .min()
                .unwrap_or(0),
            max_y: instructions
                .iter()
                .map(|i| i.bounds.max_y)
                .max()
                .unwrap_or(0),
            min_z: instructions
                .iter()
                .map(|i| i.bounds.min_z)
                .min()
                .unwrap_or(0),
            max_z: instructions
                .iter()
                .map(|i| i.bounds.max_z)
                .max()
                .unwrap_or(0),
        }
        .expand_to_power_cube();

        // This chunking and adding partial solutions is necessary
        // because I can't fit the whole thing in memory at once :(
        // It runs really slowly.
        let mut count = 0;
        for chunk_index in 0..8 {
            let problem_boundary = problem_boundary.oct_chunk(chunk_index);
            for chunk_index in 0..8 {
                let problem_boundary = problem_boundary.oct_chunk(chunk_index);
                for chunk_index in 0..8 {
                    let problem_boundary = problem_boundary.oct_chunk(chunk_index);
                    let mut octtree = OctTree::default();
                    for instruction in &instructions {
                        octtree.set_block(
                            &problem_boundary,
                            &instruction.bounds,
                            instruction.new_state,
                        );
                    }
                    count += dbg!(octtree.count_on_blocks(&problem_boundary));
                }
            }
        }
        dbg!(count);
    }

    Ok(())
}

#[derive(Default, Clone)]
struct OctTree {
    data: OctTreeData,
}

impl OctTree {
    fn set_block(&mut self, self_bounds: &Block, bounds: &Block, new_val: bool) {
        if bounds.completely_covers(self_bounds) {
            self.data = new_val.into();
        } else if bounds.intersects(self_bounds) {
            match &mut self.data {
                OctTreeData::AllOff => {
                    if new_val {
                        self.split(self_bounds);
                        self.set_block(self_bounds, bounds, new_val);
                    }
                }
                OctTreeData::AllOn => {
                    if !new_val {
                        self.split(self_bounds);
                        self.set_block(self_bounds, bounds, new_val);
                    }
                }
                OctTreeData::BitSet(ref mut bits) => {
                    let min_x = cmp::max(self_bounds.min_x, bounds.min_x);
                    let max_x = cmp::min(self_bounds.max_x, bounds.max_x);
                    let min_x_index = (min_x - self_bounds.min_x) as usize;
                    let max_x_index = min_x_index + (max_x - min_x) as usize;

                    let min_y = cmp::max(self_bounds.min_y, bounds.min_y);
                    let max_y = cmp::min(self_bounds.max_y, bounds.max_y);
                    let min_y_index = (min_y - self_bounds.min_y) as usize;
                    let max_y_index = min_y_index + (max_y - min_y) as usize;

                    let min_z = cmp::max(self_bounds.min_z, bounds.min_z);
                    let max_z = cmp::min(self_bounds.max_z, bounds.max_z);
                    let min_z_index = (min_z - self_bounds.min_z) as usize;
                    let max_z_index = min_z_index + (max_z - min_z) as usize;

                    for z_index in min_z_index..max_z_index {
                        let z_bit_index = z_index << 4;
                        for y_index in min_y_index..max_y_index {
                            let y_bit_index = y_index << 2;
                            for x_index in min_x_index..max_x_index {
                                let x_bit_index = x_index;
                                let bit_mask = 1u64 << (z_bit_index + y_bit_index + x_bit_index);
                                if new_val {
                                    *bits |= bit_mask;
                                } else {
                                    *bits &= !bit_mask;
                                }
                            }
                        }
                    }

                    if *bits == 0 {
                        self.data = OctTreeData::AllOff;
                    } else if *bits == !0u64 {
                        self.data = OctTreeData::AllOn;
                    }
                }
                OctTreeData::Diverse(ref mut subtrees) => {
                    for (sub_index, sub) in subtrees.iter_mut().enumerate() {
                        sub.set_block(&self_bounds.oct_chunk(sub_index as u8), bounds, new_val);
                    }
                    if subtrees
                        .iter()
                        .all(|sub| matches!(sub.data, OctTreeData::AllOn))
                    {
                        self.data = OctTreeData::AllOn;
                    } else if subtrees
                        .iter()
                        .all(|sub| matches!(sub.data, OctTreeData::AllOff))
                    {
                        self.data = OctTreeData::AllOff;
                    }
                }
            };
        }
    }

    fn split(&mut self, self_bounds: &Block) {
        assert!(!matches!(self.data, OctTreeData::Diverse(_)));
        if self_bounds.volume() == 64 {
            let new_bitset = match self.data {
                OctTreeData::AllOn => !0u64,
                OctTreeData::AllOff => 0,
                _ => panic!("weird split"),
            };
            self.data = OctTreeData::BitSet(new_bitset);
        } else {
            let template = OctTree {
                data: self.data.clone(),
            };
            self.data = OctTreeData::Diverse(Box::new([
                template.clone(),
                template.clone(),
                template.clone(),
                template.clone(),
                template.clone(),
                template.clone(),
                template.clone(),
                template.clone(),
            ]));
        }
    }

    fn count_on_blocks(&self, self_bounds: &Block) -> usize {
        match &self.data {
            OctTreeData::AllOff => 0,
            OctTreeData::AllOn => self_bounds.volume(),
            OctTreeData::BitSet(bitset) => bitset.count_ones() as usize,
            OctTreeData::Diverse(subtrees) => subtrees
                .iter()
                .enumerate()
                .map(|(index, sub)| sub.count_on_blocks(&self_bounds.oct_chunk(index as u8)))
                .sum(),
        }
    }
}

#[derive(Clone)]
enum OctTreeData {
    AllOff,
    AllOn,
    BitSet(u64),
    Diverse(Box<[OctTree; 8]>),
}

impl Default for OctTreeData {
    fn default() -> OctTreeData {
        Self::AllOff
    }
}

impl From<bool> for OctTreeData {
    fn from(b: bool) -> Self {
        if b {
            Self::AllOn
        } else {
            Self::AllOff
        }
    }
}

#[derive(Debug)]
struct Instruction {
    new_state: bool,
    bounds: Block,
}

#[derive(Debug, Clone)]
struct Block {
    min_x: i32,
    max_x: i32,
    min_y: i32,
    max_y: i32,
    min_z: i32,
    max_z: i32,
}

impl Block {
    fn volume(&self) -> usize {
        let x = (self.max_x - self.min_x) as usize;
        let y = (self.max_y - self.min_y) as usize;
        let z = (self.max_z - self.min_z) as usize;
        x * y * z
    }

    fn completely_covers(&self, other: &Self) -> bool {
        self.min_x <= other.min_x
            && self.max_x >= other.max_x
            && self.min_y <= other.min_y
            && self.max_y >= other.max_y
            && self.min_z <= other.min_z
            && self.max_z >= other.max_z
    }

    fn intersects(&self, other: &Self) -> bool {
        if self.max_x <= other.min_x
            || self.min_x >= other.max_x
            || self.max_y <= other.min_y
            || self.min_y >= other.max_y
            || self.max_z <= other.min_z
            || self.min_z >= other.max_z
        {
            false
        } else {
            true
        }
    }

    fn oct_chunk(&self, chunk: u8) -> Block {
        let lower_x = (chunk & 1) == 0;
        let lower_y = (chunk & 2) == 0;
        let lower_z = (chunk & 4) == 0;

        let mid_x = (self.min_x + self.max_x) / 2;
        let mid_y = (self.min_y + self.max_y) / 2;
        let mid_z = (self.min_z + self.max_z) / 2;

        Block {
            min_x: if lower_x { self.min_x } else { mid_x },
            max_x: if lower_x { mid_x } else { self.max_x },
            min_y: if lower_y { self.min_y } else { mid_y },
            max_y: if lower_y { mid_y } else { self.max_y },
            min_z: if lower_z { self.min_z } else { mid_z },
            max_z: if lower_z { mid_z } else { self.max_z },
        }
    }

    fn expand_to_power_cube(&self) -> Block {
        let mag_x = self.max_x - self.min_x;
        let mag_y = self.max_y - self.min_y;
        let mag_z = self.max_z - self.min_z;
        let mag_max = cmp::max(mag_x, cmp::max(mag_y, mag_z));
        let first_power_of_2 = (0..)
            .map(|pow| 2_i32.pow(pow))
            .filter(|pow_size| *pow_size >= mag_max)
            .next()
            .unwrap();

        Block {
            min_x: self.min_x,
            max_x: self.max_x + first_power_of_2 - mag_x,
            min_y: self.min_y,
            max_y: self.max_y + first_power_of_2 - mag_y,
            min_z: self.min_z,
            max_z: self.max_z + first_power_of_2 - mag_z,
        }
    }
}

fn parse_instructions(input: &str) -> IResult<&str, Vec<Instruction>> {
    separated_list1(line_ending, parse_instruction)(input)
}

fn parse_instruction(input: &str) -> IResult<&str, Instruction> {
    map(
        tuple((
            alt((value(true, tag("on ")), value(false, tag("off ")))),
            parse_block,
        )),
        |(new_state, bounds)| Instruction { new_state, bounds },
    )(input)
}

fn parse_block(input: &str) -> IResult<&str, Block> {
    map(
        tuple((
            tag("x="),
            nom_i32,
            tag(".."),
            nom_i32,
            tag(",y="),
            nom_i32,
            tag(".."),
            nom_i32,
            tag(",z="),
            nom_i32,
            tag(".."),
            nom_i32,
        )),
        |(
            _,
            min_x,
            _,
            max_x_inclusive,
            _,
            min_y,
            _,
            max_y_inclusive,
            _,
            min_z,
            _,
            max_z_inclusive,
        )| Block {
            min_x,
            max_x: max_x_inclusive + 1,
            min_y,
            max_y: max_y_inclusive + 1,
            min_z,
            max_z: max_z_inclusive + 1,
        },
    )(input)
}