// ---------------------------------------------------------------------------
// load — PE32+ executable loader (ADR 0006 Phase 2).
//
// Parses a PE32+ image read through the VFS, allocates identity-mapped
// pages for it (the kernel runs identity-mapped, so "user" code lives at
// its link address), applies base relocations, and returns the entry
// point.  The spawn path (exec.rs glue in sched) builds an Exports table,
// copies the argv block into the image's memory, and starts a task whose
// trampoline calls entry(&Exports, argc, argv).
//
// Scope: what rustc/lld emits for a no_std UEFI-ish binary — DOS header,
// COFF header, optional header PE32+, section table, .reloc base
// relocations.  No imports, no TLS, no exceptions.
// ---------------------------------------------------------------------------
use crate::paging;

/// Where binaries are loaded: above the heap, inside conventional RAM.
/// The page allocator's span covers the conventional-memory range of the
/// boot memory map; LOAD_BASE must sit inside it (heap ends ~0xbb6c000
/// with 256 MiB guests; conventional RAM typically ends near 0x0FF00000).
const LOAD_BASE: usize = 0x0C00_0000; // 192 MiB
const LOAD_LIMIT: usize = 0x0F00_0000; // stop before top-of-RAM

const PAGE_SIZE: usize = 4096;

/// A loaded image: pages to free on exit + the entry point.
pub struct Image {
    /// First allocated page (free this..+page_count on exit).
    pub base: usize,
    pub page_count: usize,
    pub entry: usize,
}

#[derive(Debug)]
pub enum LoadError {
    ReadFailed,
    NotPe,
    NotPe32Plus,
    BadSectionTable,
    NoEntryPoint,
    AllocFailed,
    MapFailed,
    BadReloc,
    TooBig,
}

// ---------------------------------------------------------------------------
// PE file-format structures (packed reads via byte slices — no unaligned
// struct pointer derefs).
// ---------------------------------------------------------------------------

fn u16_at(b: &[u8], off: usize) -> Option<u16> {
    b.get(off..off + 2)
        .map(|s| u16::from_le_bytes([s[0], s[1]]))
}
fn u32_at(b: &[u8], off: usize) -> Option<u32> {
    b.get(off..off + 4)
        .map(|s| u32::from_le_bytes([s[0], s[1], s[2], s[3]]))
}
fn u64_at(b: &[u8], off: usize) -> Option<u64> {
    b.get(off..off + 8).map(|s| {
        let mut t = [0u8; 8];
        t.copy_from_slice(s);
        u64::from_le_bytes(t)
    })
}

const DOS_MAGIC: u16 = 0x5A4D; // "MZ"
const PE_MAGIC: u32 = 0x0000_4550; // "PE\0\0"
const OPT_PE32PLUS: u16 = 0x20B;
const DIR_BASERELOC: usize = 5;

/// Map an RVA to its raw file offset via the section table.
fn rva_to_file(sections: &[Option<Section>], count: usize, rva: usize) -> Option<usize> {
    for s in sections.iter().take(count).flatten() {
        if rva >= s.virt_addr && rva < s.virt_addr + s.raw_size.max(s.virt_size) {
            return Some(s.raw_ptr + (rva - s.virt_addr));
        }
    }
    None
}

#[derive(Clone, Copy)]
struct Section {
    virt_addr: usize,
    virt_size: usize,
    raw_ptr: usize,
    raw_size: usize,
}

