use super::ElfParser;
use crate::header;

impl ElfParser {
    /// parses an 64-bit ELF header.
    pub(super) fn parse_elf64_header<'a>(
        &'a self,
        b: &'a [u8],
    ) -> nom::IResult<&'a [u8], header::Elf64Header> {
        self.parse_raw_elf64_header(b)
            .map(|(b, raw_header)| (b, header::Elf64Header::from(raw_header)))
    }

    /// parses an 64-bit ELF header as the raw representation.
    pub(super) fn parse_raw_elf64_header<'a>(
        &'a self,
        b: &'a [u8],
    ) -> nom::IResult<&'a [u8], header::RawElf64Header> {
        let (b, e_ident) = self.parse_raw_elf_identification()(b)?;
        let (b, e_type) = self.parse_elf64_half(e_ident.e_data, b)?;
        let (b, e_machine) = self.parse_elf64_half(e_ident.e_data, b)?;
        let (b, e_version) = self.parse_elf64_word(e_ident.e_data, b)?;
        let (b, e_entry) = self.parse_elf64_address(e_ident.e_data, b)?;
        let (b, e_phoff) = self.parse_elf64_offset(e_ident.e_data, b)?;
        let (b, e_shoff) = self.parse_elf64_offset(e_ident.e_data, b)?;
        let (b, e_flags) = self.parse_elf64_word(e_ident.e_data, b)?;
        let (b, e_ehsize) = self.parse_elf64_half(e_ident.e_data, b)?;
        let (b, e_phentsize) = self.parse_elf64_half(e_ident.e_data, b)?;
        let (b, e_phnum) = self.parse_elf64_half(e_ident.e_data, b)?;
        let (b, e_shentsize) = self.parse_elf64_half(e_ident.e_data, b)?;
        let (b, e_shnum) = self.parse_elf64_half(e_ident.e_data, b)?;
        let (b, e_shstrndx) = self.parse_elf64_half(e_ident.e_data, b)?;

        let ehdr = header::RawElf64Header {
            e_ident,
            e_type,
            e_machine,
            e_version,
            e_entry,
            e_phoff,
            e_shoff,
            e_flags,
            e_ehsize,
            e_phentsize,
            e_phnum,
            e_shentsize,
            e_shnum,
            e_shstrndx,
        };

        Ok((b, ehdr))
    }

    /// tries to parse the e_ident field in ELF header as the raw representation.
    /// But the input bytes aren't consumed.
    pub(super) fn peek_raw_elf_identification<'a>(
        &'a self,
        b: &'a [u8],
    ) -> nom::IResult<&'a [u8], header::RawElfIdentification> {
        let mut p = nom::combinator::peek(self.parse_raw_elf_identification());
        let (b, e_ident) = p(b)?;

        Ok((b, e_ident))
    }

    /// parses the e_ident field in ELF header as the raw representation.
    fn parse_raw_elf_identification<'a>(
        &'a self,
    ) -> impl Fn(&'a [u8]) -> nom::IResult<&'a [u8], header::RawElfIdentification> {
        move |b: &[u8]| {
            let (b, e_magic) = nom::bytes::complete::tag(header::ELF_MAGICNUMBER_SIGNATURE)(b)?;

            let mut e_ident = header::RawElfIdentification::default();
            let magicnumber_indices = [
                header::ELF_IDENT_MAGICNUMBER0_INDEX,
                header::ELF_IDENT_MAGICNUMBER1_INDEX,
                header::ELF_IDENT_MAGICNUMBER2_INDEX,
                header::ELF_IDENT_MAGICNUMBER3_INDEX,
            ];

            for idx in magicnumber_indices.iter() {
                e_ident.e_magic[*idx] = e_magic[*idx];
            }

            let (b, e_class) = nom::number::complete::u8(b)?;
            e_ident.e_class = e_class;
            let (b, e_data) = nom::number::complete::u8(b)?;
            e_ident.e_data = e_data;
            let (b, e_version) = nom::number::complete::u8(b)?;
            e_ident.e_version = e_version;

            let (b, e_padding) = nom::multi::count(
                nom::number::complete::u8,
                header::ELF_IDENT_PADDING_FIELD_LENGTH,
            )(b)?;

            e_ident.e_padding[..header::ELF_IDENT_PADDING_FIELD_LENGTH]
                .copy_from_slice(&e_padding[..header::ELF_IDENT_PADDING_FIELD_LENGTH]);

            Ok((b, e_ident))
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::parser::ElfParserConfig;

    use super::*;

    #[test]
    fn test_parse_elf64_header() {
        let p = ElfParserConfig::new().build();
        let input = vec![
            0x7f, 0x45, 0x4c, 0x46, // elf magic number
            0x02, // elf class
            0x01, // elf data
            0x01, // elf version
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // padding
            0x03, 0x00, // elf type
            0x3e, 0x00, // machine
            0x01, 0x00, 0x00, 0x00, // object file version
            0xb0, 0x6a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // entry
            0x40, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // phoff
            0x20, 0x14, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, // shoff
            0x00, 0x00, 0x00, 0x00, // flags
            0x40, 0x00, // ehsize
            0x38, 0x00, // phentsize
            0x0d, 0x00, // phnum
            0x40, 0x00, // shentsize
            0x1f, 0x00, // shnum
            0x1e, 0x00, // shstrndx
        ];
        let result = p.parse_elf64_header(&input);
        assert!(result.is_ok());

        let (b, ehdr) = result.unwrap();
        assert!(b.is_empty());

        assert_eq!(header::ElfClass::Class64, ehdr.ident.class);
        assert_eq!(header::ElfType::Dyn, ehdr.elf_type);
        assert_eq!(header::ElfMachine::X86_64, ehdr.machine);
        assert_eq!(0x1, ehdr.version);
        assert_eq!(0x6ab0, ehdr.entry);
        assert_eq!(0x40, ehdr.phoff);
        assert_eq!(0x21420, ehdr.shoff);
        assert_eq!(0, ehdr.flags);
        assert_eq!(0x40, ehdr.ehsize);
        assert_eq!(0x38, ehdr.phentsize);
        assert_eq!(0xd, ehdr.phnum);
        assert_eq!(0x40, ehdr.shentsize);
        assert_eq!(0x1f, ehdr.shnum);
        assert_eq!(0x1e, ehdr.shstrndx);
    }

    #[test]
    fn test_parse_raw_elf64_header() {
        let p = ElfParserConfig::new().build();
        let input = vec![
            0x7f, 0x45, 0x4c, 0x46, // elf magic number
            0x02, // elf class
            0x01, // elf data
            0x01, // elf version
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // padding
            0x03, 0x00, // elf type
            0x3e, 0x00, // machine
            0x01, 0x00, 0x00, 0x00, // object file version
            0xb0, 0x6a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // entry
            0x40, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // phoff
            0x20, 0x14, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, // shoff
            0x00, 0x00, 0x00, 0x00, // flags
            0x40, 0x00, // ehsize
            0x38, 0x00, // phentsize
            0x0d, 0x00, // phnum
            0x40, 0x00, // shentsize
            0x1f, 0x00, // shnum
            0x1e, 0x00, // shstrndx
        ];
        let result = p.parse_raw_elf64_header(&input);
        assert!(result.is_ok());

        let (b, ehdr) = result.unwrap();
        assert!(b.is_empty());

        assert_eq!(header::ELF_MAGICNUMBER_SIGNATURE, ehdr.e_ident.e_magic);
        assert_eq!(0x2, ehdr.e_ident.e_class);
        assert_eq!(0x1, ehdr.e_ident.e_data);
        assert_eq!(0x1, ehdr.e_ident.e_version);
        assert_eq!(
            [0x00; header::ELF_IDENT_PADDING_FIELD_LENGTH],
            ehdr.e_ident.e_padding
        );
        assert_eq!(0x3, ehdr.e_type);
        assert_eq!(0x3e, ehdr.e_machine);
        assert_eq!(0x1, ehdr.e_version);
        assert_eq!(0x6ab0, ehdr.e_entry);
        assert_eq!(0x40, ehdr.e_phoff);
        assert_eq!(0x21420, ehdr.e_shoff);
        assert_eq!(0, ehdr.e_flags);
        assert_eq!(0x40, ehdr.e_ehsize);
        assert_eq!(0x38, ehdr.e_phentsize);
        assert_eq!(0xd, ehdr.e_phnum);
        assert_eq!(0x40, ehdr.e_shentsize);
        assert_eq!(0x1f, ehdr.e_shnum);
        assert_eq!(0x1e, ehdr.e_shstrndx);
    }

    #[test]
    fn test_peek_raw_elf_identification() {
        let p = ElfParserConfig::new().build();
        let input = vec![
            0x7f, 0x45, 0x4c, 0x46, // elf magic number
            0x02, 0x01, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        ];
        let result = p.peek_raw_elf_identification(&input);
        assert!(result.is_ok());

        let (b, _e_ident) = result.unwrap();
        assert_eq!(&input, b);
    }

    #[test]
    fn test_parse_raw_elf_identification() {
        let p = ElfParserConfig::new().build();
        let input = vec![
            0x7f, 0x45, 0x4c, 0x46, // elf magic number
            0x02, 0x01, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        ];
        let result = p.parse_raw_elf_identification()(&input);
        assert!(result.is_ok());

        let (b, e_ident) = result.unwrap();
        assert_eq!(0, b.len());

        assert_eq!(header::ELF_MAGICNUMBER_SIGNATURE, e_ident.e_magic);
        assert_eq!(0x2, e_ident.e_class);
        assert_eq!(0x1, e_ident.e_data);
        assert_eq!(0x1, e_ident.e_version);
        assert_eq!(
            [0x00; header::ELF_IDENT_PADDING_FIELD_LENGTH],
            e_ident.e_padding
        );
    }
}
