#![allow(static_mut_refs)]
#![no_std]
#![no_main]

pub mod config;
pub mod elf;
pub mod initrd;
pub mod mem;
pub mod protocol_ext;

use core::{arch::asm, mem::size_of, ops::Deref};

use config::Config;
use elf::Object;
use log::{debug, error, info};
use uefi::{
    prelude::*,
    proto::{
        console::gop::{self, GraphicsOutput, PixelFormat},
        device_path::DevicePath,
        loaded_image::LoadedImage,
        media::{file::Directory, fs::SimpleFileSystem},
    },
    table::{
        boot::{AllocateType, MemoryType, ScopedProtocol},
        cfg,
    },
    Error,
};
use yboot_proto::{
    v1::{self, AvailableMemoryRegion, FramebufferOption},
    LoadProtocolV1, LOADER_MAGIC,
};

use crate::mem::MemoryDescriptorExt;

fn mode_score(mode: &gop::Mode) -> usize {
    let (w, h) = mode.info().resolution();
    let mut size_score = w * h;
    if w > 1920 || h > 1080 {
        // Don't pick too large sizes
        size_score = 0;
    }
    size_score
}

fn setup_framebuffer(bs: &BootServices, fb: &mut FramebufferOption) -> Result<(), Error> {
    let gop_handle = bs.get_handle_for_protocol::<GraphicsOutput>()?;
    let mut gop = bs.open_protocol_exclusive::<GraphicsOutput>(gop_handle)?;

    let mode = gop.modes().max_by_key(mode_score).ok_or_else(|| {
        error!("No mode found");
        Error::new(Status::INVALID_PARAMETER, ())
    })?;

    gop.set_mode(&mode)?;

    let (res_width, res_height) = mode.info().resolution();

    let mut result = gop.frame_buffer();

    let format = match mode.info().pixel_format() {
        PixelFormat::Bgr => v1::PIXEL_B8G8R8A8,
        PixelFormat::Rgb => v1::PIXEL_R8G8B8A8,
        _ => 0,
    };

    fb.res_width = res_width as _;
    fb.res_height = res_height as _;
    fb.res_address = result.as_mut_ptr() as _;
    fb.res_stride = mode.info().stride() as u64 * 4;
    fb.res_size = result.size() as _;
    fb.res_format = format;

    info!(
        "Framebuffer: {}x{} @ {:#x}",
        fb.res_width, fb.res_height, fb.res_address
    );

    Ok(())
}

fn locate_rsdp(st: &SystemTable<Boot>) -> Option<u64> {
    for entry in st.config_table() {
        if entry.guid == cfg::ACPI_GUID {
            return Some(entry.address as u64);
        }
    }

    None
}

fn boot_partition(
    image: Handle,
    bs: &BootServices,
) -> Result<ScopedProtocol<'_, SimpleFileSystem>, Error> {
    let loaded_image = bs.open_protocol_exclusive::<LoadedImage>(image)?;
    let device_handle = loaded_image.device();

    let device_path = bs.open_protocol_exclusive::<DevicePath>(device_handle)?;
    let mut device_path = device_path.deref();

    let fs_handle = bs.locate_device_path::<SimpleFileSystem>(&mut device_path)?;
    bs.open_protocol_exclusive::<SimpleFileSystem>(fs_handle)
}

fn open_root(image: Handle, bs: &BootServices) -> Result<Directory, Error> {
    let mut boot_partition = boot_partition(image, bs)?;
    boot_partition.open_volume()
}

