dmio 0.1.4

/*
 * Copyright 2018 Christian Ebner
 *
 * This file is part of dmio.
 *
 * dmio is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * dmio is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with dmio.  If not, see <http://www.gnu.org/licenses/>.
 */

use std::io::{Cursor, BufReader, Read};
use std::fs::File;
use std::char;

use byteorder::{ByteOrder, BigEndian, LittleEndian, ReadBytesExt};
use anyhow::bail;
use log::debug;

use crate::taggroup_entries::{Key, Tag, TagGroup};
use crate::{DMImage, Result};
use crate::dmtypes::*;

pub trait DMImageReadByVersion {
    fn read_by_version(reader: &mut DMImageReader) ->  Result<usize>;
    fn read_tag_entry_size(reader: &mut DMImageReader) ->  Result<usize>;
}

impl DMImageReadByVersion for DM3Reader {
    #[inline]
    fn read_by_version(reader: &mut DMImageReader) -> Result<usize> {
        let value = reader.buffer.read_u32::<BigEndian>()? as usize;

        Ok(value)
    }

    // DM3 doesn't have this field, so simply return 0
    #[inline]
    fn read_tag_entry_size(_reader: &mut DMImageReader) ->  Result<usize> {
        Ok(0)
    }

}

impl DMImageReadByVersion for DM4Reader {
    #[inline]
    fn read_by_version(reader: &mut DMImageReader) -> Result<usize> {
        let value = reader.buffer.read_u64::<BigEndian>()? as usize;

        Ok(value)
    }

    #[inline]
    fn read_tag_entry_size(reader: &mut DMImageReader) ->  Result<usize> {
        let value = reader.buffer.read_u64::<BigEndian>()? as usize;

        Ok(value)
    }
}


pub struct DMImageReader {
    buffer: Cursor<Vec<u8>>,
}

impl DMImageReader {
    /// Create a new image parser, loading the full file into a memory buffer and interpreting it
    #[inline]
    pub fn new(file: File) -> Result<DMImageReader> {
        debug!("DMImageReader::new() on {:?}", file);
        // Read the full file into a byte buffer
        let mut bufrdr = BufReader::new(file);
        let mut buffer = Vec::new();
        let bytes = bufrdr.read_to_end(&mut buffer)?;
        debug!("DMImageReader::new() -> {} bytes read from file", bytes);

        Ok(DMImageReader {
            buffer: Cursor::new(buffer),
        })
    }

    /// Parse the buffer loaded by DMImageParser::new() and return the DMImage.
    /// On success return the DMImage wrapped into a Result.
    /// On error the error specifics are returned wrapped into the Result.
    pub fn parse(&mut self) -> Result<DMImage> {
        debug!("DMImageReader::parse()");
        // The first 4 bytes define the image version used for encoding
        let version = self.buffer.read_u32::<BigEndian>()? as usize;
        // The size of the TagGroup has to be read version dependent
        let tg_size = match version {
            DM3 => DM3Reader::read_by_version(self)?,
            DM4 => DM4Reader::read_by_version(self)?,
            _ => bail!("Unsupported DM image version detected."),
        };
        // Finally the header states the endianess of the image encoding
        let bigendian = match self.buffer.read_u32::<BigEndian>()? as usize {
                BIG_ENDIAN => true,
                _ => false,
        };
        let root = match (version, bigendian) {
            (DM3, true) => self.read_taggroup::<DM3Reader, BigEndian>()?,
            (DM4, true) => self.read_taggroup::<DM4Reader, BigEndian>()?,
            (DM3, false) => self.read_taggroup::<DM3Reader, LittleEndian>()?,
            (DM4, false) => self.read_taggroup::<DM4Reader, LittleEndian>()?,
            _ => bail!("Unsupported DM image version detected."),
        };

        Ok(DMImage {
            path: None,
            root,
            data: Tag::Empty,
            sizex: Tag::Empty,
            sizey: Tag::Empty,
            sizez: Tag::Empty,
            version,
            tg_size,
            bigendian,
        })
    }
    
    #[inline]
    fn read_taggroup<R: DMImageReadByVersion, T: ByteOrder>(&mut self) -> Result<TagGroup> {
        debug!("DMImageReader::read_taggroup()");
        // TagGroups consist of a header of 6 bytes followed by the data
        let is_sorted = match self.buffer.read_u8()? {
            B_TRUE => true,
            _ => false,
        };

        let is_open = match self.buffer.read_u8()? {
            B_TRUE => true,
            _ => false,
        };

        let nr_tags = R::read_by_version(self)?;
        let mut tg = TagGroup::new(is_sorted, is_open);
        // Read the individual entries of the TagGroup
        for _ in 0..nr_tags {
            let (key, tag) = self.read_entry::<R, T>(tg.len())?;
            tg.insert(key, tag);
        }

        Ok(tg)
    }

