//! ELF binary format support
use core::ops::DerefMut;

use alloc::sync::Arc;
use cfg_if::cfg_if;
use elf::{
    endian::AnyEndian,
    relocation::{Rel, Rela},
    segment::ProgramHeader,
    ElfStream, ParseError,
};
use libk_mm::{
    pointer::PhysicalRefMut,
    process::{ProcessAddressSpace, VirtualRangeBacking},
    table::MapAttributes,
    L3_PAGE_SIZE,
};
use libk_util::io::{Read, Seek};
use yggdrasil_abi::{error::Error, io::SeekFrom, path::PathBuf};

use crate::{
    random,
    task::{process::ProcessImage, types::TlsImage},
};

cfg_if! {
    if #[cfg(target_arch = "x86_64")] {
        const EXPECTED_ELF_MACHINE: u16 = elf::abi::EM_X86_64;
    } else if #[cfg(target_arch = "aarch64")] {
        const EXPECTED_ELF_MACHINE: u16 = elf::abi::EM_AARCH64;
    } else if #[cfg(target_arch = "x86")] {
        const EXPECTED_ELF_MACHINE: u16 = elf::abi::EM_386;
    } else if #[cfg(target_arch = "riscv64")] {
        const EXPECTED_ELF_MACHINE: u16 = elf::abi::EM_RISCV;
    }
}

cfg_if! {
    if #[cfg(target_pointer_width = "64")] {
        const EXPECTED_ELF_CLASS: elf::file::Class = elf::file::Class::ELF64;
    } else {
        const EXPECTED_ELF_CLASS: elf::file::Class = elf::file::Class::ELF32;
    }
}

#[derive(Clone)]
pub struct FileReader<F: Read + Seek> {
    file: F,
}

#[derive(Debug, PartialEq)]
pub enum ElfSegmentType {
    Load,
    Dynamic,
    Tls,
}

pub enum ElfKind<F: Read + Seek> {
    Dynamic(Option<PathBuf>),
    Static(ElfStream<AnyEndian, FileReader<Arc<F>>>, FileReader<Arc<F>>),
}

pub struct ElfSegment {
    pub ty: ElfSegmentType,
    pub attrs: MapAttributes,
    pub offset: u64,
    pub align: u64,
    pub vaddr: u64,
    pub full_size: u64,
    pub data_size: u64,
}

impl<F: Read + Seek> elf::io_traits::InputStream for FileReader<F> {
    fn read_exact(&mut self, buf: &mut [u8]) -> Result<(), elf::io_traits::StreamError> {
        self.file.read_exact(buf).map_err(conv_stream_error)
    }

    fn seek(&mut self, pos: elf::io_traits::SeekFrom) -> Result<u64, elf::io_traits::StreamError> {
        self.file
            .seek(conv_seek_from(pos))
            .map_err(conv_stream_error)
    }
}

impl ElfSegmentType {
    pub fn from_phdr(phdr: &ProgramHeader) -> Option<Self> {
        match phdr.p_type {
            elf::abi::PT_LOAD => Some(Self::Load),
            elf::abi::PT_DYNAMIC => Some(Self::Dynamic),
            elf::abi::PT_TLS => Some(Self::Tls),
            _ => None,
        }
    }
}

impl ElfSegment {
    pub fn from_phdr(phdr: &ProgramHeader) -> Option<Self> {
        let ty = ElfSegmentType::from_phdr(phdr)?;

        let attrs = match (phdr.p_flags & elf::abi::PF_W, phdr.p_flags & elf::abi::PF_X) {
            (0, 0) => MapAttributes::USER_READ,
            (_, 0) => MapAttributes::USER_WRITE | MapAttributes::USER_READ,
            (0, _) => MapAttributes::USER_READ,
            (_, _) => MapAttributes::USER_WRITE | MapAttributes::USER_READ,
        } | MapAttributes::NON_GLOBAL;

        Some(Self {
            ty,
            attrs,
            offset: phdr.p_offset,
            vaddr: phdr.p_vaddr,
            align: phdr.p_align,
            full_size: phdr.p_memsz,
            data_size: phdr.p_filesz,
        })
    }

    fn aligned_start(&self) -> usize {
        (self.vaddr & !(self.align - 1)) as usize
    }

    fn aligned_end(&self) -> usize {
        ((self.vaddr + self.full_size + self.align - 1) & !(self.align - 1)) as usize
    }

    fn vaddr_bounds(&self) -> (usize, usize) {
        (self.aligned_start(), self.aligned_end())
    }

    // For LOAD type segments, return the range covered by vaddr..vaddr+memsz, alignment applied
    fn into_vaddr_bounds(self) -> Option<(usize, usize)> {
        if self.ty != ElfSegmentType::Load {
            return None;
        }
        Some((self.aligned_start(), self.aligned_end()))
    }
}

