msg_parser 0.3.6

use std::collections::HashMap;

use crate::ole::{Entry, EntryType, Reader};

use super::{
    constants::PropIdNameMap, decode::DataType, named_prop::NamedPropertyMap, stream::Stream,
};

// StorageType refers to major components in Message object.
// Refer to MS-OXPROPS 1.3.3
#[derive(Debug, Clone, PartialEq)]
pub enum StorageType {
    // u32 refers to its index
    Recipient(u32),
    // u32 refers to its index
    Attachment(u32),
    RootEntry,
}

impl StorageType {
    fn convert_id_to_u32(id: &str) -> Option<u32> {
        // id is 8 digits hexadecimal sequence.
        if id.len() != 8 {
            return None;
        }
        u32::from_str_radix(id, 16).ok()
    }

    pub fn create(name: &str) -> Option<Self> {
        if name.starts_with("__recip_version1.0_") {
            // Extract the digits after '#' in __recip_version1.0_#00000000
            // Remaining digits is the index of Recipient.
            let id = name.split('#').nth(1)?;
            let id_as_num = StorageType::convert_id_to_u32(id)?;
            return Some(StorageType::Recipient(id_as_num));
        }
        if name.starts_with("__attach_version1.0_") {
            let id = name.split('#').nth(1)?;
            let id_as_num = StorageType::convert_id_to_u32(id)?;
            return Some(StorageType::Attachment(id_as_num));
        }
        None
    }
}

// EntryStorageMap represents HashMap of ole::Entry id and its StorageType
#[derive(Debug)]
struct EntryStorageMap {
    map: HashMap<u32, StorageType>,
}

impl EntryStorageMap {
    pub fn new(parser: &Reader) -> Self {
        let mut storage_map: HashMap<u32, StorageType> = HashMap::new();
        for entry in parser.iterate() {
            match entry._type() {
                EntryType::RootStorage => {
                    storage_map.insert(entry.id(), StorageType::RootEntry);
                }
                EntryType::UserStorage => {
                    StorageType::create(entry.name())
                        .and_then(|storage| storage_map.insert(entry.id(), storage));
                }
                _ => {
                    continue;
                }
            }
        }
        Self { map: storage_map }
    }

    pub fn get_storage_type(&self, parent_id: Option<u32>) -> Option<&StorageType> {
        self.map.get(&parent_id?)
    }
}

// Properties is a Map is a collection of Message object elements.
pub type Properties = HashMap<String, DataType>;

// Recipients represent array of Recipient objects in Message.
pub type Recipients = Vec<Properties>;

// Attachments represent array of Attachment object in Message
pub type Attachments = Vec<Properties>;

// Storages is a collection of Storage
// object containing their decoded stream
// values for respective properties.
#[derive(Debug)]
pub struct Storages {
    storage_map: EntryStorageMap,
    prop_map: &'static PropIdNameMap,
    named_props: NamedPropertyMap,
    pub attachments: Attachments,
    pub recipients: Recipients,
    // Mail properties
    pub root: Properties,
}

impl Storages {
    fn to_arr(map: HashMap<u32, Properties>) -> Vec<Properties> {
        let mut tuples: Vec<(u32, Properties)> =
            map.into_iter().collect::<Vec<(u32, Properties)>>();
        tuples.sort_by(|a, b| a.0.cmp(&b.0));
        tuples.into_iter().map(|x| x.1).collect::<Vec<Properties>>()
    }

    fn create_stream(&self, parser: &Reader, entry: &Entry) -> Option<Stream> {
        let parent = self.storage_map.get_storage_type(entry.parent_node())?;
        let mut slice = parser.get_entry_slice(entry).ok()?;
        Stream::create(
            entry.name(),
            &mut slice,
            self.prop_map,
            &self.named_props,
            parent,
        )
    }

