emod-cli/src/debug/nbt.rs

use std::{
    collections::HashSet,
    time::{SystemTime, UNIX_EPOCH},
};

use rand::RngCore;

use crate::error::{CliError, Result};

use super::config::DebugConfig;

#[path = "level_template.rs"]
mod level_template;

const TAG_END: u8 = 0;
const TAG_BYTE: u8 = 1;
const TAG_SHORT: u8 = 2;
const TAG_INT: u8 = 3;
const TAG_LONG: u8 = 4;
const TAG_FLOAT: u8 = 5;
const TAG_DOUBLE: u8 = 6;
const TAG_BYTE_ARRAY: u8 = 7;
const TAG_STRING: u8 = 8;
const TAG_LIST: u8 = 9;
const TAG_COMPOUND: u8 = 10;
const TAG_INT_ARRAY: u8 = 11;
const TAG_LONG_ARRAY: u8 = 12;

#[derive(Clone, Debug, PartialEq)]
pub enum Tag {
    Byte(i8),
    Short(i16),
    Int(i32),
    Long(i64),
    Float(f32),
    Double(f64),
    ByteArray(Vec<i8>),
    String(String),
    List { element_type: u8, items: Vec<Tag> },
    Compound(Vec<(String, Tag)>),
    IntArray(Vec<i32>),
    LongArray(Vec<i64>),
}

#[derive(Clone, Debug, PartialEq)]
pub struct LevelDat {
    version: u32,
    root_name: String,
    root: Tag,
}

impl LevelDat {
    pub fn parse(bytes: &[u8]) -> Result<Self> {
        if bytes.len() < 8 {
            return Err(CliError::InvalidData(
                "level.dat is shorter than Bedrock header".to_string(),
            ));
        }
        let version = u32::from_le_bytes(bytes[0..4].try_into().unwrap());
        let declared_len = u32::from_le_bytes(bytes[4..8].try_into().unwrap()) as usize;
        let payload = &bytes[8..];
        if declared_len != payload.len() {
            return Err(CliError::InvalidData(format!(
                "level.dat payload length mismatch: header={declared_len}, actual={}",
                payload.len()
            )));
        }
        let mut reader = Reader {
            bytes: payload,
            pos: 0,
        };
        let tag_type = reader.u8()?;
        if tag_type != TAG_COMPOUND {
            return Err(CliError::InvalidData(
                "level.dat root tag is not a compound".to_string(),
            ));
        }
        let root_name = reader.string()?;
        let root = reader.tag_payload(TAG_COMPOUND)?;
        if reader.pos != payload.len() {
            return Err(CliError::InvalidData(
                "trailing bytes after level.dat root compound".to_string(),
            ));
        }
        Ok(Self {
            version,
            root_name,
            root,
        })
    }

    pub fn to_bytes(&self) -> Result<Vec<u8>> {
        let mut payload = Vec::new();
        payload.push(TAG_COMPOUND);
        write_string(&mut payload, &self.root_name)?;
        write_tag_payload(&mut payload, &self.root)?;

        let mut bytes = Vec::with_capacity(8 + payload.len());
        bytes.extend_from_slice(&self.version.to_le_bytes());
        bytes.extend_from_slice(&(payload.len() as u32).to_le_bytes());
        bytes.extend_from_slice(&payload);
        Ok(bytes)
    }

    pub fn set_root(&mut self, key: &str, value: Tag) -> Result<()> {
        let Tag::Compound(items) = &mut self.root else {
            return Err(CliError::InvalidData(
                "level.dat root is not compound".to_string(),
            ));
        };
        set_compound_item(items, key, value);
        Ok(())
    }

    #[cfg(test)]
    pub fn root_value(&self, key: &str) -> Option<&Tag> {
        let Tag::Compound(items) = &self.root else {
            return None;
        };
        items
            .iter()
            .find(|(name, _)| name == key)
            .map(|(_, tag)| tag)
    }

