//! ELF64 Loader для запуска userspace-процессов в Ring 3.

use x86_64::structures::paging::PageTableFlags;
use x86_64::VirtAddr;

const ELF_MAGIC: [u8; 4] = [0x7F, b'E', b'L', b'F'];
const PT_LOAD: u32 = 1;
const PF_X: u32 = 1;
const PF_W: u32 = 2;
const PF_R: u32 = 4;

/// Стандартный стек пользователя (располагается в верхушке канонического userspace).
pub const USER_STACK_TOP: usize = 0x0000_7FFF_FFFF_0000;
pub const USER_STACK_PAGES: usize = 4; // 16 KiB

static PROCESS_STACK_COUNTER: core::sync::atomic::AtomicUsize = core::sync::atomic::AtomicUsize::new(0);

#[derive(Debug)]
pub struct LoadedElf {
    pub entry_point: usize,
    pub user_stack_top: usize,
}

#[repr(C, packed)]
struct Elf64Header {
    magic: [u8; 4],
    class: u8,
    data: u8,
    version: u8,
    osabi: u8,
    abiversion: u8,
    pad: [u8; 7],
    e_type: u16,
    e_machine: u16,
    e_version: u32,
    e_entry: u64,
    e_phoff: u64,
    e_shoff: u64,
    e_flags: u32,
    e_ehsize: u16,
    e_phentsize: u16,
    e_phnum: u16,
    e_shentsize: u16,
    e_shnum: u16,
    e_shstrndx: u16,
}

#[repr(C, packed)]
struct Elf64ProgramHeader {
    p_type: u32,
    p_flags: u32,
    p_offset: u64,
    p_vaddr: u64,
    p_paddr: u64,
    p_filesz: u64,
    p_memsz: u64,
    p_align: u64,
}

/// Разобрать ELF64 бинарник, отобразить сегменты в память и подготовить структуру.
pub fn load_elf(elf_bytes: &[u8]) -> Result<LoadedElf, &'static str> {
    if elf_bytes.len() < core::mem::size_of::<Elf64Header>() {
        return Err("ELF: File too small");
    }

    let header = unsafe { &*(elf_bytes.as_ptr() as *const Elf64Header) };

    if header.magic != ELF_MAGIC {
        return Err("ELF: Invalid magic header");
    }
    if header.class != 2 {
        return Err("ELF: Not 64-bit");
    }
    if header.e_machine != 0x3E {
        return Err("ELF: Not x86_64");
    }

    let ph_offset = header.e_phoff as usize;
    let ph_count = header.e_phnum as usize;
    let ph_size = header.e_phentsize as usize;

    // 1. Отображаем сегменты PT_LOAD
    for i in 0..ph_count {
        let offset = ph_offset + i * ph_size;
        if offset + ph_size > elf_bytes.len() {
            return Err("ELF: Program header out of bounds");
        }

        let ph = unsafe { &*(elf_bytes.as_ptr().add(offset) as *const Elf64ProgramHeader) };

        if ph.p_type == PT_LOAD {
            let vaddr = ph.p_vaddr as usize;
            let memsz = ph.p_memsz as usize;
            let filesz = ph.p_filesz as usize;
            let file_offset = ph.p_offset as usize;

            let mut flags = PageTableFlags::PRESENT | PageTableFlags::USER_ACCESSIBLE;
            if (ph.p_flags & PF_W) != 0 {
                flags |= PageTableFlags::WRITABLE;
            }
            if (ph.p_flags & PF_X) == 0 {
                flags |= PageTableFlags::NO_EXECUTE;
            }

            let start_page = vaddr & !0xFFF;
            let end_page = (vaddr + memsz + 0xFFF) & !0xFFF;
            let page_count = (end_page - start_page) / 4096;

            for p in 0..page_count {
                let page_vaddr = start_page + p * 4096;
                let page_virt = VirtAddr::new(page_vaddr as u64);

                let page_ptr = if super::vmm::is_mapped(page_virt) {
                    let phys = super::vmm::translate(page_virt).ok_or("Failed to translate mapped page")?;
                    super::vmm::phys_to_virt(phys).as_mut_ptr::<u8>()
                } else {
                    let phys = super::vmm::map_alloc(page_virt, flags)?;
                    let ptr = super::vmm::phys_to_virt(phys).as_mut_ptr::<u8>();
                    unsafe {
                        core::ptr::write_bytes(ptr, 0, 4096);
                    }
                    ptr
                };

                // Копируем данные сегмента, попадающие в эту страницу, через HHDM
                if filesz > 0 {
                    let seg_start = vaddr;
                    let seg_end = vaddr + filesz;

                    let copy_start = core::cmp::max(page_vaddr, seg_start);
                    let copy_end = core::cmp::min(page_vaddr + 4096, seg_end);

                    if copy_start < copy_end {
                        let dst_offset = copy_start - page_vaddr;
                        let src_offset = (copy_start - seg_start) + file_offset;
                        let len = copy_end - copy_start;

                        if src_offset + len <= elf_bytes.len() {
                            unsafe {
                                core::ptr::copy_nonoverlapping(
                                    elf_bytes.as_ptr().add(src_offset),
                                    page_ptr.add(dst_offset),
                                    len,
                                );
                            }
                        }
                    }
                }
            }
        }
    }

    // 2. Выделяем и маппируем пользовательский стек (16 KiB)
    let stack_idx = PROCESS_STACK_COUNTER.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
    let process_stack_top = USER_STACK_TOP - (stack_idx * 0x100000);
    let stack_base = process_stack_top - (USER_STACK_PAGES * 4096);
    let stack_flags = PageTableFlags::PRESENT
        | PageTableFlags::WRITABLE
        | PageTableFlags::USER_ACCESSIBLE
        | PageTableFlags::NO_EXECUTE;

    for p in 0..USER_STACK_PAGES {
        let page_virt = VirtAddr::new((stack_base + p * 4096) as u64);
        let phys = super::vmm::map_alloc(page_virt, stack_flags)?;
        let page_ptr = super::vmm::phys_to_virt(phys).as_mut_ptr::<u8>();
        unsafe {
            core::ptr::write_bytes(page_ptr, 0, 4096);
        }
    }

    let entry = header.e_entry as usize;

    log::info!(
        "[elf] ELF loaded successfully: entry={:#x}, stack={:#x}",
        entry,
        process_stack_top
    );

    Ok(LoadedElf {
        entry_point: entry,
        user_stack_top: process_stack_top,
    })
}

/// Запустить загруженный ELF процесс в Ring 3.
pub fn spawn_user_process(name: &'static str, loaded: LoadedElf) {
    extern "C" fn user_trampoline(args_packed: usize) {
        let (entry, stack) = unsafe {
            let ptr = args_packed as *const (usize, usize);
            let tuple = *ptr;
            drop(alloc::boxed::Box::from_raw(args_packed as *mut (usize, usize)));
            tuple
        };

        log::info!("[userspace] Dropping CPU into Ring 3 (RIP: {:#x}, RSP: {:#x})...", entry, stack);

        unsafe {
            crate::arch::x86_64::context::jump_to_userspace(entry, stack);
        }
    }

    let packed = alloc::boxed::Box::into_raw(alloc::boxed::Box::new((
        loaded.entry_point,
        loaded.user_stack_top,
    ))) as usize;

    crate::sched::spawn(name, user_trampoline, packed, crate::sched::Priority::NORMAL);
}