    /// Parse fixed-size properties from a __properties_version1.0 stream.
    /// Non-root storages have an 8-byte header, then 16-byte entries.
    /// Each entry: 4 bytes prop_tag (type u16 + id u16), 4 bytes flags, 8 bytes value.
    fn parse_fixed_props(
        data: &[u8],
        prop_map: &PropIdNameMap,
        named_props: &NamedPropertyMap,
        is_root: bool,
    ) -> Properties {
        let header_size = if is_root { 32 } else { 8 };
        let mut props = Properties::new();
        if data.len() < header_size {
            return props;
        }
        let mut offset = header_size;
        while offset + 16 <= data.len() {
            let prop_type = u16::from_le_bytes([data[offset], data[offset + 1]]);
            let prop_id = u16::from_le_bytes([data[offset + 2], data[offset + 3]]);
            let value_bytes = &data[offset + 8..offset + 16];

            // Try standard prop map first, then named props for 0x8000+ range
            let id_str = format!("0x{:04X}", prop_id);
            let name: Option<&str> = prop_map
                .get_canonical_name(&id_str)
                .or_else(|| named_props.get(prop_id));
            if let Some(name) = name {
                match prop_type {
                    // PtypInteger32
                    0x0003 => {
                        let val = u32::from_le_bytes([
                            value_bytes[0],
                            value_bytes[1],
                            value_bytes[2],
                            value_bytes[3],
                        ]);
                        props.insert(name.to_string(), DataType::PtypInteger32(val));
                    }
                    // PtypTime (FILETIME, 8 bytes)
                    0x0040 => {
                        let val = u64::from_le_bytes([
                            value_bytes[0],
                            value_bytes[1],
                            value_bytes[2],
                            value_bytes[3],
                            value_bytes[4],
                            value_bytes[5],
                            value_bytes[6],
                            value_bytes[7],
                        ]);
                        props.insert(name.to_string(), DataType::PtypTime(val));
                    }
                    _ => {}
                }
            }
            offset += 16;
        }
        props
    }

    pub fn process_streams(&mut self, parser: &Reader) {
        let mut recipients_map: HashMap<u32, Properties> = HashMap::new();
        let mut attachments_map: HashMap<u32, Properties> = HashMap::new();
        for entry in parser.iterate() {
            if let EntryType::UserStream = entry._type() {
                // Parse __properties_version1.0 for fixed-size properties
                if entry.name() == "__properties_version1.0"
                    && let Some(parent) = self.storage_map.get_storage_type(entry.parent_node())
                    && let Ok(mut slice) = parser.get_entry_slice(entry)
                {
                    let mut data = vec![0u8; slice.len()];
                    if std::io::Read::read_exact(&mut slice, &mut data).is_ok() {
                        let is_root = matches!(parent, StorageType::RootEntry);
                        let fixed = Self::parse_fixed_props(
                            &data,
                            self.prop_map,
                            &self.named_props,
                            is_root,
                        );
                        match parent {
                            StorageType::Recipient(id) => {
                                recipients_map.entry(*id).or_default().extend(fixed);
                            }
                            StorageType::Attachment(id) => {
                                attachments_map.entry(*id).or_default().extend(fixed);
                            }
                            StorageType::RootEntry => {
                                self.root.extend(fixed);
                            }
                        }
                    }
                    continue;
                }

                // Decode stream from slice.
                // Skip if failed.
                let Some(stream) = self.create_stream(parser, entry) else {
                    continue;
                };

                // Populate maps accordingly
                match stream.parent {
                    StorageType::RootEntry => {
                        self.root.insert(stream.key, stream.value);
                    }
                    StorageType::Recipient(id) => {
                        let recipient_map = recipients_map.entry(id).or_default();
                        (*recipient_map).insert(stream.key, stream.value);
                    }
                    StorageType::Attachment(id) => {
                        let attachment_map = attachments_map.entry(id).or_default();
                        (*attachment_map).insert(stream.key, stream.value);
                    }
                }
            }
        }
        // Update storages
        self.recipients = Self::to_arr(recipients_map);
        self.attachments = Self::to_arr(attachments_map);
    }

    pub fn new(parser: &Reader) -> Self {
        let root: Properties = HashMap::new();
        let recipients: Recipients = vec![];
        let attachments: Attachments = vec![];
        let storage_map = EntryStorageMap::new(parser);
        let prop_map = PropIdNameMap::init();
        let named_props = Self::build_named_props(parser);
        Self {
            storage_map,
            prop_map,
            named_props,
            root,
            recipients,
            attachments,
        }
    }