    pub fn root_compound_mut(&mut self, key: &str) -> Result<&mut Vec<(String, Tag)>> {
        let Tag::Compound(items) = &mut self.root else {
            return Err(CliError::InvalidData(
                "level.dat root is not compound".to_string(),
            ));
        };
        if let Some(index) = items.iter().position(|(name, _)| name == key) {
            if !matches!(items[index].1, Tag::Compound(_)) {
                items[index].1 = Tag::Compound(Vec::new());
            }
            let Tag::Compound(child) = &mut items[index].1 else {
                unreachable!()
            };
            return Ok(child);
        }
        items.push((key.to_string(), Tag::Compound(Vec::new())));
        let Tag::Compound(child) = &mut items.last_mut().unwrap().1 else {
            unreachable!()
        };
        Ok(child)
    }
}

pub fn create_level_dat(config: &DebugConfig) -> Result<Vec<u8>> {
    let mut level = LevelDat::parse(level_template::LEVEL_DAT_TEMPLATE)?;
    apply_config(&mut level, config, true)?;
    level.to_bytes()
}

pub fn update_level_dat_world_data(
    bytes: &[u8],
    config: &DebugConfig,
    init: bool,
) -> Result<Vec<u8>> {
    let mut level = LevelDat::parse(bytes)?;
    apply_config(&mut level, config, init)?;
    level.to_bytes()
}

pub fn update_level_dat_last_played(bytes: &[u8]) -> Result<Vec<u8>> {
    let mut level = LevelDat::parse(bytes)?;
    level.set_root("LastPlayed", Tag::Long(now_seconds()))?;
    level.to_bytes()
}

fn apply_config(level: &mut LevelDat, config: &DebugConfig, init: bool) -> Result<()> {
    level.set_root("LastPlayed", Tag::Long(now_seconds()))?;
    level.set_root("LevelName", Tag::String(config.world_name.clone()))?;
    if init && config.world_type != 2 {
        level.set_root(
            "RandomSeed",
            Tag::Long(config.world_seed.unwrap_or_else(random_seed)),
        )?;
    }
    level.set_root("GameType", Tag::Int(config.game_mode))?;
    if init {
        level.set_root("Generator", Tag::Int(config.world_type))?;
    }
    level.set_root("keepInventory", Tag::Byte(config.keep_inventory as i8))?;
    level.set_root("cheatsEnabled", Tag::Byte(config.enable_cheats as i8))?;
    level.set_root("doweathercycle", Tag::Byte(config.do_weather_cycle as i8))?;
    level.set_root("dodaylightcycle", Tag::Byte(config.do_daylight_cycle as i8))?;

    if let Some(experiments) = &config.experiment_options {
        let child = level.root_compound_mut("experiments")?;
        set_compound_item(
            child,
            "data_driven_biomes",
            Tag::Byte(experiments.data_driven_biomes as i8),
        );
        set_compound_item(
            child,
            "data_driven_items",
            Tag::Byte(experiments.data_driven_items as i8),
        );
        set_compound_item(
            child,
            "experimental_molang_features",
            Tag::Byte(experiments.experimental_molang_features as i8),
        );
    }

    Ok(())
}

fn now_seconds() -> i64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_secs() as i64)
        .unwrap_or(0)
}

fn random_seed() -> i64 {
    rand::thread_rng().next_u64() as i64
}

fn set_compound_item(items: &mut Vec<(String, Tag)>, key: &str, value: Tag) {
    if let Some((_, slot)) = items.iter_mut().find(|(name, _)| name == key) {
        *slot = value;
    } else {
        items.push((key.to_string(), value));
    }
}

struct Reader<'a> {
    bytes: &'a [u8],
    pos: usize,
}

