locate_block.rs

use std::{path::PathBuf, fs::File, env, collections::HashMap, rc::Rc, cell::RefCell, fmt::Display};
use itertools::Itertools;

use fastanvil::*;
use rayon::prelude::{ParallelIterator, IntoParallelIterator};

#[derive(Clone, Eq, PartialEq, Debug)]
struct BlockDescriptor {
    name: String,
    x: i64,
    y: i64,
    z: i64,
}

#[derive(Clone, Eq, PartialEq, Debug)]
struct Vein {
    id: usize,
    blocks: Vec<BlockDescriptor>
}

struct Bounds {
    min_x: i64,
    min_y: i64,
    min_z: i64,
    max_x: i64,
    max_y: i64,
    max_z: i64,
}

impl Display for Bounds {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "({}, {}, {}) -> ({}, {}, {})", self.min_x, self.min_y, self.min_z, self.max_x, self.max_y, self.max_z)
    }
}

impl Vein {
    fn bounds(&self) -> Bounds {
        let mut min_x = i64::MAX;
        let mut min_y = i64::MAX;
        let mut min_z = i64::MAX;
        let mut max_x = i64::MIN;
        let mut max_y = i64::MIN;
        let mut max_z = i64::MIN;

        for block in &self.blocks {
            if block.x < min_x {
                min_x = block.x;
            }
            if block.y < min_y {
                min_y = block.y;
            }
            if block.z < min_z {
                min_z = block.z;
            }
            if block.x > max_x {
                max_x = block.x;
            }
            if block.y > max_y {
                max_y = block.y;
            }
            if block.z > max_z {
                max_z = block.z;
            }
        }

        Bounds {
            min_x,
            min_y,
            min_z,
            max_x,
            max_y,
            max_z,
        }
    }
}

impl Ord for Vein {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.blocks.len().cmp(&other.blocks.len())
    }
}

impl PartialOrd for Vein {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        return self.cmp(other).into();
    }
}

impl Ord for BlockDescriptor {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        if self.x < other.x {
            std::cmp::Ordering::Less
        } else if self.x > other.x {
            std::cmp::Ordering::Greater
        } else if self.y < other.y {
            std::cmp::Ordering::Less
        } else if self.y > other.y {
            std::cmp::Ordering::Greater
        } else if self.z < other.z {
            std::cmp::Ordering::Less
        } else if self.z > other.z {
            std::cmp::Ordering::Greater
        } else {
            std::cmp::Ordering::Equal
        }
    }
}

impl PartialOrd for BlockDescriptor {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        return self.cmp(other).into();
    }
}

struct ScanResult {
    regions: usize,
    chunks: usize,
    sections: usize,
    blocks: usize,
    found: Vec<BlockDescriptor>,
}

fn main() {
    let args = env::args().collect::<Vec<String>>();

    let dim_path = PathBuf::from(args[1].clone());
    let search_block = args[2].clone();
    println!("Looking for {}", search_block);

    let loader = RegionFileLoader::new(dim_path.join("region"));
    let regions = loader.list().unwrap();
    let regions_len = regions.len();
    println!("Scanning {} regions", regions_len);

    // Search for the specified block in each region
    let results = regions
        .into_par_iter()
        .map(|(rx, rz)| {
            let mut file = loader.region(rx, rz).unwrap();

            locate_in_region(search_block.clone(), &mut file, rx.0 as i64, rz.0 as i64)
        })
        .collect::<Vec<ScanResult>>();

    // Summarize the results
    let mut summary = ScanResult {
        regions: 0,
        chunks: 0,
        sections: 0,
        blocks: 0,
        found: vec![],
    };

    for result in results {
        summary.regions += result.regions;
        summary.chunks += result.chunks;
        summary.sections += result.sections;
        summary.blocks += result.blocks;

        let mut found = result.found.clone();
        summary.found.append(&mut found);
    }

    println!("Scanned {} regions", summary.regions);
    println!("Scanned {} chunks", summary.chunks);
    println!("Scanned {} sections", summary.sections);
    println!("Scanned {} blocks", summary.blocks);

    println!("Found {} blocks of {}", summary.found.len(), search_block);
    summary.found = summary.found.into_iter().sorted().collect();

    // for desc in summary.found {
    //     println!("  at {}, {}, {}", desc.x, desc.y, desc.z);
    // }

    println!("Looking for veins");
    let veins = find_veins(summary.found);

    for (num, vein) in veins.iter().sorted().enumerate() {
        println!("Vein {} ({} blocks): {}", num, vein.blocks.len(), vein.bounds());
    }
}