    /// Read the __nameid_version1.0 streams and build the named property map.
    fn build_named_props(parser: &Reader) -> NamedPropertyMap {
        use std::io::Read;

        let mut nameid_id = None;
        for entry in parser.iterate() {
            if entry.name() == "__nameid_version1.0" {
                nameid_id = Some(entry.id());
                break;
            }
        }

        let Some(nid) = nameid_id else {
            return NamedPropertyMap::default();
        };

        let mut guid_stream = Vec::new();
        let mut entry_stream = Vec::new();
        let mut string_stream = Vec::new();

        for entry in parser.iterate() {
            if entry.parent_node() != Some(nid) {
                continue;
            }
            if let Ok(mut slice) = parser.get_entry_slice(entry) {
                let mut buf = vec![0u8; slice.len()];
                let _ = slice.read(&mut buf);
                match entry.name() {
                    "__substg1.0_00020102" => guid_stream = buf,
                    "__substg1.0_00030102" => entry_stream = buf,
                    "__substg1.0_00040102" => string_stream = buf,
                    _ => {}
                }
            }
        }

        NamedPropertyMap::parse(&guid_stream, &entry_stream, &string_stream)
    }

    /// Returns the set of all resolved named property names (0x8000+ range).
    pub fn named_property_names(&self) -> std::collections::HashSet<&str> {
        self.named_props.all_names()
    }

    pub fn get_val_from_root_or_default(&self, key: &str) -> String {
        self.root.get(key).map_or(String::new(), |x| x.into())
    }

    pub fn get_root_int_prop(&self, key: &str) -> Option<u32> {
        match self.root.get(key) {
            Some(DataType::PtypInteger32(v)) => Some(*v),
            _ => None,
        }
    }

    pub fn get_recipient_int_prop(&self, idx: usize, key: &str) -> Option<u32> {
        self.recipients.get(idx).and_then(|r| match r.get(key) {
            Some(DataType::PtypInteger32(v)) => Some(*v),
            _ => None,
        })
    }

    pub fn get_attachment_int_prop(&self, idx: usize, key: &str) -> Option<u32> {
        self.attachments.get(idx).and_then(|a| match a.get(key) {
            Some(DataType::PtypInteger32(v)) => Some(*v),
            _ => None,
        })
    }

    pub fn get_bytes_from_attachment(&self, idx: usize, key: &str) -> Vec<u8> {
        self.attachments
            .get(idx)
            .and_then(|attach| match attach.get(key) {
                Some(DataType::PtypBinary(bytes)) => Some(bytes.clone()),
                _ => None,
            })
            .unwrap_or_default()
    }

    pub fn get_val_from_attachment_or_default(&self, idx: usize, key: &str) -> String {
        self.attachments
            .get(idx)
            .map(|attach| attach.get(key).map_or(String::new(), |x| x.into()))
            .unwrap_or_default()
    }
}

#[cfg(test)]
mod tests {
    use super::super::decode::DataType;
    use super::{EntryStorageMap, Properties, StorageType, Storages};
    use crate::ole::Reader;
    use std::collections::HashMap;

    #[test]
    fn test_storage_type_convert() {
        let mut id = StorageType::convert_id_to_u32("00000001");
        assert_eq!(id, Some(1u32));

        id = StorageType::convert_id_to_u32("0000000A");
        assert_eq!(id, Some(10u32));

        id = StorageType::convert_id_to_u32("00000101");
        assert_eq!(id, Some(257u32));

        id = StorageType::convert_id_to_u32("FFFFFFFF");
        assert_eq!(id, Some(u32::MAX));

        // Edge Cases
        id = StorageType::convert_id_to_u32("HELLO");
        assert_eq!(id, None);

        id = StorageType::convert_id_to_u32("00000000000000");
        assert_eq!(id, None);
    }

    #[test]
    fn test_create_storage_type() {
        let recipient = StorageType::create("__recip_version1.0_#0000000A");
        assert_eq!(recipient, Some(StorageType::Recipient(10)));

        let attachment = StorageType::create("__attach_version1.0_#0000000A");
        assert_eq!(attachment, Some(StorageType::Attachment(10)));

        let unknown_storage = StorageType::create("");
        assert_eq!(unknown_storage, None);
    }

