//! Verify every packaged native library, including JNA, against the ABI/page contract.

use std::{collections::BTreeMap, error::Error, fs::File, io::Read, path::Path};

use object::{Architecture, BinaryFormat, Object, ObjectSection, ObjectSegment, ObjectSymbol};

const PAGE_SIZE: u64 = 16_384;
const ABIS: [&str; 2] = ["arm64-v8a", "armeabi-v7a"];

pub fn package_directory(input: &Path, output: &Path) -> Result<(), Box<dyn Error>> {
    audit_directory(input)?;
    let input = input.canonicalize()?;
    let output = if output.exists() {
        if output.is_symlink() {
            return Err("native output must not be a symlink".into());
        }
        output.canonicalize()?
    } else {
        output
            .parent()
            .ok_or("output needs a parent")?
            .canonicalize()?
            .join(output.file_name().ok_or("output needs a name")?)
    };
    if input.starts_with(&output) || output.starts_with(&input) {
        return Err("native input/output must not overlap".into());
    }
    if output.exists() {
        std::fs::remove_dir_all(&output)?;
    }
    for abi in ABIS {
        std::fs::create_dir_all(output.join(abi))?;
        // cargo-ndk also emits dependency cdylibs; ship only the Android boundary.
        std::fs::copy(
            input.join(abi).join("libironstorage_android.so"),
            output.join(abi).join("libironstorage_android.so"),
        )?;
    }
    Ok(())
}

pub fn audit_directory(directory: &Path) -> Result<(), Box<dyn Error>> {
    let mut expected = None;
    for abi in ABIS {
        let bytes = std::fs::read(directory.join(abi).join("libironstorage_android.so"))?;
        audit_elf(&bytes, abi)?;
        let library = object::File::parse(bytes.as_slice())?;
        let mut metadata = BTreeMap::new();
        for symbol in library.symbols() {
            let name = symbol.name()?;
            if name.starts_with("UNIFFI_META_") {
                let section = library.section_by_index(
                    symbol
                        .section_index()
                        .ok_or("UniFFI metadata has no section")?,
                )?;
                let value = section
                    .data_range(symbol.address(), symbol.size())?
                    .ok_or("UniFFI metadata is outside its section")?;
                metadata.insert(name.to_owned(), value.to_vec());
            }
        }
        if metadata.is_empty() {
            return Err("native library has no UniFFI metadata".into());
        }
        if let Some(expected) = &expected {
            if expected != &metadata {
                return Err("UniFFI metadata differs across Android ABIs".into());
            }
        } else {
            expected = Some(metadata);
        }
    }
    println!("Verified both ARM native ABIs, 16 KB ELF alignment and identical UniFFI metadata");
    Ok(())
}

pub fn audit_elf(bytes: &[u8], abi: &str) -> Result<(), Box<dyn Error>> {
    let expected = match abi {
        "arm64-v8a" => Architecture::Aarch64,
        "armeabi-v7a" => Architecture::Arm,
        _ => return Err("unapproved native ABI".into()),
    };
    let library = object::File::parse(bytes)?;
    if library.format() != BinaryFormat::Elf || library.architecture() != expected {
        return Err("native library architecture does not match its ABI directory".into());
    }
    let mut segments = 0;
    for segment in library.segments() {
        let (offset, _) = segment.file_range();
        if segment.align() < PAGE_SIZE || segment.address() % PAGE_SIZE != offset % PAGE_SIZE {
            return Err("native ELF LOAD segment is not compatible with 16 KB pages".into());
        }
        segments += 1;
    }
    if segments == 0 {
        return Err("native ELF has no loadable segments".into());
    }
    Ok(())
}

pub fn audit_package(path: &Path) -> Result<(), Box<dyn Error>> {
    audit_package_abis(path, &ABIS)
}

pub fn audit_split(path: &Path, abi: &str) -> Result<(), Box<dyn Error>> {
    if !ABIS.contains(&abi) || !path.extension().is_some_and(|ext| ext == "apk") {
        return Err("expected an APK split for an approved ABI".into());
    }
    audit_package_abis(path, &[abi])
}

pub fn audit_release(package: &Path, symbols: &Path) -> Result<(), Box<dyn Error>> {
    audit_package(package)?;
    let mut package = zip::ZipArchive::new(File::open(package)?)?;
    let mut symbols = zip::ZipArchive::new(File::open(symbols)?)?;
    for abi in ABIS {
        for library in ["libironstorage_android.so", "libjnidispatch.so"] {
            let mut bytes = Vec::new();
            package
                .by_name(&format!("lib/{abi}/{library}"))?
                .read_to_end(&mut bytes)?;
            let mut debug = Vec::new();
            let mut entry = symbols.by_name(&format!("{abi}/{library}.dbg"))?;
            if entry.size() > 512 * 1024 * 1024 {
                return Err("native debug symbols exceed the audit bound".into());
            }
            entry.read_to_end(&mut debug)?;
            audit_symbol_pair(&bytes, &debug, abi, library == "libironstorage_android.so")?;
        }
    }
    println!("Release libraries are stripped; FULL symbols match both ARM build IDs");
    Ok(())
}

