nwn-lib-rs 0.4.0

//! Talktable file format (.tlk)
//!
//! Example
//! ```rust
//! #![feature(generic_const_exprs)]
//! #![allow(incomplete_features)]
//! use nwn_lib_rs::tlk::{Tlk, Resolver, Gender};
//!
//! // Load a TLK file
//! let tlk_data = std::fs::read("unittest/dialog.tlk").unwrap();
//! let (_, tlk) = Tlk::from_bytes(&tlk_data, false).unwrap();
//!
//! // Resolve a strref
//! let text = tlk.get_str(4).unwrap();
//!
//! // Get complete entry information (for sounds)
//! let entry = tlk.get_entry(4).unwrap();
//!
//! // Setup a resolver that aggregates multiple TLK files
//! let main_tlk = tlk;
//! let (_, user_tlk) = Tlk::from_bytes(
//!         &std::fs::read("unittest/user.tlk").unwrap(),
//!         false
//!     ).unwrap();
//! let resolver = Resolver::new(&main_tlk, None, Some(&user_tlk));
//!
//! // Resolve strrefs
//! let text = resolver.get_str(1, Gender::Male).unwrap();
//! let text = resolver.get_str(16_777_217, Gender::Male).unwrap();
//!
//! ```
//!

#![allow(dead_code)]

use crate::parsing::*;
use nom::multi::count;
use nom::number::streaming::{le_f32, le_u32};
use serde::{Deserialize, Serialize};

pub use crate::parsing::FixedSizeString;

pub const TLK_STRREF_USER_OFFSET: u32 = 0x0100_0000;

#[derive(Debug, Deserialize, Clone)]
#[serde(untagged)]
enum TlkString {
    #[serde(skip)]
    Ptr {
        offset: u32,
        size: u32,
    },
    Owned(String),
}

/// TLK entry data
///
/// Note: when the TLK is parsed using Tlk::from_bytes, the Entry's text is stored as an offset in `Tlk::strings_data: Vec<u8>`
#[derive(Debug, Default, Deserialize, Clone)]
pub struct Entry {
    /// Sound file name, if any
    pub sound_resref: Option<FixedSizeString<16>>,
    /// Unused
    pub _volume_variance: u32,
    /// Unused
    pub _pitch_variance: u32,
    /// Sound length, if any
    pub sound_length: Option<f32>,
    text: Option<TlkString>,
}
impl Entry {
    fn from_bytes(input: &[u8]) -> NWNParseResult<Self> {
        let (input, flags) = le_u32(input)?;
        let (input, sound_resref) = FixedSizeString::<16>::from_bytes(input)?;
        let (input, _volume_variance) = le_u32(input)?;
        let (input, _pitch_variance) = le_u32(input)?;
        let (input, offset_to_string) = le_u32(input)?;
        let (input, string_size) = le_u32(input)?;
        let (input, sound_length) = le_f32(input)?;
        Ok((
            input,
            Self {
                sound_resref: if (flags & 2) > 0 {
                    Some(sound_resref)
                } else {
                    None
                },
                _volume_variance,
                _pitch_variance,
                text: if (flags & 1) > 0 {
                    Some(TlkString::Ptr {
                        offset: offset_to_string,
                        size: string_size,
                    })
                } else {
                    None
                },
                sound_length: if (flags & 4) > 0 {
                    Some(sound_length)
                } else {
                    None
                },
            },
        ))
    }

    /// Returns true if the entry does not contain any text, sound or sound length
    pub fn is_empty(&self) -> bool {
        self.sound_resref.is_none() && self.text.is_none() && self.sound_length.is_none()
    }

    fn get_flags(&self) -> u32 {
        self.text.is_some() as u32
            + self.sound_resref.is_some() as u32 * 2
            + self.sound_length.is_some() as u32 * 4
    }