    #[test]
    fn test_storage_map() {
        let parser = Reader::from_path("data/test_email.msg").unwrap();
        let storage_map = EntryStorageMap::new(&parser);

        let mut expected_map = HashMap::new();
        expected_map.insert(0, StorageType::RootEntry);
        expected_map.insert(73, StorageType::Recipient(0));
        expected_map.insert(85, StorageType::Recipient(1));
        expected_map.insert(97, StorageType::Recipient(2));
        expected_map.insert(108, StorageType::Recipient(3));
        expected_map.insert(120, StorageType::Recipient(4));
        expected_map.insert(132, StorageType::Recipient(5));
        expected_map.insert(143, StorageType::Attachment(0));
        expected_map.insert(260, StorageType::Recipient(0));
        expected_map.insert(310, StorageType::Attachment(1));
        expected_map.insert(323, StorageType::Attachment(2));
        assert_eq!(storage_map.map, expected_map);
    }

    #[test]
    fn test_storage_to_arr() {
        let mut map_apple: Properties = HashMap::new();
        map_apple.insert("A".to_string(), DataType::PtypString("Apple".to_string()));
        let mut map_bagel: Properties = HashMap::new();
        map_bagel.insert("B".to_string(), DataType::PtypString("Bagel".to_string()));

        let mut basket: HashMap<u32, Properties> = HashMap::new();
        basket.insert(1, map_apple);
        basket.insert(0, map_bagel);

        let res = Storages::to_arr(basket);
        assert_eq!(
            res[0].get("B"),
            Some(&DataType::PtypString("Bagel".to_string()))
        );
        assert_eq!(
            res[1].get("A"),
            Some(&DataType::PtypString("Apple".to_string()))
        );
    }

    #[test]
    fn test_create_storage_test_email() {
        let parser = Reader::from_path("data/test_email.msg").unwrap();
        let mut storages = Storages::new(&parser);
        storages.process_streams(&parser);

        let sender = storages.root.get("SenderEmailAddress");
        assert!(sender.is_none());

        // Check attachments
        assert_eq!(storages.attachments.len(), 3);

        // Check recipients
        assert_eq!(storages.recipients.len(), 6);

        // Check Display name
        let display_name = storages.recipients[0].get("DisplayName").unwrap();
        assert_eq!(
            display_name,
            &DataType::PtypString("marirs@outlook.com".to_string())
        );
    }

    #[test]
    fn test_create_storage_outlook_attachments() {
        let parser = Reader::from_path("data/test_email.msg").unwrap();
        let mut storages = Storages::new(&parser);
        storages.process_streams(&parser);

        // Check attachment
        assert_eq!(storages.attachments.len(), 3);

        let attachment_name = storages.attachments[0].get("DisplayName");
        assert_eq!(
            attachment_name,
            Some(&DataType::PtypString(
                "1 Days Left—35% off cloud space, upgrade now!".to_string()
            ))
        );

        let attachment_name = storages.attachments[1].get("AttachFilename");
        assert_eq!(
            attachment_name,
            Some(&DataType::PtypString("milky-~1.jpg".to_string()))
        );

        let attachment_name = storages.attachments[2].get("AttachFilename");
        assert_eq!(
            attachment_name,
            Some(&DataType::PtypString("TestEm~1.msg".to_string()))
        );

        // Check recipients
        assert_eq!(storages.recipients.len(), 6);
        let display_name = storages.recipients[1].get("DisplayName").unwrap();
        assert_eq!(
            display_name,
            &DataType::PtypString("Sriram Govindan".to_string())
        );
    }

    #[test]
    fn test_named_properties_resolved() {
        let parser = Reader::from_path("data/test_email.msg").unwrap();
        let mut storages = Storages::new(&parser);
        storages.process_streams(&parser);

        // Named properties (0x8000+ range) should be resolved and present in root
        // test_email.msg has well-known named props like InternetAccountName, ReminderSet, etc.
        let has_named = storages.root.keys().any(|k| {
            matches!(
                k.as_str(),
                "InternetAccountName"
                    | "InternetAccountStamp"
                    | "ReminderSet"
                    | "SmartNoAttach"
                    | "SideEffects"
                    | "Private"
            )
        });
        assert!(
            has_named,
            "Expected at least one well-known named property in root. Keys: {:?}",
            storages.root.keys().collect::<Vec<_>>()
        );
    }
}