#![feature(vec_into_raw_parts)]
#![no_main]
#![no_std]

extern crate alloc;

use alloc::{boxed::Box, vec, vec::Vec};
use core::{arch::asm, slice};
use goblin::elf;
use nel_os_common::{gop, memory};
use uefi::{
    allocator::Allocator,
    boot::{AllocateType, MemoryType, ScopedProtocol},
    mem::memory_map::MemoryMap,
    prelude::*,
    println,
    proto::{
        console::gop::{GraphicsOutput, PixelFormat},
        media::{
            file::{Directory, File, FileAttribute, FileInfo, FileMode},
            fs::SimpleFileSystem,
        },
    },
    CStr16,
};

#[global_allocator]
static ALLOCATOR: Allocator = Allocator;

fn hlt_loop() -> ! {
    loop {
        unsafe {
            asm!("hlt");
        }
    }
}

fn get_fs() -> Directory {
    let mut fs: ScopedProtocol<SimpleFileSystem> =
        uefi::boot::get_image_file_system(uefi::boot::image_handle()).unwrap();

    fs.open_volume().unwrap()
}

fn read_file(name: &CStr16) -> Box<[u8]> {
    let mut root = get_fs();
    let file_info = root
        .open(name, FileMode::Read, FileAttribute::empty())
        .unwrap();
    let mut file = file_info.into_regular_file().unwrap();

    let file_size = file.get_boxed_info::<FileInfo>().unwrap().file_size();
    let mut buf = vec![0; file_size as usize];

    let read_size = file.read(&mut buf).unwrap();
    println!("file {} size: {}", name, read_size);

    buf.into_boxed_slice()
}

fn load_file_to_laoder_data(name: &CStr16) -> (u64, u64) {
    let mut root = get_fs();
    let file_info = root
        .open(name, FileMode::Read, FileAttribute::empty())
        .expect("Failed to open file");
    let mut file = file_info
        .into_regular_file()
        .expect("Failed to convert to regular file");

    let file_size = file
        .get_boxed_info::<FileInfo>()
        .expect("Failed to get file info")
        .file_size();

    let page_ptr = unsafe {
        uefi::boot::allocate_pages(
            AllocateType::AnyPages,
            MemoryType::LOADER_DATA,
            file_size.div_ceil(4096) as usize,
        )
        .expect("Failed to allocate pages")
        .as_mut()
    };

    let buf = unsafe { slice::from_raw_parts_mut(page_ptr, file_size as usize) };
    let read_size = file.read(buf).expect("Failed to read file");
    println!("file {} size: {}", name, read_size);

    (page_ptr as *mut u8 as u64, file_size)
}

fn load_elf(bin: Box<[u8]>) -> u64 {
    let elf = elf::Elf::parse(&bin).expect("Failed to parse elf");
    let mut dest_start = u64::MAX;
    let mut dest_end = 0u64;

    elf.program_headers
        .iter()
        .filter(|header| header.p_type == elf::program_header::PT_LOAD)
        .for_each(|header| {
            dest_start = dest_start.min(header.p_vaddr);
            dest_end = dest_end.max(header.p_vaddr + header.p_memsz);
        });

    uefi::boot::allocate_pages(
        AllocateType::Address(dest_start),
        MemoryType::LOADER_DATA,
        dest_end
            .checked_sub(dest_start)
            .and_then(|size| size.checked_add(4095))
            .map(|size| size / 4096)
            .unwrap_or(0) as usize,
    )
    .expect("Failed to allocate pages");

    elf.program_headers
        .iter()
        .filter(|header| header.p_type == elf::program_header::PT_LOAD)
        .for_each(|header| {
            let dest = unsafe {
                slice::from_raw_parts_mut(header.p_vaddr as *mut u8, header.p_memsz as usize)
            };

            let file_size = header.p_filesz as usize;
            let offset = header.p_offset as usize;

            dest[..file_size].copy_from_slice(&bin[offset..offset + file_size]);
            dest[file_size..].fill(0);
        });

    elf.entry
}

fn get_frame_buffer() -> Option<gop::FrameBuffer> {
    let gop_handle = if let Ok(gop_handle) = uefi::boot::get_handle_for_protocol::<GraphicsOutput>()
    {
        gop_handle
    } else {
        println!("GraphicsOutput protocol not found");
        return None;
    };
    let mut gop = if let Ok(gop) = boot::open_protocol_exclusive::<GraphicsOutput>(gop_handle) {
        gop
    } else {
        println!("Failed to open GraphicsOutput protocol");
        return None;
    };

    let info = gop.current_mode_info();
    let (width, height) = info.resolution();
    let frame_buffer = gop.frame_buffer().as_mut_ptr();
    let stride = info.stride();
    let pixel_format = info.pixel_format();

    Some(gop::FrameBuffer {
        frame_buffer,
        width,
        height,
        stride,
        pixl_format: match pixel_format {
            PixelFormat::Rgb => gop::PixelFormat::Rgb,
            PixelFormat::Bgr => gop::PixelFormat::Bgr,
            format => panic!("Unsupported pixel_format: {:?}", format),
        },
    })
}

fn get_rsdp() -> Option<u64> {
    uefi::system::with_config_table(move |c| {
        c.iter()
            .find(|config| config.guid == uefi::table::cfg::ACPI_GUID)
            .map(|config| config.address as u64)
            .or(None)
    })
}

#[entry]
fn main() -> Status {
    uefi::helpers::init().unwrap();

    uefi::system::with_stdout(|stdout| stdout.clear().unwrap());

    println!("{} v{}", env!("CARGO_PKG_NAME"), env!("CARGO_PKG_VERSION"));

    let kernel = read_file(cstr16!("nel_os_kernel.elf"));

    let (bzimage_addr, bzimage_size) = load_file_to_laoder_data(cstr16!("bzImage"));
    let (rootfs_addr, rootfs_size) = load_file_to_laoder_data(cstr16!("rootfs-n.cpio.gz"));

    let entry_point = load_elf(kernel);

    println!("Entry point: {:#x}", entry_point);

    let entry: extern "sysv64" fn(&nel_os_common::BootInfo) =
        unsafe { core::mem::transmute(entry_point) };

    let frame_buffer = get_frame_buffer();

    let rsdp = get_rsdp();

    let size = uefi::boot::memory_map(MemoryType::LOADER_DATA)
        .unwrap()
        .len()
        + 8 * core::mem::size_of::<memory::Range>();
    let mut ranges: Vec<memory::Range> = Vec::with_capacity(size);

    println!("Usable memory table size: {}", size);

    let memory_map = unsafe { uefi::boot::exit_boot_services(Some(MemoryType::LOADER_DATA)) };

    memory_map
        .entries()
        .filter(|entry| {
            matches!(
                entry.ty,
                MemoryType::CONVENTIONAL
                    | MemoryType::BOOT_SERVICES_CODE
                    | MemoryType::BOOT_SERVICES_DATA
            )
        })
        .for_each(|entry| {
            ranges.push(memory::Range {
                start: entry.phys_start,
                end: entry.phys_start + entry.page_count * 4096,
            })
        });

    let usable_memory = {
        let (ptr, len, _) = ranges.into_raw_parts();
        memory::UsableMemory {
            ranges: ptr as *const memory::Range,
            len: len as u64,
        }
    };

    entry(&nel_os_common::BootInfo {
        usable_memory,
        frame_buffer,
        rsdp,
        bzimage_addr,
        bzimage_size,
        rootfs_addr,
        rootfs_size,
    });

    hlt_loop();
}
