use core::mem::size_of;

use bytemuck::Zeroable;
use log::{error, info};
use types::{Rela, SHT_RELA};
// TODO use 'elf' crate
use uefi::{
    prelude::BootServices,
    proto::media::file::{File, FileAttribute, FileMode, RegularFile},
    table::boot::{AllocateType, MemoryType},
    CStr16, Error, Status,
};
use yboot_proto::LoadProtocolV1;

use crate::elf::types::{PT_LOAD, SHF_ALLOC, SHF_WRITE, SHT_PROGBITS};

use self::types::{Ehdr, Phdr, Shdr};

#[allow(unused)]
mod types {
    use bytemuck::{Pod, Zeroable};

    pub type Off = u64;
    pub type Addr = u64;
    pub type Half = u16;
    pub type Word = u32;
    pub type XWord = u64;
    pub type SXWord = i64;

    pub const PT_LOAD: Word = 1;

    pub const SHT_PROGBITS: Word = 1;
    pub const SHT_RELA: Word = 4;

    pub const SHF_WRITE: XWord = 1 << 0;
    pub const SHF_ALLOC: XWord = 1 << 1;

    pub const R_X86_64_RELATIVE: u32 = 8;

    #[derive(Clone, Copy, Zeroable, Pod)]
    #[repr(C)]
    pub struct Ehdr {
        pub ident: [u8; 16],
        pub type_: Half,
        pub machine: Half,
        pub version: Word,
        pub entry: Addr,
        pub phoff: Off,
        pub shoff: Off,
        pub flags: Word,
        pub ehsize: Half,
        pub phentsize: Half,
        pub phnum: Half,
        pub shentsize: Half,
        pub shnum: Half,
        pub shstrndx: Half,
    }

    #[derive(Clone, Copy, Zeroable, Pod)]
    #[repr(C)]
    pub struct Shdr {
        pub name: Word,
        pub type_: Word,
        pub flags: XWord,
        pub addr: Addr,
        pub offset: Off,
        pub size: XWord,
        pub link: Word,
        pub info: Word,
        pub addralign: XWord,
        pub entsize: XWord,
    }

    #[derive(Clone, Copy, Zeroable, Pod)]
    #[repr(C)]
    pub struct Phdr {
        pub type_: Word,
        pub flags: Word,
        pub offset: Off,
        pub vaddr: Addr,
        pub paddr: Addr,
        pub filesz: XWord,
        pub memsz: XWord,
        pub align: XWord,
    }

    #[derive(Clone, Copy, Zeroable, Pod)]
    #[repr(C)]
    pub struct Rela {
        pub offset: Addr,
        pub info: XWord,
        pub addend: SXWord,
    }

    impl Rela {
        pub fn r_type(&self) -> u32 {
            self.info as u32
        }
    }
}

// Maximum address this loader can map in the target kernel
pub const MAX_KERNEL_PHYS: u64 = 0x100000000;

pub struct Object {
    file: RegularFile,
    ehdr: Ehdr,
}

#[derive(Debug)]
pub struct LoadedObject {
    pub image_start: u64,
    pub image_end: u64,

    pub load_address: u64,

    pub entry: u64,

    pub protocol_struct_paddr: u64,
    pub protocol_version: u32,
}

#[derive(Debug)]
struct LocatedProtocol {
    address: usize,
    offset: u64,
    size: usize,
}

struct RelaSection {
    offset: u64,
    entry_count: usize,
    entry_size: usize,
}

trait ReadExact {
    fn read_exact(&mut self, buf: &mut [u8]) -> Result<(), Error>;
}

impl ReadExact for RegularFile {
    fn read_exact(&mut self, buf: &mut [u8]) -> Result<(), Error> {
        if self.read(buf)? != buf.len() {
            todo!();
        } else {
            Ok(())
        }
    }
}

impl RelaSection {
    pub fn from_shdr(shdr: &Shdr) -> Option<Self> {
        if shdr.type_ != SHT_RELA {
            return None;
        }

        let entry_size = shdr.entsize as usize;
        let entry_count = shdr.size as usize / entry_size;

        Some(Self {
            offset: shdr.offset,
            entry_size,
            entry_count,
        })
    }
}