fn audit_symbol_pair(
    bytes: &[u8],
    debug: &[u8],
    abi: &str,
    require_debug_info: bool,
) -> Result<(), Box<dyn Error>> {
    audit_elf(bytes, abi)?;
    let release = object::File::parse(bytes)?;
    let symbols = object::File::parse(debug)?;
    if symbols.format() != BinaryFormat::Elf || symbols.architecture() != release.architecture() {
        return Err("debug archive architecture does not match release".into());
    }
    if release.sections().any(|s| {
        s.name()
            .is_ok_and(|n| n.starts_with(".debug_") || n == ".symtab")
    }) {
        return Err("release library contains unstripped debug symbols".into());
    }
    if (require_debug_info && symbols.section_by_name(".debug_info").is_none())
        || symbols.symbols().next().is_none()
    {
        return Err("native symbol archive lacks required symbols/debug information".into());
    }
    let id = release
        .build_id()?
        .ok_or("release library lacks a build ID")?;
    if id.is_empty() || Some(id) != symbols.build_id()? {
        return Err("native symbol archive does not match release build ID".into());
    }
    Ok(())
}

fn audit_package_abis(path: &Path, abis: &[&str]) -> Result<(), Box<dyn Error>> {
    let apk = path.extension().is_some_and(|ext| ext == "apk");
    if !apk && !path.extension().is_some_and(|ext| ext == "aab") {
        return Err("expected an APK or AAB".into());
    }
    let mut archive = zip::ZipArchive::new(File::open(path)?)?;
    let mut libraries = BTreeMap::<String, Vec<String>>::new();
    for index in 0..archive.len() {
        let mut file = archive.by_index(index)?;
        let name = file.name().to_owned();
        let native = name.strip_prefix(if apk { "lib/" } else { "base/lib/" });
        let Some((abi, library)) = native.and_then(|name| name.split_once('/')) else {
            continue;
        };
        if !library.ends_with(".so") {
            continue;
        }
        if library == "libironstorage_apple.so" {
            return Err("Android package includes the redundant Apple boundary library".into());
        }
        if !abis.contains(&abi) || library.contains('/') {
            return Err("package contains an unapproved native library path".into());
        }
        if ![
            "libironstorage_android.so",
            "libjnidispatch.so",
            "libimage_processing_util_jni.so",
            "libsurface_util_jni.so",
            "libandroidx.graphics.path.so",
        ]
        .contains(&library)
            || libraries
                .get(abi)
                .is_some_and(|names| names.iter().any(|name| name == library))
        {
            return Err("package contains an unreviewed or duplicate native library".into());
        }
        if apk
            && (file.compression() != zip::CompressionMethod::Stored
                || file.data_start() % PAGE_SIZE != 0)
        {
            return Err(format!("{name}: APK native entry is not stored/aligned to 16 KB").into());
        }
        if file.size() > 512 * 1024 * 1024 {
            return Err("native library exceeds the package audit bound".into());
        }
        let mut bytes = Vec::new();
        file.read_to_end(&mut bytes)?;
        audit_elf(&bytes, abi).map_err(|error| format!("{name}: {error}"))?;
        libraries
            .entry(abi.to_owned())
            .or_default()
            .push(library.to_owned());
    }
    for abi in abis {
        let found = libraries
            .get(*abi)
            .ok_or("package is missing a supported ABI")?;
        for required in ["libironstorage_android.so", "libjnidispatch.so"] {
            if !found.iter().any(|name| name == required) {
                return Err(format!("{abi}: missing {required}").into());
            }
        }
    }
    println!(
        "Verified all packaged native libraries for {abis:?} in {}",
        path.display()
    );
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    fn elf(alignment: u64, address: u64) -> Vec<u8> {
        let mut bytes = vec![0; 120];
        bytes[..7].copy_from_slice(b"\x7fELF\x02\x01\x01");
        bytes[16..18].copy_from_slice(&3_u16.to_le_bytes());
        bytes[18..20].copy_from_slice(&183_u16.to_le_bytes());
        bytes[20..24].copy_from_slice(&1_u32.to_le_bytes());
        bytes[32..40].copy_from_slice(&64_u64.to_le_bytes());
        bytes[52..54].copy_from_slice(&64_u16.to_le_bytes());
        bytes[54..56].copy_from_slice(&56_u16.to_le_bytes());
        bytes[56..58].copy_from_slice(&1_u16.to_le_bytes());
        bytes[64..68].copy_from_slice(&1_u32.to_le_bytes());
        bytes[68..72].copy_from_slice(&4_u32.to_le_bytes());
        bytes[80..88].copy_from_slice(&address.to_le_bytes());
        bytes[96..104].copy_from_slice(&120_u64.to_le_bytes());
        bytes[104..112].copy_from_slice(&120_u64.to_le_bytes());
        bytes[112..120].copy_from_slice(&alignment.to_le_bytes());
        bytes
    }

    #[test]
    fn checks_real_load_alignment_congruence_and_directory_abi() {
        assert!(audit_elf(&elf(PAGE_SIZE, 0), "arm64-v8a").is_ok());
        assert!(audit_elf(&elf(4096, 0), "arm64-v8a").is_err());
        assert!(audit_elf(&elf(PAGE_SIZE, 4), "arm64-v8a").is_err());
        assert!(audit_elf(&elf(PAGE_SIZE, 0), "x86_64").is_err());
    }

    #[test]
    fn rejects_non_elf_unapproved_abi_and_truncated_headers() {
        assert!(audit_elf(b"not an ELF", "arm64-v8a").is_err());
        assert!(audit_elf(b"\x7fELF", "x86_64").is_err());
        assert!(audit_elf(&[], "mips").is_err());
    }

    #[test]
    fn symbol_gate_rejects_missing_full_symbols_and_build_identity() {
        assert!(
            audit_symbol_pair(&elf(PAGE_SIZE, 0), &elf(PAGE_SIZE, 0), "arm64-v8a", true).is_err()
        );
        assert!(
            audit_symbol_pair(&elf(PAGE_SIZE, 0), b"wrong archive", "arm64-v8a", false).is_err()
        );
    }
}