fn load_kernel<'a>(
    config: &Config,
    root: &mut Directory,
    st: &SystemTable<Boot>,
) -> Result<(u64, u64, u64, &'a mut LoadProtocolV1), Error> {
    let bs = st.boot_services();

    let mut kernel_obj = Object::open(root, cstr16!("kernel.elf"))?;
    let loaded_obj = kernel_obj.load(bs)?;

    debug!("Loaded object: {:#x?}", loaded_obj);

    // Load initrd
    let (initrd_start, initrd_size) = initrd::load_somewhere(bs, root, cstr16!("initrd.img"))?;
    debug!(
        "Loaded initrd: {:#x?}",
        initrd_start..initrd_start + initrd_size
    );

    // Load cmdline
    let cmdline = if config.cmdline_len != 0 {
        let address = bs.allocate_pages(AllocateType::AnyPages, MemoryType::LOADER_DATA, 1)?;

        let dst =
            unsafe { core::slice::from_raw_parts_mut(address as *mut u8, config.cmdline_len) };
        dst.copy_from_slice(&config.cmdline[..config.cmdline_len]);

        debug!("Cmdline at {:#x?}", address);

        address
    } else {
        0
    };

    // Other versions are not existent yet
    assert_eq!(loaded_obj.protocol_version, 1);
    let proto_data = unsafe { &mut *(loaded_obj.protocol_struct_paddr as *mut LoadProtocolV1) };

    let rsdp = locate_rsdp(st).ok_or_else(|| {
        error!("Failed to find a RSDP config entry");
        Error::new(Status::NOT_FOUND, ())
    })?;
    info!("RSDP at {:#x}", rsdp);

    proto_data.cmdline = cmdline;
    proto_data.cmdline_len = config.cmdline_len as _;

    proto_data.rsdp_address = rsdp;
    proto_data.initrd_address = initrd_start;
    proto_data.initrd_size = initrd_size;

    // 32768 bytes should be enough
    let mmap_memory = bs.allocate_pages(AllocateType::AnyPages, MemoryType::LOADER_DATA, 8)?;
    info!("Memory map is placed at {:#x}", mmap_memory);

    // Setup framebuffer if needed
    if proto_data.opt_framebuffer.req_width != 0 && proto_data.opt_framebuffer.req_height != 0 {
        setup_framebuffer(bs, &mut proto_data.opt_framebuffer)?;
    }

    // Not yet supported
    if proto_data.kernel_virt_offset == 0 {
        todo!();
    }

    // TODO handle other offsets
    assert_eq!(proto_data.kernel_virt_offset, 0xFFFFFF8000000000);
    info!(
        "Kernel entry will be at {:#x}",
        loaded_obj.entry + proto_data.kernel_virt_offset
    );

    let entry = loaded_obj.entry + proto_data.kernel_virt_offset;

    Ok((entry, loaded_obj.load_address, mmap_memory, proto_data))
}

unsafe fn map_and_enter_kernel(
    st: SystemTable<Boot>,
    proto_data: &mut LoadProtocolV1,
    mmap_memory: u64,
    load_base: u64,
    entry: u64,
) -> ! {
    let (_, mmap) = st.exit_boot_services();
    let mut index = 0;

    let mmap_data = core::slice::from_raw_parts_mut(
        mmap_memory as *mut AvailableMemoryRegion,
        8 * 0x1000 / size_of::<AvailableMemoryRegion>(),
    );

    for entry in mmap.entries() {
        if entry.is_runtime_usable() {
            mmap_data[index] = AvailableMemoryRegion {
                start_address: entry.phys_start,
                page_count: entry.page_count,
            };
            index += 1;
        }
    }

    proto_data.memory_map.address = mmap_memory;
    proto_data.memory_map.len = index as _;

    let cr3 = mem::map_image();
    asm!("cli; wbinvd; mov {0}, %cr3", in(reg) cr3, options(att_syntax));

    asm!("jmp *{0}", in(reg) entry, in("eax") LOADER_MAGIC, in("ecx") load_base, options(noreturn, att_syntax));
}

#[entry]
fn efi_main(image_handle: Handle, mut system_table: SystemTable<Boot>) -> Status {
    if uefi_services::init(&mut system_table).is_err() {
        return Status::LOAD_ERROR;
    }

    let bs = system_table.boot_services();

    let mut root = match open_root(image_handle, bs) {
        Ok(root) => root,
        Err(error) => {
            error!("Could not open boot partition root: {error:?}");
            return Status::LOAD_ERROR;
        }
    };

    let config = match Config::load(&mut root, cstr16!("yboot.cfg")) {
        Ok(config) => config,
        Err(error) => {
            error!("Malformed yboot.cfg: {error:?}");
            return Status::LOAD_ERROR;
        }
    };

    let (entry, load_base, mmap_memory, proto_data) =
        match load_kernel(&config, &mut root, &system_table) {
            Ok(e) => e,
            Err(error) => {
                error!("Failed to load the kernel/initrd: {error:?}");
                return Status::LOAD_ERROR;
            }
        };

    unsafe {
        map_and_enter_kernel(system_table, proto_data, mmap_memory, load_base, entry);
    }
}