impl Reader<'_> {
    fn take(&mut self, len: usize) -> Result<&[u8]> {
        let end = self
            .pos
            .checked_add(len)
            .ok_or_else(|| CliError::InvalidData("NBT offset overflow".to_string()))?;
        if end > self.bytes.len() {
            return Err(CliError::InvalidData(
                "unexpected end of NBT data".to_string(),
            ));
        }
        let out = &self.bytes[self.pos..end];
        self.pos = end;
        Ok(out)
    }

    fn u8(&mut self) -> Result<u8> {
        Ok(self.take(1)?[0])
    }
    fn i8(&mut self) -> Result<i8> {
        Ok(self.u8()? as i8)
    }
    fn i16(&mut self) -> Result<i16> {
        Ok(i16::from_le_bytes(self.take(2)?.try_into().unwrap()))
    }
    fn i32(&mut self) -> Result<i32> {
        Ok(i32::from_le_bytes(self.take(4)?.try_into().unwrap()))
    }
    fn i64(&mut self) -> Result<i64> {
        Ok(i64::from_le_bytes(self.take(8)?.try_into().unwrap()))
    }
    fn f32(&mut self) -> Result<f32> {
        Ok(f32::from_le_bytes(self.take(4)?.try_into().unwrap()))
    }
    fn f64(&mut self) -> Result<f64> {
        Ok(f64::from_le_bytes(self.take(8)?.try_into().unwrap()))
    }

    fn string(&mut self) -> Result<String> {
        let len = u16::from_le_bytes(self.take(2)?.try_into().unwrap()) as usize;
        let bytes = self.take(len)?;
        String::from_utf8(bytes.to_vec())
            .map_err(|err| CliError::InvalidData(format!("NBT string is not UTF-8: {err}")))
    }

    fn tag_payload(&mut self, tag_type: u8) -> Result<Tag> {
        match tag_type {
            TAG_BYTE => Ok(Tag::Byte(self.i8()?)),
            TAG_SHORT => Ok(Tag::Short(self.i16()?)),
            TAG_INT => Ok(Tag::Int(self.i32()?)),
            TAG_LONG => Ok(Tag::Long(self.i64()?)),
            TAG_FLOAT => Ok(Tag::Float(self.f32()?)),
            TAG_DOUBLE => Ok(Tag::Double(self.f64()?)),
            TAG_BYTE_ARRAY => {
                let len = checked_len(self.i32()?)?;
                let mut values = Vec::with_capacity(len);
                for _ in 0..len {
                    values.push(self.i8()?);
                }
                Ok(Tag::ByteArray(values))
            }
            TAG_STRING => Ok(Tag::String(self.string()?)),
            TAG_LIST => {
                let element_type = self.u8()?;
                let len = checked_len(self.i32()?)?;
                let mut items = Vec::with_capacity(len);
                for _ in 0..len {
                    items.push(self.tag_payload(element_type)?);
                }
                Ok(Tag::List {
                    element_type,
                    items,
                })
            }
            TAG_COMPOUND => {
                let mut items = Vec::new();
                loop {
                    let child_type = self.u8()?;
                    if child_type == TAG_END {
                        break;
                    }
                    let name = self.string()?;
                    let value = self.tag_payload(child_type)?;
                    items.push((name, value));
                }
                Ok(Tag::Compound(items))
            }
            TAG_INT_ARRAY => {
                let len = checked_len(self.i32()?)?;
                let mut values = Vec::with_capacity(len);
                for _ in 0..len {
                    values.push(self.i32()?);
                }
                Ok(Tag::IntArray(values))
            }
            TAG_LONG_ARRAY => {
                let len = checked_len(self.i32()?)?;
                let mut values = Vec::with_capacity(len);
                for _ in 0..len {
                    values.push(self.i64()?);
                }
                Ok(Tag::LongArray(values))
            }
            _ => Err(CliError::InvalidData(format!(
                "unsupported NBT tag type {tag_type}"
            ))),
        }
    }
}

fn checked_len(value: i32) -> Result<usize> {
    usize::try_from(value)
        .map_err(|_| CliError::InvalidData(format!("negative NBT collection length {value}")))
}

fn write_named_tag(out: &mut Vec<u8>, name: &str, tag: &Tag) -> Result<()> {
    out.push(tag_id(tag));
    write_string(out, name)?;
    write_tag_payload(out, tag)
}

fn write_string(out: &mut Vec<u8>, value: &str) -> Result<()> {
    let len = u16::try_from(value.len())
        .map_err(|_| CliError::InvalidData("NBT string is too long".to_string()))?;
    out.extend_from_slice(&len.to_le_bytes());
    out.extend_from_slice(value.as_bytes());
    Ok(())
}