fn locate_in_region(name: String, region: &mut Region<File>, rx: i64, rz: i64) -> ScanResult {
    let chunk_coords = (0..32i64).flat_map(|z| (0..32i64).map(move |x| (x, z)));

    let mut result = ScanResult {
        regions: 1,
        chunks: 0,
        sections: 0,
        blocks: 0,
        found: vec![],
    };

    for (cx, cz) in chunk_coords {
        let chunk = region.read_chunk(cx as usize, cz as usize);
        let chunk = match chunk {
            Ok(Some(data)) => Some(data),
            _ => None,
        };
        let chunk = chunk.and_then(|data| {
            match JavaChunk::from_bytes(&data) {
                Ok(JavaChunk::Post18(chunk)) => Some(chunk),
                _ => None,
            }
        });

        let tower = chunk.and_then(|chunk| chunk.sections);
        if let None = tower {
            continue;
        }

        let mut has_non_air = false;

        let tower = tower.unwrap();
        for section in tower.sections() {
            let palette = section.block_states.palette();
            let is_air = palette.len() == 1 && palette.get(0).unwrap().name() == "minecraft:air";
            let iter = section.block_states.try_iter_indices();
            let sy = (section.y() * 16) as i64;
            has_non_air = has_non_air || !is_air;

            if let None = iter {
                continue
            }

            // Skip air sections completely
            if is_air {
                continue
            }

            result.sections += 1;

            for (i, block_index) in iter.unwrap().enumerate() {
                let block = palette.get(block_index).unwrap();

                // Skip air blocks
                if block.name() == "minecraft:air" {
                    continue;
                }

                result.blocks += 1;

                if name != block.name() {
                    continue;
                }

                let x = (i & 0x000F) as i64;
                let y = ((i & 0x0F00) >> 8) as i64;
                let z = ((i & 0x00F0) >> 4) as i64;

                // 512 blocks per region + 16 blocks per chunk + x
                result.found.push(BlockDescriptor {
                    name: name.clone(),
                    x: rx * 512 + cx * 16 + x,
                    y: sy + y,
                    z: rz * 512 + cz * 16 + z,
                });
            }
        }

        if has_non_air {
            result.chunks += 1;
        }
    }

    result
}

type VeinRef = Rc<RefCell<Vein>>;

// TODO: This approach is both flawed (misses some veins) and slow because O(n^2)
// But it's "good enough" for now
fn find_veins(blocks: Vec<BlockDescriptor>) -> Vec<Vein> {
    let mut next_vein_id = 0usize;

    fn vein_new(id: usize) -> (VeinRef, usize) {
        return (
            Rc::new(RefCell::new(Vein { id, blocks: vec![] })),
            id + 1,
        );
    }

    // Create hashmaps for each coordinate
    let mut x_map = HashMap::new();
    let mut y_map = HashMap::new();
    let mut z_map = HashMap::new();

    for block in &blocks {
        x_map.entry(block.x).or_insert(vec![]).push(block.clone());
        y_map.entry(block.y).or_insert(vec![]).push(block.clone());
        z_map.entry(block.z).or_insert(vec![]).push(block.clone());
    }

    // Create hashmaps pointing coordinates to their vein
    let mut xyz_map = HashMap::<(i64, i64, i64), VeinRef>::new();

    // Create a list of veins
    let mut veins = vec![];

    for block in &blocks {
        let vein;

        if xyz_map.contains_key(&(block.x, block.y, block.z)) {
            vein = xyz_map.get(&(block.x, block.y, block.z)).unwrap().clone();
        } else {
            (vein, next_vein_id) = vein_new(next_vein_id);
        };

        for other in &blocks {
            if block.x == other.x && block.y == other.y && block.z == other.z {
                continue;
            }

            // If it's within 1 block of the current block, add it to the vein
            let dx = block.x.abs_diff(other.x);
            let dy = block.y.abs_diff(other.y);
            let dz = block.z.abs_diff(other.z);

            let should_add = match (dx, dy, dz) {
                // Same block
                (0, 0, 0) => false,

                // Connected directly within 1 block
                (1, 0, 0) => true,
                (0, 1, 0) => true,
                (0, 0, 1) => true,

                // Not connected at all
                _ => false,
            };

            if should_add {
                let has_block = vein.borrow().blocks.contains(other);
                if !has_block {
                    vein.borrow_mut().blocks.push(other.clone())
                }
            }
        }

        if vein.borrow().blocks.len() > 0 {
            for block in &vein.borrow().blocks {
                xyz_map.insert((block.x, block.y, block.z), vein.clone());
            }

            if !veins.contains(&vein) {
                veins.push(vein);
            }
        }
    }

    let mut result = vec![];

    for vein in veins {
        result.push(vein.borrow().clone());
    }

    result
}