use std::collections::{HashMap, HashSet};
use std::fmt::{Display, Formatter};

use chrono::{DateTime, Utc};
use fuzzyhash::FuzzyHash;
use md5::{Digest, Md5};

/// ELF parsing
#[cfg(feature = "elf")]
pub mod elf;

/// Mach-O parsing
#[cfg(feature = "macho")]
pub mod macho;

/// Parsing of PE32 & EXE files, and anything else with an MZ header
#[cfg(feature = "pe32")]
pub mod pe32;

/// PEF (Preferred Executable File) parsing
#[cfg(feature = "pef")]
pub mod pef;

/// CPU Architectures
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
pub enum Architecture {
    /// DEC Alpha
    /// <https://en.wikipedia.org/wiki/DEC_Alpha>
    Alpha,

    /// ARM 32-bit
    /// <https://en.wikipedia.org/wiki/ARM_architecture_family>
    ARM,

    /// ARM Thumb
    /// <https://en.wikipedia.org/wiki/ARM_architecture_family>
    ARMThumb,

    /// ARM 64-bit, also known as Aarch64
    /// <https://en.wikipedia.org/wiki/ARM_architecture_family>
    ARM64,

    /// Hitachi SH3
    /// <https://en.wikipedia.org/wiki/SuperH>
    HitachiSH3,

    /// Hitachi SH4
    /// <https://en.wikipedia.org/wiki/SuperH>
    HitachiSH4,

    /// Hitachi SH5
    /// <https://en.wikipedia.org/wiki/SuperH>
    HitachiSH5,

    /// AT&T Hobbit
    /// <https://en.wikipedia.org/wiki/AT%26T_Hobbit>
    Hobbit,

    /// Intel Itanium
    /// <https://en.wikipedia.org/wiki/Itanium>
    Itanium,

    /// Loongson 32-bit
    /// <https://en.wikipedia.org/wiki/Loongson>
    LoongArch32,

    /// Loongson 64-bit
    /// <https://en.wikipedia.org/wiki/Loongson>
    LoongArch64,

    /// Motorola 68000 (68k)
    /// <https://en.wikipedia.org/wiki/Motorola_68000>
    M68k,

    /// Motorola 88000 (88k)
    /// <https://en.wikipedia.org/wiki/Motorola_88000>
    M88k,

    /// MIPS 32-bit Big Endian
    /// <https://en.wikipedia.org/wiki/MIPS_architecture>
    MIPS,

    /// MIPS 64-bit Big Endian
    /// <https://en.wikipedia.org/wiki/MIPS_architecture>
    MIPS64,

    /// MIPS 32-bit Little Endian
    /// <https://en.wikipedia.org/wiki/MIPS_architecture>
    MIPSEL,

    /// MIPS 64-bit Little Endian
    /// <https://en.wikipedia.org/wiki/MIPS_architecture>
    MIPSEL64,

    /// IBM PowerPC 32-bit Big Endian
    /// <https://en.wikipedia.org/wiki/PowerPC>
    PowerPC,

    /// IBM PowerPC 64-bit Big Endian
    /// <https://en.wikipedia.org/wiki/PowerPC>
    PowerPC64,

    /// IBM PowerPC 32-bit Little Endian
    /// <https://en.wikipedia.org/wiki/PowerPC>
    PowerPCLE,

    /// IBM PowerPC 64-bit Little Endian
    /// <https://en.wikipedia.org/wiki/PowerPC>
    PowerPC64LE,

    /// RISC-V 32-bit
    /// <https://en.wikipedia.org/wiki/RISC-V>
    RISCV,

    /// RISC-V 64-bit
    /// <https://en.wikipedia.org/wiki/RISC-V>
    RISCV64,

    /// RISC-V 128-bit
    /// <https://en.wikipedia.org/wiki/RISC-V>
    RISCV128,

    /// Sun (now Oracle) Sparc 32-bit
    /// <https://en.wikipedia.org/wiki/SPARC>
    Sparc,

    /// Sun (now Oracle) Sparc 64-bit
    /// <https://en.wikipedia.org/wiki/SPARC>
    Sparc64,

    /// IBM s390 mainframe 32-bit
    /// <https://en.wikipedia.org/wiki/IBM_System/390>
    S390,

    /// IBM s390x mainframe 64-bit
    /// <https://en.wikipedia.org/wiki/IBM_System/390>
    S390x,

    /// Intel (or AMD) x86 32-bit
    /// <https://en.wikipedia.org/wiki/X86-64>
    X86,

    /// Intel (or AMD) x86 64-bit
    /// <https://en.wikipedia.org/wiki/X86-64>
    X86_64,

    /// Other CPU type
    Other(u16),

    /// Unknown CPU
    Unknown,
}