/// Total virtual span of a PE image: the highest section RVA+virt_size
/// (never smaller than one page for the headers).
fn image_span_bytes(img: &[u8]) -> usize {
    let e_lfanew = u32_at(img, 0x3C).unwrap_or(0) as usize;
    let coff = e_lfanew + 4;
    let num_sections = u16_at(img, coff + 2).unwrap_or(0) as usize;
    let opt_size = u16_at(img, coff + 16).unwrap_or(0) as usize;
    let sect_tab = coff + 20 + opt_size;
    let mut end = PAGE_SIZE; // at least one page for headers
    for i in 0..num_sections {
        let s = sect_tab + i * 40;
        let vsz = u32_at(img, s + 8).unwrap_or(0) as usize;
        let va = u32_at(img, s + 12).unwrap_or(0) as usize;
        end = end.max(va + vsz);
    }
    end
}
fn parse_and_relocate(img: &[u8], mem: usize) -> Result<usize, LoadError> {
    // ---- DOS header ----
    if u16_at(img, 0) != Some(DOS_MAGIC) {
        return Err(LoadError::NotPe);
    }
    let e_lfanew = u32_at(img, 0x3C).ok_or(LoadError::NotPe)? as usize;

    // ---- COFF header ----
    if u32_at(img, e_lfanew) != Some(PE_MAGIC) {
        return Err(LoadError::NotPe);
    }
    let coff = e_lfanew + 4;
    let num_sections = u16_at(img, coff + 2).ok_or(LoadError::NotPe)? as usize;
    let opt_size = u16_at(img, coff + 16).ok_or(LoadError::NotPe)? as usize;
    if opt_size == 0 {
        return Err(LoadError::NotPe);
    }

    // ---- Optional header (PE32+) ----
    let opt = coff + 20;
    if u16_at(img, opt) != Some(OPT_PE32PLUS) {
        return Err(LoadError::NotPe32Plus);
    }
    let entry_rva = u32_at(img, opt + 16).ok_or(LoadError::NotPe32Plus)? as usize;
    let image_base = u64_at(img, opt + 24).ok_or(LoadError::NotPe32Plus)? as usize;

    let num_dirs = u32_at(img, opt + 108).ok_or(LoadError::NotPe32Plus)? as usize;
    let dir_base = opt + 112;
    let (reloc_rva, reloc_size) = if num_dirs > DIR_BASERELOC {
        (
            u32_at(img, dir_base + DIR_BASERELOC * 8).ok_or(LoadError::BadReloc)? as usize,
            u32_at(img, dir_base + DIR_BASERELOC * 8 + 4).ok_or(LoadError::BadReloc)? as usize,
        )
    } else {
        (0, 0)
    };

    // ---- Section table ----
    let sect_tab = opt + opt_size;
    let mut sections: [Option<Section>; 96] = [None; 96];
    if num_sections > sections.len() {
        return Err(LoadError::BadSectionTable);
    }
    for i in 0..num_sections {
        let s = sect_tab + i * 40;
        let virt_size = u32_at(img, s + 8).ok_or(LoadError::BadSectionTable)? as usize;
        let virt_addr = u32_at(img, s + 12).ok_or(LoadError::BadSectionTable)? as usize;
        let raw_size = u32_at(img, s + 16).ok_or(LoadError::BadSectionTable)? as usize;
        let raw_ptr = u32_at(img, s + 20).ok_or(LoadError::BadSectionTable)? as usize;
        sections[i] = Some(Section {
            virt_addr,
            virt_size,
            raw_ptr,
            raw_size,
        });
    }

    // ---- Copy section contents ----
    for s in sections.iter_mut().take(num_sections).flatten() {
        let n = s.raw_size.min(s.virt_size);
        if n == 0 {
            continue;
        }
        let src = img
            .get(s.raw_ptr..s.raw_ptr + n)
            .ok_or(LoadError::BadSectionTable)?;
        unsafe {
            let dst = (mem + s.virt_addr) as *mut u8;
            core::ptr::copy_nonoverlapping(src.as_ptr(), dst, n);
        }
    }

    // ---- Apply base relocations ----
    if reloc_size > 0 && reloc_rva != 0 {
        // The .reloc directory is an RVA; the section table maps it to a
        // raw file offset.
        let mut off = rva_to_file(&sections, num_sections, reloc_rva).ok_or(LoadError::BadReloc)?;
        let end = off + reloc_size;
        while off + 8 <= end {
            let page_rva = u32_at(img, off).ok_or(LoadError::BadReloc)? as usize;
            let block_size = u32_at(img, off + 4).ok_or(LoadError::BadReloc)? as usize;
            if block_size < 8 || off + block_size > end {
                break;
            }
            let entries = (block_size - 8) / 2;
            for i in 0..entries {
                let e = u16_at(img, off + 8 + i * 2).ok_or(LoadError::BadReloc)?;
                let r#type = e >> 12;
                let pos = e & 0xFFF;
                if r#type == 0 {
                    continue; // IMAGE_REL_BASED_ABSOLUTE — padding
                }
                let target = mem + page_rva + pos as usize;
                match r#type {
                    10 => {
                        // DIR64: 8-bit image base delta
                        unsafe {
                            let p = target as *mut u64;
                            let v = p.read_unaligned();
                            p.write_unaligned(
                                v.wrapping_add((mem as u64).wrapping_sub(image_base as u64)),
                            );
                        }
                    }
                    _ => return Err(LoadError::BadReloc),
                }
            }
            off += block_size;
        }
    }

    Ok(entry_rva)
}