    /// Returns a TLK entry text. May fail if the provided entry comes from another TLK
    ///
    /// # Panics
    /// if the entry references data that is out of `strings_data` bounds or is invalid utf8 data
    pub fn get_str<'a>(&'a self, strings_data: &'a [u8]) -> Option<&'a str> {
        if let Some(text) = &self.text {
            match text {
                TlkString::Owned(s) => Some(s),
                TlkString::Ptr { offset, size } => {
                    let start = *offset as usize;
                    let end = start + *size as usize;

                    Some(
                        std::str::from_utf8(
                            strings_data
                                .get(start..end)
                                .expect("entry references out of bounds data"),
                        )
                        .expect("entry references invalid utf8 data"),
                    )
                }
            }
        } else {
            None
        }
    }

    /// Returns a TLK entry text length in bytes
    pub fn get_str_len(&self) -> usize {
        if let Some(text) = &self.text {
            match text {
                TlkString::Owned(s) => s.len(),
                TlkString::Ptr { offset: _, size } => *size as usize,
            }
        } else {
            0
        }
    }

    /// Remove any reference to external data
    //
    /// # Panics
    /// if the entry references data that is out of `strings_data` bounds or is invalid utf8 data
    pub fn into_owned(mut self, strings_data: &[u8]) -> Self {
        if let Some(TlkString::Ptr { offset: _, size: _ }) = &self.text {
            self.text = Some(TlkString::Owned(
                self.get_str(strings_data)
                    .expect("should not fail")
                    .to_string(),
            ));
        }
        self
    }

    /// Returns a copy of this entry, that does not contain any reference to external data
    //
    /// # Panics
    /// if the entry references data that is out of `strings_data` bounds or is invalid utf8 data
    pub fn to_owned(&self, strings_data: &[u8]) -> Self {
        self.clone().into_owned(strings_data)
    }

    /// Formats a single TLK entry to text (used by `Tlk::to_string_pretty`)
    ///
    /// * `index`: Entry strref
    /// * `index_padding`: Character width of the left side containing the strref
    /// * `user_indices`: true to display `index` as a user strref
    /// * `strings_data`: string data vector for non-owned TLK entries
    ///
    /// # Panics
    /// if the entry references data that is out of `strings_data` bounds or is invalid utf8 data
    pub fn to_string_pretty(
        &self,
        index: u32,
        index_padding: usize,
        user_indices: bool,
        strings_data: &[u8],
    ) -> String {
        let mut ret = String::new();

        let index = if user_indices {
            index + TLK_STRREF_USER_OFFSET
        } else {
            index
        };

        ret += &format!(
            "{0:1$}> sound_resref={2:?} _volume_variance={3} _pitch_variance={4} \
             sound_length={5:?}",
            index,
            index_padding,
            self.sound_resref,
            self._volume_variance,
            self._pitch_variance,
            self.sound_length,
        );
        if let Some(text) = self.get_str(strings_data) {
            ret += "\n";
            ret += &text
                .lines()
                .map(|line| {
                    format!(
                        "{}|{}{}",
                        " ".repeat(index_padding),
                        if line.is_empty() { "" } else { " " },
                        line
                    )
                })
                .collect::<Vec<_>>()
                .join("\n");
        }
        ret
    }
}
/// Structure for serializing TLK entries that contains TlkString::Ptr
struct EntrySerde<'a> {
    entry: &'a Entry,
    strings_data: &'a [u8],
    /// The strref used to reference this entry
    strref: u32,
}

impl<'a> Serialize for EntrySerde<'a> {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        use serde::ser::SerializeStruct;
        let mut ser_struct = serializer.serialize_struct("Entry", 6)?;
        ser_struct.serialize_field("__hint_strref", &self.strref)?;
        ser_struct.serialize_field("sound_resref", &self.entry.sound_resref)?;
        ser_struct.serialize_field("_volume_variance", &self.entry._volume_variance)?;
        ser_struct.serialize_field("_pitch_variance", &self.entry._pitch_variance)?;
        ser_struct.serialize_field("sound_length", &self.entry.sound_length)?;
        ser_struct.serialize_field("text", &self.entry.get_str(self.strings_data))?;
        ser_struct.end()
    }
}

fn default_false() -> bool {
    false
}

