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

use anyhow::Result;
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,
    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 Display for Architecture {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            Architecture::Alpha => write!(f, "DEC Alpha"),
            Architecture::ARM => write!(f, "ARM"),
            Architecture::ARMThumb => write!(f, "ARM Thumb"),
            Architecture::ARM64 => write!(f, "ARM64"),
            Architecture::HitachiSH3 => write!(f, "Hitachi SH3"),
            Architecture::HitachiSH4 => write!(f, "Hitachi SH4"),
            Architecture::HitachiSH5 => write!(f, "Hitachi SH5"),
            Architecture::Itanium => write!(f, "Intel Itanium"),
            Architecture::LoongArch32 => write!(f, "LoongArch"),
            Architecture::LoongArch64 => write!(f, "LoongArch64"),
            Architecture::M68k => write!(f, "M68k"),
            Architecture::M88k => write!(f, "M88k"),
            Architecture::MIPS => write!(f, "MIPS"),
            Architecture::MIPS64 => write!(f, "MIPS64"),
            Architecture::MIPSEL => write!(f, "MIPSEL"),
            Architecture::MIPSEL64 => write!(f, "MIPSEL64"),
            Architecture::PowerPC => write!(f, "PowerPC"),
            Architecture::PowerPC64 => write!(f, "PowerPC64"),
            Architecture::PowerPCLE => write!(f, "PowerPCLE"),
            Architecture::PowerPC64LE => write!(f, "PowerPC64LE"),
            Architecture::RISCV => write!(f, "RISC-V"),
            Architecture::RISCV64 => write!(f, "RISC-V 64"),
            Architecture::RISCV128 => write!(f, "RISC-V 128"),
            Architecture::Sparc => write!(f, "Sparc"),
            Architecture::Sparc64 => write!(f, "Sparc64"),
            Architecture::S390 => write!(f, "S390"),
            Architecture::S390x => write!(f, "S390x"),
            Architecture::X86 => write!(f, "x86"),
            Architecture::X86_64 => write!(f, "x86_64"),
            Architecture::Other(other) => write!(f, "Other: 0x{other:02X}"),
            Architecture::Unknown => write!(f, "Unknown architecture, or architecture-independent"),
        }
    }
}

#[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 Display for OperatingSystem {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            OperatingSystem::AIX => write!(f, "AIX"),
            OperatingSystem::Linux => write!(f, "Linux"),
            OperatingSystem::FreeBSD => write!(f, "FreeBSD"),
            OperatingSystem::OpenBSD => write!(f, "OpenBSD"),
            OperatingSystem::NetBSD => write!(f, "NetBSD"),
            OperatingSystem::HPUX => write!(f, "HP-UX"),
            OperatingSystem::Irix => write!(f, "Irix"),
            OperatingSystem::Solaris => write!(f, "Solaris"),
            OperatingSystem::UnknownUnixLike => write!(f, "Unknown Unix or Unix-like"),
            OperatingSystem::Haiku => write!(f, "Haiku"),
            OperatingSystem::MacOS => write!(f, "Mac OS (or maybe iOS)"),
            OperatingSystem::MacOS_Classic => write!(f, "Classic Mac OS (7.0 - 9.2)"),
            OperatingSystem::DOS => write!(f, "MS-DOS or compatible"),
            OperatingSystem::Windows => write!(f, "Windows"),
            OperatingSystem::Other(other) => write!(f, "Other: 0x{other:02X}"),
        }
    }
}

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum Ordering {
    /// Big Endian, Most Significant Byte (MSB) is first
    BigEndian,

    /// Little Endian, Least Significant Byte (LSB) is first
    LittleEndian,

    /// An application which may use both in the same file
    BiEndian,
}

impl Display for Ordering {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            Ordering::BigEndian => write!(f, "Big Endian"),
            Ordering::LittleEndian => write!(f, "Little Endian"),
            Ordering::BiEndian => write!(f, "Bi-Endian"),
        }
    }
}

#[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 Display for ExecutableType {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            ExecutableType::Core => write!(f, "Core file"),
            ExecutableType::Library => write!(f, "Shared library"),
            ExecutableType::Program => write!(f, "Program/Application"),
            ExecutableType::Unknown(other) => write!(f, "Unknown 0x{other:02X}"),
        }
    }
}

pub trait ExecutableFile {
    fn type_name(&self) -> String;

    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) -> Result<&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,

    pub data: Option<&'a [u8]>,
}

#[derive(Clone, Debug, Default, PartialEq)]
pub struct 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(())
    }
}

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

impl<'a> Deref for Sections<'a> {
    type Target = Sections<'a>;

    fn deref(&self) -> &Self::Target {
        // Seems silly, but seems required for making an Option<Sections> into Result<&Sections>
        self
    }
}

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