    #[inline]
    fn read_entry<R: DMImageReadByVersion, T: ByteOrder>(&mut self, index: usize) -> Result<(Key, Tag)> {
        debug!("DMImageReader::read_entry()");
        // Entries consist of a 1 byte for type, 2 bytes for string label lenght, n bytes of the
        // label string followed by the tag instance, which is a TagGroup instance or a TagType
        // (data) instance

        let tag_type = self.buffer.read_u8()?;
        let label_length = self.buffer.read_u16::<BigEndian>()? as usize;
        let mut label_buf = vec![0u8; label_length];
        self.buffer.read_exact(&mut label_buf)?;

        let label = if label_buf.is_empty() {
            Key::Index(index)
        } else {
            Key::Label(unsafe { String::from_utf8_unchecked(label_buf) })
        };
        debug!("Label: {}", label);

        // DM4 has an additional field which specifies the size of the following
        // tagentry so it could be skipped if wanted.
        // TODO For now simply ignoring this, might be an interesting feature
        // for the future to speed up parsing, by avoiding unneeded tags
        let _tag_entry_size = R::read_tag_entry_size(self)?;

        let value = match tag_type {
            T_TAG => self.read_tag::<R, T>()?,
            T_GROUP => Tag::TagGroupEntry(self.read_taggroup::<R, T>()?),
            _ => bail!("Can not identify tag type"),
        };

        Ok((label, value))
    }

    #[inline]
    fn read_tag<R: DMImageReadByVersion, T: ByteOrder>(&mut self) -> Result<Tag> {
        debug!("DMImageReader::read_tag()");
        // Tags consist of a 4 bytes for spacer ')%%%%'
        // 4 or 8 bytes n = length of definition of encoded type (DM3, DM4)
        // where simple type = 1, string = 2, array = 3, structs 1+2*f with f=number of struct fields
        // 4*n or 8*n bytes where n = number of encoded types from before; each 4 or 8 bytes (DM3, DM4)
        // define type by ID, finally the data as definde in the TagType

        // First 4 bytes are spacer '%%%%'... Lets ignore them
        let _spacer = self.buffer.read_u32::<BigEndian>()?;

        let enc_info_length = R::read_by_version(self)?;
        match enc_info_length {
            T_SIMPLE => self.read_simple_data::<R, T>(),
            T_STRING => self.read_string::<R>(),
            T_ARRAY => self.read_simple_array::<R, T>(),
            _ => {
                match R::read_by_version(self)? {
                    D_STRUCT => self.read_struct::<R, T>(),
                    _ => self.read_complex_array::<R, T>(),
                }
            }
        }
    }

    #[inline]
    fn read_simple_data<R: DMImageReadByVersion, T: ByteOrder>(&mut self) -> Result<Tag> {
        debug!("DMImageReader::read_simple_data()");
        let data_type = R::read_by_version(self)?;
        debug!("data type: {}", data_type);
        let value = match data_type {
            D_SHORT => Tag::Short(self.buffer.read_i16::<T>()?),
            D_LONG => Tag::Long(self.buffer.read_i32::<T>()?),
            D_USHORT => Tag::UShort(self.buffer.read_u16::<T>()?),
            D_ULONG => Tag::ULong(self.buffer.read_u32::<T>()?),
            D_FLOAT => Tag::Float(self.buffer.read_f32::<T>()?),
            D_DOUBLE => Tag::Double(self.buffer.read_f64::<T>()?),
            D_BOOLEAN => Tag::Boolean(match self.buffer.read_u8()? {
                B_TRUE => true,
                _ => false,
            }),
            D_CHAR => Tag::Char({
                let val = self.buffer.read_u8()?;
                match char::from_u32(u32::from(val)) {
                    Some(ch) => ch,
                    None => panic!("Cannot convert from u32 to char"),
                }
            }),
            D_OCTET => Tag::Octet(self.buffer.read_u8()?),
            D_ULONGLONG => Tag::ULongLong(self.buffer.read_u64::<T>()?),
            D_UNKNOWN64 => Tag::Unknown64(self.buffer.read_u64::<T>()?),
            _ => bail!("Type not known"),
        };

        Ok(value)
    }

