#![no_std]

#[cfg(test)]
mod tests;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ElfError {
    TooSmall,
    InvalidMagic,
    Not64Bit,
    NotLittleEndian,
    NotExecutable,
    NoProgramHeaders,
    HeaderOutOfBounds,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Segment {
    pub vaddr: usize,
    pub file_offset: usize,
    pub file_size: usize,
    pub mem_size: usize,
    pub flags: SegmentFlags,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SegmentFlags {
    pub read: bool,
    pub write: bool,
    pub execute: bool,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ElfInfo {
    pub entry_point: usize,
    pub segments: [Option<Segment>; MAX_SEGMENTS],
    pub segment_count: usize,
}

pub const MAX_SEGMENTS: usize = 8;

const ELF_MAGIC: [u8; 4] = [0x7f, b'E', b'L', b'F'];
const ELFCLASS64: u8 = 2;
const ELFDATA2LSB: u8 = 1;
const ET_EXEC: u16 = 2;
const PT_LOAD: u32 = 1;

const ELF_HEADER_SIZE: usize = 64;
const PHDR_SIZE: usize = 56;

pub fn parse(data: &[u8]) -> Result<ElfInfo, ElfError> {
    if data.len() < ELF_HEADER_SIZE {
        return Err(ElfError::TooSmall);
    }

    if data[0..4] != ELF_MAGIC {
        return Err(ElfError::InvalidMagic);
    }

    if data[4] != ELFCLASS64 {
        return Err(ElfError::Not64Bit);
    }

    if data[5] != ELFDATA2LSB {
        return Err(ElfError::NotLittleEndian);
    }

    let e_type = read_u16(data, 16);
    if e_type != ET_EXEC {
        return Err(ElfError::NotExecutable);
    }

    let entry_point = read_u64(data, 24) as usize;
    let ph_offset = read_u64(data, 32) as usize;
    let ph_entry_size = read_u16(data, 54) as usize;
    let ph_count = read_u16(data, 56) as usize;

    if ph_count == 0 {
        return Err(ElfError::NoProgramHeaders);
    }

    let mut segments = [None; MAX_SEGMENTS];
    let mut segment_count = 0;

    for i in 0..ph_count {
        let offset = ph_offset + i * ph_entry_size;
        if offset + PHDR_SIZE > data.len() {
            return Err(ElfError::HeaderOutOfBounds);
        }

        let p_type = read_u32(data, offset);
        if p_type != PT_LOAD {
            continue;
        }

        if segment_count >= MAX_SEGMENTS {
            break;
        }

        let p_flags = read_u32(data, offset + 4);
        let p_offset = read_u64(data, offset + 8) as usize;
        let p_vaddr = read_u64(data, offset + 16) as usize;
        let p_filesz = read_u64(data, offset + 32) as usize;
        let p_memsz = read_u64(data, offset + 40) as usize;

        segments[segment_count] = Some(Segment {
            vaddr: p_vaddr,
            file_offset: p_offset,
            file_size: p_filesz,
            mem_size: p_memsz,
            flags: SegmentFlags {
                execute: p_flags & 1 != 0,
                write: p_flags & 2 != 0,
                read: p_flags & 4 != 0,
            },
        });
        segment_count += 1;
    }

    Ok(ElfInfo {
        entry_point,
        segments,
        segment_count,
    })
}

impl ElfInfo {
    pub fn segments(&self) -> impl Iterator<Item = &Segment> {
        self.segments[..self.segment_count]
            .iter()
            .filter_map(|s| s.as_ref())
    }

    pub fn total_pages(&self) -> usize {
        let mut total = 0;
        for seg in self.segments() {
            total += (seg.mem_size + 4095) / 4096;
        }
        total
    }
}

fn read_u16(data: &[u8], offset: usize) -> u16 {
    u16::from_le_bytes([data[offset], data[offset + 1]])
}

fn read_u32(data: &[u8], offset: usize) -> u32 {
    u32::from_le_bytes([
        data[offset],
        data[offset + 1],
        data[offset + 2],
        data[offset + 3],
    ])
}

fn read_u64(data: &[u8], offset: usize) -> u64 {
    u64::from_le_bytes([
        data[offset],
        data[offset + 1],
        data[offset + 2],
        data[offset + 3],
        data[offset + 4],
        data[offset + 5],
        data[offset + 6],
        data[offset + 7],
    ])
}