impl Object {
    pub fn open<D: File>(root: &mut D, path: &CStr16) -> Result<Self, Error> {
        let file = root.open(path, FileMode::Read, FileAttribute::empty())?;
        let mut file = file.into_regular_file().unwrap();

        let mut ehdr = Ehdr::zeroed();
        file.set_position(0)?;
        file.read_exact(bytemuck::bytes_of_mut(&mut ehdr))?;

        // Validate the the image is indeed ELF
        if &ehdr.ident[..4] != b"\x7FELF" {
            error!("Image has invalid ELF magic");
            return Err(Error::new(Status::LOAD_ERROR, ()));
        }

        // Validate that we're loading a x86_64 image
        if ehdr.ident[4] != 2 {
            error!("Image is not 64-bit");
            return Err(Error::new(Status::LOAD_ERROR, ()));
        }

        // // Check that the entry point is set
        // if ehdr.entry == 0 {
        //     error!("Image does not have a valid entry point");
        //     return Err(Error::new(Status::LOAD_ERROR, ()));
        // }

        Ok(Self { file, ehdr })
    }

    // Returns the image entry point
    pub fn load(&mut self, bs: &BootServices) -> Result<LoadedObject, Error> {
        // 1. Locate the protocol data structure
        let Some(loc_proto) = self.locate_protocol_data()? else {
            error!("The image does not have a valid protocol structure");
            todo!();
        };

        info!("Protocol structure @ {:#x?}", loc_proto);
        // No other protocol versions currently implemented, so just try to interpret this as V1
        if loc_proto.size < size_of::<LoadProtocolV1>() {
            error!("The image's protocol structure has invalid size");
            todo!();
        }

        // TODO avoid reading the whole protocol struct just to obtain the virtual offset
        let mut proto_data = LoadProtocolV1::zeroed();
        self.file.set_position(loc_proto.offset)?;
        self.file
            .read_exact(bytemuck::bytes_of_mut(&mut proto_data))?;

        // 2. Find the kernel's range and check that the loaded physical addresses are actually
        //    usable from UEFI
        let mut image_start = u64::MAX;
        let mut image_end = 0;

        for i in 0..self.ehdr.phnum {
            let phdr = self.read_phdr(i as _)?;

            if phdr.type_ != PT_LOAD {
                continue;
            }
            if phdr.paddr + phdr.memsz >= MAX_KERNEL_PHYS {
                error!(
                    "Kernel segment cannot be mapped (above 4GiB): {:#x?}",
                    phdr.paddr..phdr.paddr + phdr.memsz
                );
                todo!();
            }

            let aligned_start = phdr.paddr & !0xFFF;
            let aligned_end = (phdr.paddr + phdr.memsz + 0xFFF) & !0xFFF;

            if aligned_start < image_start {
                image_start = aligned_start;
            }
            if aligned_end > image_end {
                image_end = aligned_end;
            }
        }

        assert!(image_end > image_start);
        assert_eq!(image_start & 0xFFF, 0);
        assert_eq!(image_end & 0xFFF, 0);

        // Allocate memory to load the kernel into
        let kernel_load_address = bs
            .allocate_pages(
                AllocateType::MaxAddress(0xFFFFFFFF),
                MemoryType::LOADER_DATA,
                (image_end - image_start) as usize / 0x1000,
            )
            .expect("Could not allocate memory for the kernel");

        // Print info
        info!("Image start: {:#x}, end: {:#x}", image_start, image_end);
        info!(
            "Kernel virtual offset: {:#x}",
            proto_data.kernel_virt_offset
        );
        info!("Kernel load address: {kernel_load_address:#x}");

        // 3. Load the segments
        for i in 0..self.ehdr.phnum {
            let phdr = self.read_phdr(i as _)?;

            if phdr.type_ != PT_LOAD {
                continue;
            }

            let segment_load_base = phdr.paddr + kernel_load_address;

            info!(
                "[{i}] Load {:#x?}",
                segment_load_base..segment_load_base + phdr.memsz
            );

            if phdr.filesz > 0 {
                let dst_slice = unsafe {
                    core::slice::from_raw_parts_mut(
                        segment_load_base as *mut u8,
                        phdr.filesz as usize,
                    )
                };

                self.file.set_position(phdr.offset)?;
                self.file.read_exact(dst_slice)?;
            }

            if phdr.memsz > phdr.filesz {
                let dst_slice = unsafe {
                    core::slice::from_raw_parts_mut(
                        (segment_load_base + phdr.filesz) as *mut u8,
                        (phdr.memsz - phdr.filesz) as usize,
                    )
                };

                dst_slice.fill(0);
            }
        }

        // 4. Perform kernel relocation
        let mut rela_section = None;
        for i in 0..self.ehdr.shnum as usize {
            let shdr = self.read_shdr(i)?;
            if let Some(rela) = RelaSection::from_shdr(&shdr) {
                rela_section = Some(rela);
                break;
            }
        }

        if let Some(rela_section) = rela_section {
            info!("Relocating kernel: {image_start:#x} -> {kernel_load_address:#x}");
            info!("({} relocations)", rela_section.entry_count);

            let b = (kernel_load_address + proto_data.kernel_virt_offset) as i64;

            for i in 0..rela_section.entry_count {
                let mut rela = Rela::zeroed();
                self.file
                    .set_position(rela_section.offset + (i * rela_section.entry_size) as u64)?;
                self.file.read_exact(bytemuck::bytes_of_mut(&mut rela))?;

                match rela.r_type() {
                    types::R_X86_64_RELATIVE => {
                        let qword = (rela.offset + kernel_load_address) as *mut i64;
                        let value = rela.addend + b;
                        unsafe { qword.write_volatile(value) };
                    }
                    other => todo!("Unsupported relocation type: {other}"),
                }
            }
        }

        // Now that the image is in memory, protocol structure can be written in the further steps
        let protocol_struct_paddr = loc_proto.address as u64 + kernel_load_address; // (loc_proto.address as u64) - proto_data.kernel_virt_offset;
        let protocol_version = proto_data.header.version;

        let entry = self.ehdr.entry + kernel_load_address;

        Ok(LoadedObject {
            image_start,
            image_end,
            load_address: kernel_load_address,
            entry,
            protocol_struct_paddr,
            protocol_version,
        })
    }