impl Architecture {
    /// Static string representation
    pub fn as_str(&self) -> &'static str {
        match self {
            Architecture::Alpha => "DEC Alpha",
            Architecture::ARM => "ARM",
            Architecture::ARMThumb => "ARM Thumb",
            Architecture::ARM64 => "ARM64",
            Architecture::HitachiSH3 => "Hitachi SH3",
            Architecture::HitachiSH4 => "Hitachi SH4",
            Architecture::HitachiSH5 => "Hitachi SH5",
            Architecture::Hobbit => "AT&T Hobbit",
            Architecture::Itanium => "Intel Itanium",
            Architecture::LoongArch32 => "LoongArch",
            Architecture::LoongArch64 => "LoongArch64",
            Architecture::M68k => "M68k",
            Architecture::M88k => "M88k",
            Architecture::MIPS => "MIPS",
            Architecture::MIPS64 => "MIPS64",
            Architecture::MIPSEL => "MIPSEL",
            Architecture::MIPSEL64 => "MIPSEL64",
            Architecture::PowerPC => "PowerPC",
            Architecture::PowerPC64 => "PowerPC64",
            Architecture::PowerPCLE => "PowerPCLE",
            Architecture::PowerPC64LE => "PowerPC64LE",
            Architecture::RISCV => "RISC-V",
            Architecture::RISCV64 => "RISC-V 64",
            Architecture::RISCV128 => "RISC-V 128",
            Architecture::Sparc => "Sparc",
            Architecture::Sparc64 => "Sparc64",
            Architecture::S390 => "S390",
            Architecture::S390x => "S390x",
            Architecture::X86 => "x86",
            Architecture::X86_64 => "x86_64",
            Architecture::Other(_) => "Other",
            Architecture::Unknown => "Unknown architecture, or architecture-independent",
        }
    }
}

impl Display for Architecture {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            Architecture::Other(other) => write!(f, "Other: 0x{other:02X}"),
            a => write!(f, "{}", a.as_str()),
        }
    }
}

/// Operating Systems
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
pub enum OperatingSystem {
    /// IBM AIX
    AIX,

    /// Linux (includes "SystemV" type in ELFs)
    Linux,

    /// FreeBSD
    FreeBSD,

    /// OpenBSD
    OpenBSD,

    /// NetBSD
    NetBSD,

    /// HP's UX
    HPUX,

    /// SGI's Irix
    Irix,

    /// Sun then Oracle Solaris
    Solaris,

    /// Unknown Unix or Unix-like
    UnknownUnixLike,

    /// Haiku, the BeOS successor
    Haiku,

    /// Apple's Mac OS X (now macOS)
    MacOS,

    /// Apple's older Mac OS, now referred to Classic Mac OS
    #[allow(non_camel_case_types)]
    MacOS_Classic,

    /// MS-DOS, IBM-DOS, or FreeDOS
    DOS,

    /// Microsoft Windows
    Windows,

    /// Something else?
    Other(u16),
}

impl OperatingSystem {
    /// Static string representation
    pub fn as_str(&self) -> &'static str {
        match self {
            OperatingSystem::AIX => "AIX",
            OperatingSystem::Linux => "Linux",
            OperatingSystem::FreeBSD => "FreeBSD",
            OperatingSystem::OpenBSD => "OpenBSD",
            OperatingSystem::NetBSD => "NetBSD",
            OperatingSystem::HPUX => "HP-UX",
            OperatingSystem::Irix => "Irix",
            OperatingSystem::Solaris => "Solaris",
            OperatingSystem::UnknownUnixLike => "Unknown Unix or Unix-like",
            OperatingSystem::Haiku => "Haiku",
            OperatingSystem::MacOS => "Mac OS (or maybe iOS)",
            OperatingSystem::MacOS_Classic => "Classic Mac OS (7.0 - 9.2)",
            OperatingSystem::DOS => "MS-DOS or compatible",
            OperatingSystem::Windows => "Windows",
            OperatingSystem::Other(_) => "Other",
        }
    }
}

impl Display for OperatingSystem {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            OperatingSystem::Other(other) => write!(f, "Other: 0x{other:02X}"),
            o => write!(f, "{}", o.as_str()),
        }
    }
}

/// Type of binary file containing machine code
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum ExecutableType {
    /// Core file, from a crash
    Core,

    /// Shared library
    Library,

    /// Directly executable program or application
    Program,

    /// Something else?
    Unknown(u16),
}

impl ExecutableType {
    /// Static string representation
    pub fn as_str(&self) -> &'static str {
        match self {
            ExecutableType::Core => "Core file",
            ExecutableType::Library => "Shared library",
            ExecutableType::Program => "Program/Application",
            ExecutableType::Unknown(_) => "Unknown",
        }
    }
}