pub fn open_elf_file<F: Read + Seek>(file: Arc<F>) -> Result<ElfKind<F>, Error> {
    let (mut elf, file) = open_elf_direct(file)?;

    if is_dynamic(&mut elf)? {
        // TODO populate interp value
        Ok(ElfKind::Dynamic(None))
    } else {
        Ok(ElfKind::Static(elf, file))
    }
}

#[allow(clippy::type_complexity)]
pub fn open_elf_direct<F: Read + Seek>(
    file: Arc<F>,
) -> Result<(ElfStream<AnyEndian, FileReader<Arc<F>>>, FileReader<Arc<F>>), Error> {
    let file = FileReader { file };
    let elf = ElfStream::<AnyEndian, _>::open_stream(file.clone()).map_err(from_parse_error)?;

    if elf.ehdr.e_machine != EXPECTED_ELF_MACHINE {
        log::warn!("e_machine in ELF ({:#x}) does not match expected value for current architecture ({:#x})", elf.ehdr.e_machine, EXPECTED_ELF_MACHINE);
        return Err(Error::UnrecognizedExecutable);
    }

    // No 32-bit executables currently supported
    if elf.ehdr.class != EXPECTED_ELF_CLASS {
        log::warn!("ELF class is not {:?}", EXPECTED_ELF_CLASS);
        return Err(Error::UnrecognizedExecutable);
    }

    if elf.ehdr.e_type != elf::abi::ET_DYN && elf.ehdr.e_type != elf::abi::ET_EXEC {
        log::warn!("ELF is not an executable");
        return Err(Error::UnrecognizedExecutable);
    }

    Ok((elf, file))
}

/// Loads an ELF binary from `file` into the target address space
pub fn load_elf_from_file<F: Read + Seek>(
    space: &ProcessAddressSpace,
    mut elf: ElfStream<AnyEndian, FileReader<F>>,
    file: FileReader<F>,
    aslr: bool,
) -> Result<ProcessImage, Error> {
    // TODO get information about interpreter: this could be a dynamic (not just relocatable)
    //      executable with required dependency libraries

    // Find out image size
    let (vaddr_min, vaddr_max) = elf_virtual_range(&elf);
    let image_load_size = vaddr_max - vaddr_min;
    let (image_load_base, ip_offset) =
        elf_load_address(aslr, elf.ehdr.e_type, vaddr_min, image_load_size);

    log::debug!(
        "Loading ELF virtual {:#x?} -> real {:#x?}",
        vaddr_min..vaddr_max,
        image_load_base..image_load_base + image_load_size
    );

    // Load the segments
    for segment in elf.segments().iter().filter_map(ElfSegment::from_phdr) {
        match segment.ty {
            ElfSegmentType::Load => {
                load_segment(space, &file, &segment, image_load_base, vaddr_min)?;
            }
            ElfSegmentType::Tls => (),
            ElfSegmentType::Dynamic => (),
        }
    }

    // Load TLS master copy somewhere into process memory
    let tls_image = handle_tls(space, &elf, &file)?;

    // Fixup relocations
    handle_relocations(&mut elf, space, image_load_base)?;

    let entry = elf.ehdr.e_entry as usize + image_load_base - vaddr_min;

    log::debug!("Entry: {:#x}", entry);

    Ok(ProcessImage {
        entry,
        ip_offset,
        load_base: image_load_base,
        tls_image,
    })
}

fn is_dynamic<F: Read + Seek>(
    elf: &mut ElfStream<AnyEndian, FileReader<F>>,
) -> Result<bool, Error> {
    if elf.ehdr.e_type != elf::abi::ET_DYN && elf.ehdr.e_type != elf::abi::ET_EXEC {
        return Ok(false);
    }

    let dynamic = elf.dynamic().map_err(from_parse_error)?;
    let Some(dynamic) = dynamic else {
        // No .dynamic section -> not dynamic
        log::debug!("No .dynamic -> not dynamic");
        return Ok(false);
    };

    Ok(dynamic.into_iter().any(|e| e.d_tag == elf::abi::DT_NEEDED))
}

pub fn elf_virtual_range<F: Read + Seek>(
    elf: &ElfStream<AnyEndian, FileReader<F>>,
) -> (usize, usize) {
    let mut vaddr_min = usize::MAX;
    let mut vaddr_max = usize::MIN;

    for (start, end) in elf
        .segments()
        .iter()
        .filter_map(ElfSegment::from_phdr)
        .filter_map(ElfSegment::into_vaddr_bounds)
    {
        debug_assert!(start < end);
        if start < vaddr_min {
            vaddr_min = start;
        }
        if end > vaddr_max {
            vaddr_max = end;
        }
    }

    assert_ne!(vaddr_min, vaddr_max);
    (vaddr_min, vaddr_max)
}