    #[inline]
    fn read_string<R: DMImageReadByVersion>(&mut self) -> Result<Tag> {
        debug!("DMImageReader::read_string()");
        // Check if this is indeed a string type
        match R::read_by_version(self)? {
            D_STRING => (),
            _ => bail!("Unexpected type for string"),
        };
        let string_length = R::read_by_version(self)?;
        let mut string_buf = vec![0u8; 2 * string_length];
        self.buffer.read_exact(&mut string_buf)?;
        // TODO need to change this to utf16
        let string = String::from_utf8(string_buf).unwrap();

        Ok(Tag::CharSeq(string))
    }

    #[inline]
    fn read_simple_array<R: DMImageReadByVersion, T: ByteOrder>(&mut self) -> Result<Tag> {
        debug!("DMImageReader::read_simple_array()");
        // Read and ignore type info, should be 20
        let _type = R::read_by_version(self)?;
        let element_t = R::read_by_version(self)?;
        let size = R::read_by_version(self)?;
        let value = match element_t {
            D_SHORT => {
                let mut array = vec![0i16; size];
                for elem in &mut array {
                    *elem = self.buffer.read_i16::<T>()?;
                }
                Tag::ArrayShort(array)
            }
            D_LONG => {
                let mut array = vec![0i32; size];
                for elem in &mut array {
                    *elem = self.buffer.read_i32::<T>()?;
                }
                Tag::ArrayLong(array)
            }
            D_USHORT => {
                let mut array = vec![0u16; size];
                for elem in &mut array {
                    *elem = self.buffer.read_u16::<T>()?;
                }
                Tag::ArrayUShort(array)
            }
            D_ULONG => {
                let mut array = vec![0u32; size];
                for elem in &mut array {
                    *elem = self.buffer.read_u32::<T>()?;
                }
                Tag::ArrayULong(array)
            }
            D_FLOAT => {
                let mut array = vec![0f32; size];
                for elem in &mut array {
                    *elem = self.buffer.read_f32::<T>()?;
                }
                Tag::ArrayFloat(array)
            }
            D_DOUBLE => {
                let mut array = vec![0f64; size];
                for elem in &mut array {
                    *elem = self.buffer.read_f64::<T>()?;
                }
                Tag::ArrayDouble(array)
            }
            D_BOOLEAN => {
                let mut array = vec![false; size];
                for elem in &mut array {
                    *elem = match self.buffer.read_u8()? {
                        B_TRUE => true,
                        _ => false,
                    };
                }
                Tag::ArrayBoolean(array)
            }
            D_CHAR => {
                let mut string_buf = vec![0u8; size];
                self.buffer.read_exact(&mut string_buf)?;
                let string = String::from_utf8(string_buf).unwrap();
                Tag::ArrayChar(string)
            }
            D_OCTET => {
                let mut val_buf = vec![0u8; size];
                self.buffer.read_exact(&mut val_buf)?;
                Tag::ArrayOctet(val_buf)
            }
            D_ULONGLONG => {
                let mut array = vec![0u64; size];
                for elem in &mut array {
                    *elem = self.buffer.read_u64::<T>()?;
                }
                Tag::ArrayULongLong(array)
            }
            D_UNKNOWN64 => {
                let mut array = vec![0u64; size];
                for elem in &mut array {
                    *elem = self.buffer.read_u64::<T>()?;
                }
                Tag::ArrayUnknown64(array)
            }
            _ => bail!("Type not known"),
        };

        Ok(value)
    }