impl Display for ExecutableType {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            ExecutableType::Unknown(other) => write!(f, "Unknown 0x{other:02X}"),
            x => write!(f, "{}", x.as_str()),
        }
    }
}

/// Common functions for executable files
pub trait ExecutableFile {
    /// Get the architecture type
    fn architecture(&self) -> Architecture;

    /// Get the pointer size, 32- or 64-bit
    fn pointer_size(&self) -> usize;

    /// Get the operating system type for a binary
    fn operating_system(&self) -> OperatingSystem;

    /// Get the compilation timestamp, if available
    fn compiled_timestamp(&self) -> Option<DateTime<Utc>>;

    /// Number of sections for a binary
    fn num_sections(&self) -> u32;

    /// Vec of sections for the binary
    fn sections(&self) -> Option<&Sections>;

    /// Import hash of the binary
    fn import_hash(&self) -> Option<String>;

    /// SSDeep fuzzy hash of the binary
    fn fuzzy_imports(&self) -> Option<String>;
}

/// Section of an executable file
#[derive(Clone, Debug, PartialEq)]
pub struct Section<'a> {
    /// Name of the section, can be empty, not a reliable way to identify attributes of it
    pub name: String,

    /// Whether an execute bit was set
    pub is_executable: bool,

    /// Size of the section
    pub size: usize,

    /// Offset in the file where the section starts
    pub offset: usize,

    /// Address of the section once loaded into memory, not for all executable types
    pub virtual_address: u32,

    /// Size of the section once loaded into memory, not for all executable types
    pub virtual_size: u32,

    /// Entropy of the section
    pub entropy: f32,

    /// Slice of this section's bytes
    pub data: Option<&'a [u8]>,
}

type Sections<'a> = Vec<Section<'a>>;

impl<'a> Display for Section<'a> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "{} at 0x{:02x}, size 0x{:02x}, entropy {:.2}",
            self.name, self.offset, self.size, self.entropy
        )?;
        if self.virtual_address > 0 {
            write!(f, ", v address: 0x{:02x}", self.virtual_address)?;
        }
        if self.is_executable {
            write!(f, " - executable")?;
        }
        Ok(())
    }
}

/// Import data, normalized across various executable file formats
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Import {
    /// Library file, .dll in Windows, .so in Unix/Linux, .dylib in macOS
    pub library: String,

    /// Function name imported
    pub function: String,
}

impl Display for Import {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}:{}", self.library, self.function)
    }
}

/// Collection of import data, normalized across various executable file formats
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct Imports {
    /// The collection of found imports
    pub imports: Vec<Import>,

    /// The total number of imports which should have been found, in case some couldn't be parsed
    pub expected_imports: u32,
}

impl Display for Imports {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        for import in self.imports.iter() {
            writeln!(f, "{import}")?;
        }
        Ok(())
    }
}

impl Imports {
    /// Build a string with library.function for each pair, sorted.
    pub fn build_import_string(&self) -> String {
        // A HashSet probably isn't needed, but malware might do something funny.
        let mut imports_map: HashMap<String, HashSet<String>> = HashMap::new();

        // Collect all function names by library
        for import in self.imports.iter() {
            let mut lib = import.library.clone().to_lowercase();
            if lib.ends_with(".dll") {
                lib = lib.replace(".dll", "");
            } else if lib.ends_with(".sys") {
                lib = lib.replace(".sys", "");
            } else if let Some(idx) = lib.find(".so") {
                lib.truncate(lib.len() - idx);
            }

            if !imports_map.contains_key(&lib) {
                imports_map.insert(lib.clone(), HashSet::new());
            }

            imports_map
                .get_mut(&lib)
                .unwrap()
                .insert(import.function.to_lowercase());
        }

        // Sort the libraries
        let mut libs: Vec<&String> = imports_map.keys().collect();
        libs.sort();

        // Get the mapping of lib.func
        let mut imports_string = Vec::new();
        for lib in libs {
            // Sort the functions
            let functions = imports_map.get(lib).unwrap();
            let mut functions = Vec::from_iter(functions);
            functions.sort();
            for function in functions.iter() {
                imports_string.push(format!("{lib}.{function}"));
            }
        }

        imports_string.join(",")
    }

    /// The Import Hash, or "ImpHash" is the MD5 of the imports string
    pub fn hash(&self) -> Vec<u8> {
        let mut hasher = Md5::new();
        hasher.update(self.build_import_string());
        let result = hasher.finalize();
        result.to_vec()
    }

    /// The fuzzy import hash is the SSDeep hash of the import string
    pub fn fuzzy_hash(&self) -> String {
        let import_string = self.build_import_string();
        let fuzzy = FuzzyHash::new(import_string.into_bytes());
        fuzzy.to_string()
    }
}