/// Talktable file format (.tlk)
///
/// TLK entries can be stored at two locations:
/// - In each [`Entry`], as a String.
/// - Or in [`strings_data`], where everything is packed together. This is the
///   default storage location when using [`Tlk::from_bytes`], as it reduces
///   memory allocations. Entry text needs to be reallocated using
///   [`Entry::into_owned`] in order to be modified.
#[derive(Debug, Deserialize)]
pub struct Tlk {
    /// File type header
    pub file_type: FixedSizeString<4>,
    /// File version header
    pub file_version: FixedSizeString<4>,
    /// TLK language ID. IDs are different depending on NWN1/2 versions
    pub language_id: u32,
    /// TLK entries list
    pub entries: Vec<Entry>,
    /// When parsed using `Tlk::from_bytes`, contains all strings data referenced by entries
    #[serde(skip)]
    pub strings_data: Vec<u8>,
    ///
    #[serde(default = "default_false")]
    pub hint_is_user: bool,
}

impl Tlk {
    /// Parse TLK binary file format
    ///
    /// * `repair`: attempt to repair the TLK data (invalid UTF8 strings and out of bound pointers)
    pub fn from_bytes(input: &[u8], repair: bool) -> NWNParseResult<Self> {
        // Header
        let (input, file_type) = FixedSizeString::<4>::from_bytes(input)?;
        let (input, file_version) = FixedSizeString::<4>::from_bytes(input)?;
        let (input, language_id) = le_u32(input)?;
        let (input, string_count) = le_u32(input)?;
        let (input, string_entries_offset) = le_u32(input)?;

        let (input, mut entries) = nom_parse_context(
            "while parsing tlk entries",
            count(Entry::from_bytes, string_count as usize)(input),
        )?;
        let strings_data = input.to_vec();

        // Check utf8 strings, convert to TlkString::Owned if invalid
        // Also do strref mapping in preparation for possible JSON/YAML output
        for (strref, entry) in entries.iter_mut().enumerate() {
            if let Some(ref mut text) = &mut entry.text {
                if let TlkString::Ptr { offset, size } = text {
                    let start = *offset as usize;
                    let mut end = start + *size as usize;

                    if repair && end > strings_data.len() {
                        if start < strings_data.len() {
                            end = strings_data.len();
                            *size = (end - start) as u32;
                        } else {
                            entry.text = None;
                            continue;
                        }
                    }

                    if let Some(data) = strings_data.get(start..end) {
                        if let Err(e) = std::str::from_utf8(data) {
                            if repair {
                                entry.text = Some(TlkString::Owned(
                                    String::from_utf8_lossy(data).to_string(),
                                ));
                            } else {
                                return Err(nom::Err::Error(NWNParseError::from(e)));
                            }
                        }
                    } else {
                        return Err(nom::Err::Error(
                            format!(
                                "TLK strref {} is pointing outside of strings_data: text \
                                 address={}, size={}",
                                strref,
                                string_entries_offset + *offset,
                                size
                            )
                            .into(),
                        ));
                    }
                } else {
                    panic!("unexpected owned string");
                }
            }
        }

        Ok((
            input,
            Tlk {
                file_type,
                file_version,
                language_id,
                entries,
                strings_data,
                hint_is_user: false,
            },
        ))
    }

    pub fn write_stream<T>(&self, output: &mut T) -> Result<(), Box<dyn std::error::Error>>
    where
        T: std::io::Write + std::io::Seek,
    {
        output.write_all(self.file_type.as_bytes())?;
        output.write_all(self.file_version.as_bytes())?;
        output.write_all(&self.language_id.to_le_bytes())?;
        output.write_all(&self.len().to_le_bytes())?; // string_count
        output.write_all(&(20u32 + self.len() * 40u32).to_le_bytes())?; // string_entries_offset

        let mut strings_end_offset = 0u32;
        for entry in &self.entries {
            output.write_all(&entry.get_flags().to_le_bytes())?;
            output.write_all(
                entry
                    .sound_resref
                    .as_ref()
                    .unwrap_or(&Default::default())
                    .as_bytes(),
            )?;
            output.write_all(&entry._volume_variance.to_le_bytes())?;
            output.write_all(&entry._pitch_variance.to_le_bytes())?;

            let offset: u32 = strings_end_offset;
            let size: u32 = entry.get_str_len() as u32;
            output.write_all(&offset.to_le_bytes())?;
            output.write_all(&size.to_le_bytes())?;

            output.write_all(&entry.sound_length.unwrap_or(0f32).to_le_bytes())?;

            strings_end_offset += size;
        }

        // Preallocate
        output.seek(std::io::SeekFrom::Current(strings_end_offset as i64))?;
        output.write_all(&[])?;
        output.seek(std::io::SeekFrom::Current(-(strings_end_offset as i64)))?;

        for entry in &self.entries {
            output.write_all(entry.get_str(&self.strings_data).unwrap_or("").as_bytes())?;
        }

        Ok(())
    }

