use crate::static_vec::StaticVec;
use zerocopy::{FromBytes, Immutable, KnownLayout};

const ELF_MAGIC: [u8; 4] = [0x7f, b'E', b'L', b'F'];
const PT_LOAD: u32 = 1;
const MAX_LOAD_SEGMENTS: usize = 8;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ElfError {
    Truncated,
    BadMagic,
    UnsupportedFormat,
    InvalidPhdrSize,
    TooManySegments,
    MalformedSegment,
}

fn read_struct<T: FromBytes + KnownLayout + Immutable>(data: &[u8], offset: usize) -> Option<T> {
    let size = core::mem::size_of::<T>();
    let end = offset.checked_add(size)?;
    T::read_from_bytes(data.get(offset..end)?).ok()
}

#[repr(C)]
#[derive(Clone, Copy, FromBytes, KnownLayout, Immutable)]
pub struct Elf64Header {
    pub e_ident: [u8; 16],
    pub e_type: u16,
    pub e_machine: u16,
    pub e_version: u32,
    pub e_entry: u64,
    pub e_phoff: u64,
    pub e_shoff: u64,
    pub e_flags: u32,
    pub e_ehsize: u16,
    pub e_phentsize: u16,
    pub e_phnum: u16,
    pub e_shentsize: u16,
    pub e_shnum: u16,
    pub e_shstrndx: u16,
}

#[repr(C)]
#[derive(Clone, Copy, FromBytes, KnownLayout, Immutable)]
pub struct Elf64Phdr {
    pub p_type: u32,
    pub p_flags: u32,
    pub p_offset: u64,
    pub p_vaddr: u64,
    pub p_paddr: u64,
    pub p_filesz: u64,
    pub p_memsz: u64,
    pub p_align: u64,
}

#[derive(Debug)]
pub struct ElfInfo {
    pub entry: u64,
    pub segments: StaticVec<LoadSegment, MAX_LOAD_SEGMENTS>,
}

#[derive(Debug, Clone, Copy)]
pub struct LoadSegment {
    pub vaddr: u64,
    pub file_offset: u64,
    pub filesz: u64,
    pub memsz: u64,
    pub writable: bool,
    pub executable: bool,
}

pub fn parse(data: &[u8]) -> Result<ElfInfo, ElfError> {
    let header: Elf64Header = read_struct(data, 0).ok_or(ElfError::Truncated)?;

    if header.e_ident[0..4] != ELF_MAGIC {
        return Err(ElfError::BadMagic);
    }

    if header.e_ident[4] != 2 || header.e_ident[5] != 1 || header.e_machine != 0x3E {
        return Err(ElfError::UnsupportedFormat);
    }

    if (header.e_phentsize as usize) < core::mem::size_of::<Elf64Phdr>() {
        return Err(ElfError::InvalidPhdrSize);
    }

    let mut segments: StaticVec<LoadSegment, MAX_LOAD_SEGMENTS> = StaticVec::new();

    let phdr_base = header.e_phoff as usize;
    let phdr_size = header.e_phentsize as usize;

    (0..header.e_phnum as usize)
        .filter_map(|i| {
            let offset = phdr_base.checked_add(i.checked_mul(phdr_size)?)?;
            let phdr: Elf64Phdr = read_struct(data, offset)?;
            let valid = phdr.p_type == PT_LOAD
                && phdr.p_memsz > 0
                && phdr.p_memsz >= phdr.p_filesz
                && phdr
                    .p_offset
                    .checked_add(phdr.p_filesz)
                    .is_some_and(|end| end <= data.len() as u64);
            if valid {
                Some(LoadSegment {
                    vaddr: phdr.p_vaddr,
                    file_offset: phdr.p_offset,
                    filesz: phdr.p_filesz,
                    memsz: phdr.p_memsz,
                    writable: (phdr.p_flags & 2) != 0,
                    executable: (phdr.p_flags & 1) != 0,
                })
            } else {
                None
            }
        })
        .try_fold(&mut segments, |segs, seg| match segs.push(seg) {
            Ok(()) => Ok(segs),
            Err(_) => Err(ElfError::TooManySegments),
        })?;

    Ok(ElfInfo {
        entry: header.e_entry,
        segments,
    })
}

#[cfg(kani)]
mod kani_proofs {
    use super::*;

    #[kani::proof]
    fn parse_rejects_empty() {
        assert!(parse(&[]).is_err());
    }

    #[kani::proof]
    fn read_struct_rejects_truncated_header() {
        let len: usize = kani::any();
        kani::assume(len < core::mem::size_of::<Elf64Header>());
        let data = [0u8; 63];
        let result: Option<Elf64Header> = read_struct(&data[..len], 0);
        assert!(result.is_none());
    }

    #[kani::proof]
    fn read_struct_bounds_check() {
        let data: [u8; 16] = kani::any();
        let offset: usize = kani::any();
        kani::assume(offset <= 20);
        let result: Option<u64> = read_struct(&data, offset);
        if offset.checked_add(8).is_none_or(|end| end > 16) {
            assert!(result.is_none());
        }
    }

    #[kani::proof]
    #[kani::unwind(10)]
    fn phdr_offset_arithmetic_no_overflow() {
        let base: usize = kani::any();
        let entry_size: usize = kani::any();
        let count: usize = kani::any();
        kani::assume(count <= 8);
        kani::assume(entry_size >= core::mem::size_of::<Elf64Phdr>());
        kani::assume(entry_size <= 256);

        (0..count).for_each(|i| {
            let product = i.checked_mul(entry_size);
            let offset = product.and_then(|p| base.checked_add(p));
            assert!(product.is_none() || offset.is_some() || base > usize::MAX - product.unwrap());
        });
    }

    #[kani::proof]
    fn segment_file_bounds_check() {
        let data_len: u64 = kani::any();
        let p_offset: u64 = kani::any();
        let p_filesz: u64 = kani::any();

        kani::assume(data_len <= 0x1_0000);

        let in_bounds = p_offset
            .checked_add(p_filesz)
            .is_some_and(|end| end <= data_len);

        if in_bounds {
            assert!(p_offset <= data_len);
            assert!(p_filesz <= data_len);
        }
    }

    #[kani::proof]
    fn segment_memsz_ge_filesz_check() {
        let p_memsz: u64 = kani::any();
        let p_filesz: u64 = kani::any();

        let valid = p_memsz > 0 && p_memsz >= p_filesz;
        if valid {
            assert!(p_memsz >= p_filesz);
        }
    }
}