    #[inline]
    fn read_struct<R: DMImageReadByVersion, T: ByteOrder>(&mut self) -> Result<Tag> {
        debug!("DMImageReader::read_struct()");
        let _namelength = R::read_by_version(self)?;
        let number_fields = R::read_by_version(self)?;
        let mut fields = vec![(0usize, 0usize); number_fields];
        for field in &mut fields {
            let namelength = R::read_by_version(self)?;
            let field_t = R::read_by_version(self)?;
            *field = (namelength, field_t);
        }
        debug!("fields: {:?}", fields);

        let mut tg_struct: Vec<Tag> = Vec::new();
        for (namelength, field_t) in fields {
            let mut name_buf = vec![0u8; namelength];
            self.buffer.read_exact(&mut name_buf)?;

            match field_t {
                D_SHORT => {
                    let value = self.buffer.read_i16::<T>()?;
                    tg_struct.push(Tag::Short(value));
                }
                D_LONG => {
                    let value = self.buffer.read_i32::<T>()?;
                    tg_struct.push(Tag::Long(value));
                }
                D_USHORT => {
                    let value = self.buffer.read_u16::<T>()?;
                    tg_struct.push(Tag::UShort(value));
                }
                D_ULONG => {
                    let value = self.buffer.read_u32::<T>()?;
                    tg_struct.push(Tag::ULong(value));
                }
                D_FLOAT => {
                    let value = self.buffer.read_f32::<T>()?;
                    tg_struct.push(Tag::Float(value));
                }
                D_DOUBLE => {
                    let value = self.buffer.read_f64::<T>()?;
                    tg_struct.push(Tag::Double(value));
                }
                D_BOOLEAN => {
                    let value = self.buffer.read_u8()? == B_TRUE;
                    tg_struct.push(Tag::Boolean(value));
                }
                D_CHAR => {
                    let character = char::from_u32(u32::from(self.buffer.read_u8()?)).unwrap();
                    tg_struct.push(Tag::Char(character));
                }
                D_OCTET => {
                    let value = self.buffer.read_u8()?;
                    tg_struct.push(Tag::Octet(value));
                }
                D_ULONGLONG => {
                    let value = self.buffer.read_u64::<T>()?;
                    tg_struct.push(Tag::ULongLong(value));
                }
                D_UNKNOWN64 => {
                    let value = self.buffer.read_u64::<T>()?;
                    tg_struct.push(Tag::Unknown64(value));
                }
                _ => bail!("Type not known"),
            }
        }

        Ok(Tag::Struct(tg_struct))
    }

    #[inline]
    fn read_complex_array<R: DMImageReadByVersion, T: ByteOrder>(&mut self) -> Result<Tag> {
        debug!("DMImageReader::read_complex_array()");
        let _type = R::read_by_version(self)?;
        let _namelength = R::read_by_version(self)?;
        let struct_elements = R::read_by_version(self)?;

        let mut fields = vec![(0usize, 0usize); struct_elements];
        for field in &mut fields {
            let namelength = R::read_by_version(self)?;
            let field_t = R::read_by_version(self)?;
            *field = (namelength, field_t);
        }

        let array_elements = R::read_by_version(self)?;
        let mut tg_cplx_arr: Vec<Tag> = Vec::new();
        for _ in 0..array_elements {
            let mut tg_cplx_entry: Vec<Tag> = Vec::new();
            for &(_, t) in &fields {
                match t {
                    D_SHORT => {
                        let value = self.buffer.read_i16::<T>()?;
                        tg_cplx_entry.push(Tag::Short(value));
                    }
                    D_LONG => {
                        let value = self.buffer.read_i32::<T>()?;
                        tg_cplx_entry.push(Tag::Long(value));
                    }
                    D_USHORT => {
                        let value = self.buffer.read_u16::<T>()?;
                        tg_cplx_entry.push(Tag::UShort(value));
                    }
                    D_ULONG => {
                        let value = self.buffer.read_u32::<T>()?;
                        tg_cplx_entry.push(Tag::ULong(value));
                    }
                    D_FLOAT => {
                        let value = self.buffer.read_f32::<T>()?;
                        tg_cplx_entry.push(Tag::Float(value));
                    }
                    D_DOUBLE => {
                        let value = self.buffer.read_f64::<T>()?;
                        tg_cplx_entry.push(Tag::Double(value));
                    }
                    D_BOOLEAN => {
                        let value = match self.buffer.read_u8()? {
                            B_TRUE => true,
                            _ => false,
                        };
                        tg_cplx_entry.push(Tag::Boolean(value));
                    }
                    D_CHAR => {
                        let character = char::from_u32(u32::from(self.buffer.read_u8()?)).unwrap();
                        tg_cplx_entry.push(Tag::Char(character));
                    }
                    D_OCTET => {
                        let val_buf = self.buffer.read_u8()?;
                        tg_cplx_entry.push(Tag::Octet(val_buf));
                    }
                    D_ULONGLONG => {
                        let value = self.buffer.read_u64::<T>()?;
                        tg_cplx_entry.push(Tag::ULongLong(value));
                    }
                    D_UNKNOWN64 => {
                        let value = self.buffer.read_u64::<T>()?;
                        tg_cplx_entry.push(Tag::Unknown64(value));
                    }
                    _ => bail!("Type not known"),
                }
            }
            tg_cplx_arr.push(Tag::Struct(tg_cplx_entry));
        }

        Ok(Tag::ComplexArray(tg_cplx_arr))
    }
}