fn elf_load_address(
    aslr: bool,
    elf_type: u16,
    virtual_address: usize,
    #[allow(unused)] size: usize,
) -> (usize, usize) {
    match elf_type {
        elf::abi::ET_EXEC => (virtual_address, 0),
        elf::abi::ET_DYN => {
            let base = if aslr {
                #[cfg(any(target_pointer_width = "32", rust_analyzer))]
                {
                    let ceiling = (0x30000000 - size) / 0x1000 - 0x500;

                    let index = random::range(0x500, ceiling as _);
                    (index as usize) * 0x1000
                }
                #[cfg(any(target_pointer_width = "64", rust_analyzer))]
                {
                    let index = random::range(0x5000, 0x20000);
                    (index as usize) * 0x1000
                }
            } else {
                0x400000
            };

            (base, base)
        }
        // Handled in load_elf_from_file()
        _ => unreachable!(),
    }
}

fn handle_relocations<F: Read + Seek>(
    elf: &mut ElfStream<AnyEndian, FileReader<F>>,
    space: &ProcessAddressSpace,
    image_load_base: usize,
) -> Result<(), Error> {
    let rel_section = elf
        .section_headers()
        .iter()
        .find(|sh| sh.sh_type == elf::abi::SHT_REL);

    if let Some(&rel_section) = rel_section {
        let rel_iter = elf
            .section_data_as_rels(&rel_section)
            .map_err(from_parse_error)?;

        for rel in rel_iter {
            write_rel(&rel, space, image_load_base)?;
        }
    }

    let rela_section = elf
        .section_headers()
        .iter()
        .find(|sh| sh.sh_type == elf::abi::SHT_RELA);

    if let Some(&rela_section) = rela_section {
        let rela_iter = elf
            .section_data_as_relas(&rela_section)
            .map_err(from_parse_error)?;

        for rela in rela_iter {
            write_rela(&rela, space, image_load_base)?;
        }
    }

    Ok(())
}

fn handle_tls<F: Read + Seek>(
    space: &ProcessAddressSpace,
    elf: &ElfStream<AnyEndian, FileReader<F>>,
    file: &FileReader<F>,
) -> Result<Option<TlsImage>, Error> {
    let mut tls_segments = elf
        .segments()
        .iter()
        .filter(|s| s.p_type == elf::abi::PT_TLS);
    let Some(tls_segment) = tls_segments.next() else {
        return Ok(None);
    };
    if tls_segments.next().is_some() {
        todo!("Handle ELFs with more than one PT_TLS?");
    }

    if !tls_segment.p_align.is_power_of_two() {
        log::warn!("PT_TLS align not power of two: {}", tls_segment.p_align);
        return Err(Error::UnrecognizedExecutable);
    }
    if tls_segment.p_align > 0x1000 {
        todo!("PT_TLS align > 0x1000");
    }

    // Load the master copy somewhere.
    // This will be the pointer passed to the program as an argument so that it can
    // clone the TLS master copy when, for example, spawning a new thread.
    //
    // Only do the allocation if there's any actual data in the master copy,
    // i.e., don't allocate a master copy if there's only a zeroed .tbss
    let master_copy_base = if tls_segment.p_filesz > 0 {
        // NOTE: don't need to load more than filesz bytes, rest will get zeroed anyway
        let page_aligned_size =
            (tls_segment.p_filesz as usize + L3_PAGE_SIZE - 1) & !(L3_PAGE_SIZE - 1);
        log::debug!(
            "Allocate TLS master copy for {:#x} bytes",
            page_aligned_size
        );
        // User doesn't need to write/execute a master copy, it's for cloning only
        let master_copy_base = space.allocate(
            None,
            page_aligned_size,
            VirtualRangeBacking::Anonymous,
            MapAttributes::USER_READ,
        )?;

        load_bytes(
            space,
            master_copy_base,
            tls_segment.p_filesz as usize,
            |off, mut dst| {
                file.file
                    .seek(SeekFrom::Start(tls_segment.p_offset + off as u64))?;
                file.file.read_exact(dst.deref_mut())
            },
        )?;

        master_copy_base
    } else {
        0
    };

    Ok(Some(TlsImage {
        master_copy_base,
        align: tls_segment.p_align as usize,
        data_size: tls_segment.p_filesz as usize,
        full_size: tls_segment.p_memsz as usize,
    }))
}