fn write_tag_payload(out: &mut Vec<u8>, tag: &Tag) -> Result<()> {
    match tag {
        Tag::Byte(value) => out.push(*value as u8),
        Tag::Short(value) => out.extend_from_slice(&value.to_le_bytes()),
        Tag::Int(value) => out.extend_from_slice(&value.to_le_bytes()),
        Tag::Long(value) => out.extend_from_slice(&value.to_le_bytes()),
        Tag::Float(value) => out.extend_from_slice(&value.to_le_bytes()),
        Tag::Double(value) => out.extend_from_slice(&value.to_le_bytes()),
        Tag::ByteArray(values) => {
            write_len(out, values.len())?;
            out.extend(values.iter().map(|value| *value as u8));
        }
        Tag::String(value) => write_string(out, value)?,
        Tag::List {
            element_type,
            items,
        } => {
            out.push(*element_type);
            write_len(out, items.len())?;
            for item in items {
                if tag_id(item) != *element_type {
                    return Err(CliError::InvalidData(
                        "NBT list contains mixed element types".to_string(),
                    ));
                }
                write_tag_payload(out, item)?;
            }
        }
        Tag::Compound(items) => {
            let mut names = HashSet::with_capacity(items.len());
            for (name, value) in items {
                if !names.insert(name) {
                    return Err(CliError::InvalidData(format!(
                        "duplicate NBT compound key {name}"
                    )));
                }
                write_named_tag(out, name, value)?;
            }
            out.push(TAG_END);
        }
        Tag::IntArray(values) => {
            write_len(out, values.len())?;
            for value in values {
                out.extend_from_slice(&value.to_le_bytes());
            }
        }
        Tag::LongArray(values) => {
            write_len(out, values.len())?;
            for value in values {
                out.extend_from_slice(&value.to_le_bytes());
            }
        }
    }
    Ok(())
}

fn write_len(out: &mut Vec<u8>, len: usize) -> Result<()> {
    let len = i32::try_from(len)
        .map_err(|_| CliError::InvalidData("NBT collection is too long".to_string()))?;
    out.extend_from_slice(&len.to_le_bytes());
    Ok(())
}

fn tag_id(tag: &Tag) -> u8 {
    match tag {
        Tag::Byte(_) => TAG_BYTE,
        Tag::Short(_) => TAG_SHORT,
        Tag::Int(_) => TAG_INT,
        Tag::Long(_) => TAG_LONG,
        Tag::Float(_) => TAG_FLOAT,
        Tag::Double(_) => TAG_DOUBLE,
        Tag::ByteArray(_) => TAG_BYTE_ARRAY,
        Tag::String(_) => TAG_STRING,
        Tag::List { .. } => TAG_LIST,
        Tag::Compound(_) => TAG_COMPOUND,
        Tag::IntArray(_) => TAG_INT_ARRAY,
        Tag::LongArray(_) => TAG_LONG_ARRAY,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::debug::config::{DebugConfig, NeteaseConfig};
    use serde_json::{Map, Value};

    fn config() -> DebugConfig {
        DebugConfig {
            included_mod_dirs: Vec::new(),
            world_seed: Some(123),
            reset_world: false,
            world_name: "UnitWorld".to_string(),
            world_folder_name: "UnitWorld".to_string(),
            auto_join_game: true,
            include_debug_mod: true,
            auto_hot_reload_mods: true,
            world_type: 1,
            game_mode: 1,
            enable_cheats: true,
            keep_inventory: true,
            do_weather_cycle: true,
            do_daylight_cycle: true,
            game_executable_path: String::new(),
            debug_options: Value::Object(Map::new()),
            netease_config: NeteaseConfig::default(),
            user_name: "developer".to_string(),
            skin_info: None,
            experiment_options: None,
        }
    }

    #[test]
    fn template_round_trips() {
        let level = LevelDat::parse(level_template::LEVEL_DAT_TEMPLATE).unwrap();
        let bytes = level.to_bytes().unwrap();
        let reparsed = LevelDat::parse(&bytes).unwrap();
        assert_eq!(level, reparsed);
    }

    #[test]
    fn updates_required_world_fields() {
        let bytes = create_level_dat(&config()).unwrap();
        let level = LevelDat::parse(&bytes).unwrap();
        assert_eq!(
            level.root_value("LevelName"),
            Some(&Tag::String("UnitWorld".to_string()))
        );
        assert_eq!(level.root_value("GameType"), Some(&Tag::Int(1)));
        assert_eq!(level.root_value("RandomSeed"), Some(&Tag::Long(123)));
        assert!(matches!(level.root_value("LastPlayed"), Some(Tag::Long(value)) if *value > 0));
    }
}