    /// Serialize to TLK binary file format
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut buf = std::io::Cursor::new(vec![]);
        self.write_stream(&mut buf)
            .expect("failed to allocate memory?");
        buf.into_inner()
    }
    /// returns the number of entries in the TLK
    pub fn len(&self) -> u32 {
        self.entries.len() as u32
    }
    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    /// returns a a TLK entry matching the strref. None if the strref does not exist.
    pub fn get_entry(&self, strref: u32) -> Option<&Entry> {
        self.entries.get(strref as usize)
    }
    /// returns a mutable reference to a TLK entry. The entries list is resized if needed
    pub fn get_entry_mut(&mut self, strref: u32) -> &mut Entry {
        if strref as usize >= self.entries.len() {
            self.entries
                .resize_with(strref as usize + 1, Default::default);
        }
        &mut self.entries[strref as usize]
    }
    /// Returns a TLK strref text
    pub fn get_str(&self, strref: u32) -> Option<&str> {
        self.get_entry(strref)?.get_str(&self.strings_data)
    }
    /// Set a TLK strref text
    pub fn set_string(&mut self, strref: u32, text: Option<String>) {
        let entry = &mut self.get_entry_mut(strref);
        entry.text = text.map(TlkString::Owned);
    }

    /// Formats the TLK to text.
    pub fn to_string_pretty(&self) -> String {
        let mut ret = String::new();

        ret += "============================ Header ============================\n";
        ret += &format!(
            "file_type={:?} file_version={:?}\n",
            self.file_type, self.file_version
        );
        ret += &format!("language_id={}\n", self.language_id);
        ret += "============================ Entries ============================\n";

        let max_index = if self.hint_is_user {
            self.entries.len() + TLK_STRREF_USER_OFFSET as usize
        } else {
            self.entries.len()
        };
        let index_padding = (max_index.ilog10() + 1) as usize;

        ret += &self
            .entries
            .iter()
            .enumerate()
            .map(|(index, entry)| {
                entry.to_string_pretty(
                    index as u32,
                    index_padding,
                    self.hint_is_user,
                    &self.strings_data,
                )
            })
            .collect::<Vec<_>>()
            .join("\n");
        ret
    }

    /// Sets if the TLK is a base or user TLK (mostly affects serialization)
    pub fn with_user_hint(mut self, hint: bool) -> Self {
        self.hint_is_user = hint;
        self
    }

    /// Change each entry to remove any reference to `strings_data`, and empty `strings_data`
    //
    /// # Panics
    /// if any entry references data that is out of `strings_data` bounds or is invalid utf8 data
    pub fn into_owned(mut self) -> Self {
        self.entries = self
            .entries
            .into_iter()
            .map(|e| e.into_owned(&self.strings_data))
            .collect();
        self.strings_data.clear();
        self.strings_data.shrink_to_fit();
        self
    }

    /// Removes trailing empty entries
    pub fn shrink_to_fit(&mut self) {
        let mut new_size = self.entries.len();

        for (index, entry) in self.entries.iter().enumerate().rev() {
            if entry.is_empty() {
                new_size = index;
            } else {
                break;
            }
        }

        self.entries.truncate(new_size);
    }

    /// Resizes entries list by adding empty entries if needed
    pub fn resize(&mut self, new_len: u32) {
        self.entries.resize_with(new_len as usize, Default::default);
    }
}