    fn read_phdr(&mut self, index: usize) -> Result<Phdr, Error> {
        let mut phdr = Phdr::zeroed();
        self.file
            .set_position(self.ehdr.phoff + self.ehdr.phentsize as u64 * index as u64)?;
        self.file.read_exact(bytemuck::bytes_of_mut(&mut phdr))?;
        Ok(phdr)
    }

    fn read_shdr(&mut self, index: usize) -> Result<Shdr, Error> {
        let mut shdr = Shdr::zeroed();
        self.file
            .set_position(self.ehdr.shoff + self.ehdr.shentsize as u64 * index as u64)?;
        self.file.read_exact(bytemuck::bytes_of_mut(&mut shdr))?;
        Ok(shdr)
    }

    fn locate_protocol_data(&mut self) -> Result<Option<LocatedProtocol>, Error> {
        for i in 0..self.ehdr.shnum {
            let shdr = self.read_shdr(i as _)?;

            if shdr.type_ != SHT_PROGBITS {
                continue;
            }
            if (shdr.flags & (SHF_ALLOC | SHF_WRITE)) != SHF_ALLOC | SHF_WRITE {
                continue;
            }

            // The protocol structure must be located at the start of the section
            // Read the magic field from the ELF section data
            let mut magic = 0;
            self.file.set_position(shdr.offset)?;
            self.file.read_exact(bytemuck::bytes_of_mut(&mut magic))?;

            if magic != yboot_proto::KERNEL_MAGIC {
                continue;
            }

            return Ok(Some(LocatedProtocol {
                offset: shdr.offset,
                address: shdr.addr as _,
                size: shdr.size as _,
            }));
        }

        Ok(None)
    }
}
