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

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

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

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

use crate::{mem::MemoryDescriptorExt, protocol_ext::GraphicsOutputExt};

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)?;

    // Find the requested mode
    let mode = gop.match_mode(fb.req_width, fb.req_height).ok_or_else(|| {
        error!(
            "Requested mode is not supported: {}x{}",
            fb.req_width, fb.req_height
        );
        Error::new(Status::INVALID_PARAMETER, ())
    })?;

    gop.set_mode(&mode)?;

    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 = fb.req_width;
    fb.res_height = fb.req_height;
    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 load_kernel<'a>(
    ih: Handle,
    st: &SystemTable<Boot>,
) -> Result<(u64, u64, &'a mut LoadProtocolV1), Error> {
    let bs = st.boot_services();

    let mut fs = boot_partition(ih, bs)?;
    let mut root = fs.open_volume()?;

    let mut kernel_obj = Object::open(&mut 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, &mut root, cstr16!("initrd.img"))?;
    debug!(
        "Loaded initrd: {:#x?}",
        initrd_start..initrd_start + initrd_size
    );

    // 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.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, mmap_memory, proto_data))
}

unsafe fn map_and_enter_kernel(
    st: SystemTable<Boot>,
    proto_data: &mut LoadProtocolV1,
    mmap_memory: 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, options(noreturn, att_syntax));
}

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

    let (entry, mmap_memory, proto_data) = load_kernel(image_handle, &system_table).unwrap();

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