/// Load a PE32+ image from the given bytes.  Returns the loaded image
/// (identity-mapped at LOAD_BASE+) with its entry point resolved.
pub fn load_image(bytes: &[u8]) -> Result<Image, LoadError> {
    if bytes.len() > 8 * 1024 * 1024 {
        return Err(LoadError::TooBig);
    }
    // Quick header sanity before allocating.
    if u16_at(bytes, 0) != Some(DOS_MAGIC) {
        return Err(LoadError::NotPe);
    }
    let e_lfanew = u32_at(bytes, 0x3C).unwrap_or(0) as usize;
    if u32_at(bytes, e_lfanew) != Some(PE_MAGIC) {
        return Err(LoadError::NotPe);
    }

    // The memory span is the VIRTUAL layout: the highest section RVA +
    // virt_size (headers live at rva 0).  Sections are scattered across
    // RVAs, so this can exceed the file size.
    let image_span = image_span_bytes(bytes);
    let page_count = image_span.div_ceil(PAGE_SIZE);
    // Allocate physical pages at a fixed window (identity mapping makes
    // the link address = runtime address after relocation).
    let base =
        crate::memory::alloc_pages_above(page_count, LOAD_BASE).ok_or(LoadError::AllocFailed)?;
    if base >= LOAD_LIMIT {
        crate::memory::free_pages(base, page_count);
        return Err(LoadError::TooBig);
    }
    // Identity-map user+writable: single-privilege-level kernel, no W^X —
    // the image contains code that must execute and data that relocates.
    if !paging::map_range(base, page_count, paging::PAGE_USER | paging::PAGE_WRITABLE) {
        crate::memory::free_pages(base, page_count);
        return Err(LoadError::MapFailed);
    }

    let entry_rva = match parse_and_relocate(bytes, base) {
        Ok(e) => e,
        Err(e) => {
            crate::memory::free_pages(base, page_count);
            return Err(e);
        }
    };
    // Clear BSS tails: sections were copied raw-size bytes each; anything
    // between raw_size and virt_size is zero per PE spec.  parse_and_relocate
    // only copied min(raw,virt); zero-fill the remainder now.
    zero_bss_tails(bytes, base);

    Ok(Image {
        base,
        page_count,
        entry: base + entry_rva,
    })
}

/// Zero the part of each section between raw_size and virt_size (BSS).
fn zero_bss_tails(img: &[u8], mem: usize) {
    // Re-walk the section table minimally (headers are small).
    let Some(e_lfanew) = u32_at(img, 0x3C) else {
        return;
    };
    let coff = e_lfanew as usize + 4;
    let Some(num_sections) = u16_at(img, coff + 2) else {
        return;
    };
    let Some(opt_size) = u16_at(img, coff + 16) else {
        return;
    };
    let sect_tab = coff + 20 + opt_size as usize;
    for i in 0..num_sections as usize {
        let Some(virt_size) = u32_at(img, sect_tab + i * 40 + 8) else {
            return;
        };
        let Some(raw_size) = u32_at(img, sect_tab + i * 40 + 16) else {
            return;
        };
        let Some(virt_addr) = u32_at(img, sect_tab + i * 40 + 12) else {
            return;
        };
        if virt_size > raw_size {
            unsafe {
                core::ptr::write_bytes(
                    (mem + virt_addr as usize + raw_size as usize) as *mut u8,
                    0,
                    virt_size as usize - raw_size as usize,
                );
            }
        }
    }
}

/// Free a previously loaded image.
///
/// The task that ran the image must be dead — no code may still execute
/// from or reference these pages.
pub unsafe fn unload_image(img: &Image) {
    for i in 0..img.page_count {
        let _ = paging::unmap_page(img.base + i * PAGE_SIZE);
    }
    crate::memory::free_pages(img.base, img.page_count);
}