fn load_segment<F: Read + Seek>(
    space: &ProcessAddressSpace,
    file: &FileReader<F>,
    seg: &ElfSegment,
    load_base: usize,
    elf_virtual_base: usize,
) -> Result<(), Error> {
    assert_ne!(seg.full_size, 0);

    let (virt_start, virt_end) = seg.vaddr_bounds();
    let size = virt_end - virt_start;

    // For ET_EXEC (non-relocatable) elf_virtual_base = load_base, so all addresses will appear the
    // same
    // For ET_DYN (relocatable or shared object) this will setup an offset to operate on
    // "virtual" virtual addresses and "real" virtual addresses
    let relative_offset = seg.vaddr as usize - elf_virtual_base;
    let real_vaddr = load_base + relative_offset;
    let real_vaddr_aligned = real_vaddr & !(seg.align as usize - 1);

    space.map(
        real_vaddr_aligned,
        size,
        VirtualRangeBacking::anonymous(),
        seg.attrs,
    )?;

    if seg.data_size > 0 {
        // Load the actual data
        load_bytes(space, real_vaddr, seg.data_size as _, |off, mut dst| {
            file.file.seek(SeekFrom::Start(seg.offset + off as u64))?;
            file.file.read_exact(dst.deref_mut())
        })?;
    }

    if seg.full_size > seg.data_size {
        let addr = real_vaddr + seg.data_size as usize;
        let len = (seg.full_size - seg.data_size) as usize;

        // Zero the rest
        load_bytes(space, addr, len, |_off, mut dst| {
            dst.fill(0);
            Ok(())
        })?;
    }

    Ok(())
}

fn write_rel(rel: &Rel, space: &ProcessAddressSpace, b: usize) -> Result<(), Error> {
    let rel_field = rel.r_offset as usize + b;

    let (value, width) = match rel.r_type {
        // R_386_RELATIVE
        8 => (b as i32, 4),
        _ => todo!("Unsupported rel type: {:#x}", rel.r_type),
    };

    load_bytes(space, rel_field, width, |off, mut dst| {
        assert_eq!(off, 0);

        match width {
            4 => unsafe {
                let ptr = dst.as_mut_ptr() as *mut u32;
                let old = ptr.read_unaligned();
                ptr.write_unaligned((old as i32 + value) as u32);
                Ok(())
            },
            _ => todo!("Unhandled rel width: {}", width),
        }
    })
}

fn write_rela(rela: &Rela, space: &ProcessAddressSpace, b: usize) -> Result<(), Error> {
    // A: Addend (from entry)
    // B: Image base
    let a = rela.r_addend;

    let rel_field = rela.r_offset as usize + b;

    let (value, width) = match rela.r_type {
        elf::abi::R_X86_64_RELATIVE | elf::abi::R_AARCH64_RELATIVE | elf::abi::R_RISCV_RELATIVE => {
            // B + A
            // Width: qword
            (b as i64 + a, 8)
        }
        _ => todo!("Unsupported rela type: {:#x}", rela.r_type),
    };

    load_bytes(space, rel_field, width, |off, mut dst| {
        assert_eq!(off, 0);

        match width {
            8 => {
                unsafe { (dst.as_mut_ptr() as *mut u64).write_unaligned(value as u64) };
                Ok(())
            }
            _ => todo!("Unhandled rela width: {}", width),
        }
    })
}

#[inline]
fn conv_stream_error(_v: Error) -> elf::io_traits::StreamError {
    elf::io_traits::StreamError {
        message: "Elf read error",
    }
}

#[inline]
fn conv_seek_from(v: elf::io_traits::SeekFrom) -> SeekFrom {
    match v {
        elf::io_traits::SeekFrom::End(off) => SeekFrom::End(off),
        elf::io_traits::SeekFrom::Start(off) => SeekFrom::Start(off),
        _ => todo!(),
    }
}

#[inline]
fn from_parse_error(v: ParseError) -> Error {
    log::warn!("ELF loading error: {:?}", v);
    Error::InvalidFile
}

fn load_bytes<F>(
    space: &ProcessAddressSpace,
    addr: usize,
    len: usize,
    mut src: F,
) -> Result<(), Error>
where
    F: FnMut(usize, PhysicalRefMut<'_, [u8]>) -> Result<(), Error>,
{
    // TODO check for crazy addresses here
    let dst_page_off = addr & 0xFFF;
    let dst_page_aligned = addr & !0xFFF;
    let mut off = 0usize;
    let mut rem = len;

    while rem != 0 {
        let page_idx = (dst_page_off + off) / 0x1000;
        let page_off = (dst_page_off + off) % 0x1000;
        let count = core::cmp::min(rem, 0x1000 - page_off);

        let virt_page = dst_page_aligned + page_idx * 0x1000;
        assert_eq!(virt_page & 0xFFF, 0);

        let phys_page = space.translate(virt_page)?;
        let dst_slice = unsafe { PhysicalRefMut::map_slice(phys_page.add(page_off), count) };

        src(off, dst_slice)?;

        rem -= count;
        off += count;
    }

    Ok(())
}