impl Serialize for Tlk {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        use serde::ser::SerializeStruct;
        let mut ser_struct = serializer.serialize_struct("Tlk", 5)?;
        ser_struct.serialize_field("file_type", &self.file_type)?;
        ser_struct.serialize_field("file_version", &self.file_version)?;
        ser_struct.serialize_field("language_id", &self.language_id)?;
        ser_struct.serialize_field("__hint_is_user", &self.hint_is_user)?;
        ser_struct.serialize_field(
            "entries",
            &self
                .entries
                .iter()
                .enumerate()
                .map(|(i, e)| EntrySerde {
                    entry: e,
                    strings_data: &self.strings_data,
                    strref: if !self.hint_is_user {
                        i as u32
                    } else {
                        i as u32 + TLK_STRREF_USER_OFFSET
                    },
                })
                .collect::<Vec<_>>(),
        )?;
        ser_struct.end()
    }
}

/// TLK gender
#[derive(Debug, Default)]
pub enum Gender {
    #[default]
    Male = 0,
    Female = 1,
}

/// Helper struct for resolving strrefs, using a base & user TLKs
#[derive(Debug)]
pub struct Resolver<'tlk> {
    pub base: &'tlk Tlk,
    pub base_f: Option<&'tlk Tlk>,
    pub user: Option<&'tlk Tlk>,
}
impl<'tlk> Resolver<'tlk> {
    /// Create a new resolver using multiple TLK files
    pub fn new(base: &'tlk Tlk, base_f: Option<&'tlk Tlk>, user: Option<&'tlk Tlk>) -> Self {
        Resolver { base, base_f, user }
    }
    /// Returns the TLK used to resolve the given strref
    pub fn get_tlk_for_strref(&self, strref: u32, gender: Gender) -> (Option<&Tlk>, u32) {
        if strref < TLK_STRREF_USER_OFFSET {
            match gender {
                Gender::Male => (Some(self.base), strref),
                Gender::Female => (Some(self.base_f.as_ref().unwrap_or(&self.base)), strref),
            }
        } else {
            let strref = strref - TLK_STRREF_USER_OFFSET;
            if let Some(tlk) = &self.user {
                (Some(tlk), strref)
            } else {
                (None, strref)
            }
        }
    }
    /// Returns the TLK entry corresponding to the given strref
    pub fn get_entry(&self, strref: u32, gender: Gender) -> Option<&Entry> {
        let (tlk, strref) = self.get_tlk_for_strref(strref, gender);
        if let Some(tlk) = tlk {
            if strref < tlk.len() {
                tlk.get_entry(strref)
            } else {
                None
            }
        } else {
            None
        }
    }
    /// Returns the text corresponding to the given strref
    pub fn get_str(&self, strref: u32, gender: Gender) -> Option<&str> {
        let (tlk, strref) = self.get_tlk_for_strref(strref, gender);
        if let Some(tlk) = tlk {
            tlk.get_str(strref)
        } else {
            None
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_user_tlk() {
        let tlk_bytes = include_bytes!("../unittest/user.tlk");

        let tlk = Tlk::from_bytes(tlk_bytes, false).unwrap().1;
        assert_eq!(tlk.get_str(0), Some("Hello world"));
        assert_eq!(tlk.get_str(1), Some("Café liégeois"));
        assert_eq!(tlk.get_str(2), None);
        assert_eq!(tlk.get_str(3), Some("Custom sound"));
        assert_eq!(
            tlk.get_entry(3)
                .unwrap()
                .sound_resref
                .as_ref()
                .unwrap()
                .as_str(),
            "snd_custom"
        );
        assert_eq!(tlk.get_str(10), None);
        assert_eq!(tlk.len(), 4);

        let tlk_bytes_ser = tlk.to_bytes();
        assert_eq!(tlk_bytes_ser, tlk_bytes);

        // Extend tlk table
        let mut tlk = Tlk::from_bytes(&tlk_bytes_ser, false).unwrap().1;
        tlk.set_string(0, Some("Yolo".to_string()));
        tlk.set_string(1, Some("Sôm€ ütf8 🐧".to_string()));
        tlk.set_string(9, Some("Extend this tlk !".to_string()));

        let tlk_bytes_ser = tlk.to_bytes();
        let tlk = Tlk::from_bytes(&tlk_bytes_ser, false).unwrap().1;
        assert_eq!(tlk.len(), 10);
        assert_eq!(tlk.get_str(0), Some("Yolo"));
        assert_eq!(tlk.get_str(1), Some("Sôm€ ütf8 🐧"));
        assert_eq!(tlk.get_str(2), None);
        assert_eq!(tlk.get_str(9), Some("Extend this tlk !"));
    }

    #[test]
    fn test_dialog_tlk() {
        let tlk_bytes = include_bytes!("../unittest/dialog.tlk");

        let tlk = Tlk::from_bytes(tlk_bytes, false).unwrap().1;
        assert_eq!(tlk.get_str(0), Some("Bad Strref"));
        assert_eq!(tlk.get_str(1), Some("Barbares"));
        assert_eq!(tlk.get_str(10), Some("Moine"));
        assert_eq!(tlk.get_str(30), Some("Demi-elfe"));
        assert_eq!(
            tlk.get_str(36),
            Some(concat!(
                "La dernière tentative pour détecter les paramètres de votre ",
                "système a échoué. Cela indique généralement une incompatibilité ",
                "entre Neverwinter Nights 2 et vos pilotes vidéo actuels. Veuillez ",
                "consultez le site internet du fabriquant de votre carte vidéo pour ",
                "télécharger les pilotes les plus récents.\n\nVous pouvez :\n\n1. ",
                "Cliquer sur \"Continuer\" si vous pensez que Neverwinter Nights 2 ",
                "fonctionnera sur votre système. Si le jeu ne démarre pas, relancez ",
                "le programme de configuration et essayez une autre option.\n\n2. ",
                "Cliquez sur \"Assistance\" pour lancer votre navigateur Internet ",
                "afin d'accéder au site web de l'assistance technique de ",
                "Neverwinter Nights 2.\n\n3. Cliquez sur \"LisezMoi\" et consultez ",
                "la section compatibilité pour savoir s'il n'y a pas de problèmes ",
                "de compatibilité entre Neverwinter Nights 2 et votre ",
                "système.\n\n4. Cliquez sur \"Quitter\" pour fermer le ",
                "programme."
            ))
        );
        assert_eq!(tlk.get_str(3000), None);

        let tlk_bytes_ser = tlk.to_bytes();
        assert_eq!(tlk_bytes_ser, tlk_bytes);
    }

    #[test]
    fn test_repair_tlk() {
        let tlk_bytes = include_bytes!("../unittest/user.broken.tlk");

        let tlk = Tlk::from_bytes(tlk_bytes, true).unwrap().1;
        assert_eq!(tlk.get_str(0), Some("Some in�(id UTF8 magic"));
        assert_eq!(
            tlk.get_str(1),
            Some("Some text with bad sizeAnother text but with bad offset")
        );
        assert_eq!(tlk.get_str(2), None);
    }

    #[test]
    fn test_tlk_resolver() {
        let base = Tlk::from_bytes(include_bytes!("../unittest/dialog.tlk"), false)
            .unwrap()
            .1;
        let user = Tlk::from_bytes(include_bytes!("../unittest/user.tlk"), false)
            .unwrap()
            .1;
        let resolver = Resolver::new(&base, None, Some(&user));

        assert_eq!(resolver.get_str(0, Gender::Male), Some("Bad Strref"));
        assert_eq!(resolver.get_str(1, Gender::Male), Some("Barbares"));
        assert_eq!(resolver.get_str(1_000_000, Gender::Male), None);
        assert_eq!(
            resolver.get_str(16_777_216, Gender::Male),
            Some("Hello world")
        );
        assert_eq!(
            resolver.get_str(16_777_217, Gender::Male),
            Some("Café liégeois")
        );
        assert_eq!(resolver.get_str(17_777_216, Gender::Male), None);
    }

    // extern crate test;
    // #[bench]
    // fn bench_dialog_tlk(b: &mut test::Bencher) {
    //     b.iter(test_dialog_tlk);
    // }
}