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

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

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

#[cfg(feature = "macho")]
pub mod macho;

#[cfg(feature = "pe32")]
pub mod pe32;

#[cfg(feature = "pef")]
pub mod pef;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Architecture {
    Alpha,
    ARM,
    ARMThumb,
    ARM64,
    HitachiSH3,
    HitachiSH4,
    HitachiSH5,
    Hobbit,
    Itanium,
    LoongArch32,
    LoongArch64,
    M68k,
    M88k,
    MIPS,
    MIPS64,
    MIPSEL,
    MIPSEL64,
    PowerPC,
    PowerPC64,
    PowerPCLE,
    PowerPC64LE,
    RISCV,
    RISCV64,
    RISCV128,
    Sparc,
    Sparc64,
    S390,
    S390x,
    X86,
    X86_64,
    Other(u16),
    Unknown,
}

impl Architecture {
    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()),
        }
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
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 {
    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()),
        }
    }
}

#[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 {
    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()),
        }
    }
}

pub trait ExecutableFile {
    fn architecture(&self) -> &Architecture;
    fn pointer_size(&self) -> usize;
    fn operating_system(&self) -> &OperatingSystem;
    fn compiled_timestamp(&self) -> Option<DateTime<Utc>>;

    fn num_sections(&self) -> u32;
    fn sections(&self) -> Option<&Sections>;
}

#[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 or not 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}",
            self.name, self.offset, self.size
        )?;
        if self.virtual_address > 0 {
            write!(f, ", v address: 0x{:02x}", self.virtual_address)?;
        }
        if self.is_executable {
            write!(f, " - executable")?;
        }
        Ok(())
    }
}

#[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)
    }
}

#[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()
    }
}