#![no_std]
#![no_main]

use r_efi::efi;
use turkos_bootinfo::{BOOT_INFO_MAP_CAPACITY, BootInfoV1, ExitBootServicesRetry};

#[cfg(feature = "kernel-file-test")]
mod kernel_file;
#[cfg(feature = "kernel-handoff-test")]
mod kernel_handoff;
#[cfg(feature = "kernel-hash-test")]
mod kernel_hash;

static BOOT_INFO_TEMPLATE: BootInfoV1 = BootInfoV1::empty();
const MAP_CAPACITY: usize = BOOT_INFO_MAP_CAPACITY as usize;

#[repr(C, align(8))]
struct MapStorage([u8; MAP_CAPACITY]);

static mut FIRMWARE_MAP: MapStorage = MapStorage([0; MAP_CAPACITY]);
static mut OWNED_MAP: MapStorage = MapStorage([0; MAP_CAPACITY]);

// The UART address is scoped to the QEMU virt acceptance fixture. Unlike
// UEFI text output, this diagnostic channel remains usable after exit.
fn serial_byte(byte: u8) {
    const UART: usize = 0x0900_0000;
    // SAFETY: QEMU virt maps the PL011 at this address for this acceptance fixture.
    unsafe {
        while core::ptr::read_volatile((UART + 0x18) as *const u32) & (1 << 5) != 0 {}
        core::ptr::write_volatile(UART as *mut u32, u32::from(byte));
    }
}

fn serial(text: &str) {
    for byte in text.bytes() {
        serial_byte(byte);
    }
}

fn serial_hex(value: u64) {
    for shift in (0..16).rev() {
        let nibble = ((value >> (shift * 4)) & 15) as u8;
        serial_byte(if nibble < 10 {
            b'0' + nibble
        } else {
            b'a' + nibble - 10
        });
    }
}

fn digest(bytes: *const u8, len: usize) -> u64 {
    let mut hash = 0xcbf2_9ce4_8422_2325u64;
    for index in 0..len {
        // SAFETY: caller supplies an initialized map buffer and checked length.
        hash ^= u64::from(unsafe { core::ptr::read_volatile(bytes.add(index)) });
        hash = hash.wrapping_mul(0x0000_0100_0000_01b3);
    }
    hash
}

struct MapSnapshot {
    bytes: usize,
    key: usize,
    descriptor_bytes: usize,
    descriptor_version: u32,
    digest: u64,
}

fn get_owned_map(services: *mut efi::BootServices) -> Result<MapSnapshot, efi::Status> {
    let mut bytes = MAP_CAPACITY;
    let mut key = 0;
    let mut descriptor_bytes = 0;
    let mut descriptor_version = 0;
    let firmware = core::ptr::addr_of_mut!(FIRMWARE_MAP).cast::<efi::MemoryDescriptor>();
    // SAFETY: firmware owns the validated BootServices table. The aligned
    // static storage is writable and has `bytes` bytes of capacity.
    let status = unsafe {
        ((*services).get_memory_map)(
            &mut bytes,
            firmware,
            &mut key,
            &mut descriptor_bytes,
            &mut descriptor_version,
        )
    };
    if status != efi::Status::SUCCESS {
        return Err(status);
    }
    if bytes == 0
        || bytes > MAP_CAPACITY
        || descriptor_bytes < 40
        || descriptor_bytes > u32::MAX as usize
    {
        return Err(efi::Status::COMPROMISED_DATA);
    }
    let mut info = BootInfoV1::empty();
    info.memory_map_ptr = core::ptr::addr_of!(OWNED_MAP).cast::<u8>() as u64;
    info.memory_map_bytes = bytes as u64;
    info.descriptor_bytes = descriptor_bytes as u32;
    info.descriptor_version = descriptor_version;
    if info.validate_memory_map_metadata().is_err() {
        return Err(efi::Status::COMPROMISED_DATA);
    }
    let owned = core::ptr::addr_of_mut!(OWNED_MAP).cast::<u8>();
    // SAFETY: both static regions are disjoint, aligned and at least `bytes` long.
    unsafe { core::ptr::copy_nonoverlapping(firmware.cast::<u8>(), owned, bytes) };
    let original_digest = digest(firmware.cast::<u8>(), bytes);
    let copied_digest = digest(owned, bytes);
    if original_digest != copied_digest {
        return Err(efi::Status::COMPROMISED_DATA);
    }
    Ok(MapSnapshot {
        bytes,
        key,
        descriptor_bytes,
        descriptor_version,
        digest: copied_digest,
    })
}

fn fail(status: efi::Status) -> usize {
    serial("TURKOS/UEFI-ERROR: status=0x");
    serial_hex(status.as_usize() as u64);
    serial("\n");
    status.as_usize()
}

#[panic_handler]
fn panic(_info: &core::panic::PanicInfo<'_>) -> ! {
    loop {
        core::hint::spin_loop();
    }
}

#[unsafe(export_name = "efi_main")]
pub extern "efiapi" fn entry(
    _image_handle: *mut core::ffi::c_void,
    _system_table: *mut core::ffi::c_void,
) -> usize {
    // SAFETY: this immutable, statically initialized BootInfo has a valid address.
    // A later handoff gate will populate payload fields and pass it to the kernel.
    let header = unsafe { core::ptr::read_volatile(&BOOT_INFO_TEMPLATE) };
    if header.validate_header().is_err() {
        return (1usize << (usize::BITS - 1)) | 1;
    }
    if _system_table.is_null() {
        return fail(efi::Status::INVALID_PARAMETER);
    }
    let table = _system_table.cast::<efi::SystemTable>();
    // SAFETY: the UEFI entry ABI passes a valid SystemTable pointer.
    let services = unsafe { (*table).boot_services };
    if services.is_null() {
        return fail(efi::Status::INVALID_PARAMETER);
    }
    serial("TURKOS/UEFI-ENTRY: PASS\n");
    #[cfg(feature = "lifecycle-shutdown-test")]
    {
        // Disposable firmware acceptance only. Do not install this helper as
        // TurkOS BOOT0 or claim that the kernel implements guest poweroff.
        let runtime = unsafe { (*table).runtime_services };
        if runtime.is_null() {
            return fail(efi::Status::INVALID_PARAMETER);
        }
        serial("TURKOS/P08-CLEAN: QMP-OBSERVER-WINDOW\n");
        let status = unsafe { ((*services).stall)(5_000_000) };
        if status != efi::Status::SUCCESS {
            return fail(status);
        }
        serial("TURKOS/P08-CLEAN: GUEST-RESET-SHUTDOWN\n");
        unsafe {
            ((*runtime).reset_system)(
                efi::RESET_SHUTDOWN,
                efi::Status::SUCCESS,
                0,
                core::ptr::null_mut(),
            );
        }
    }
    #[cfg(feature = "kernel-file-test")]
    let loaded_bytes = match kernel_file::load(_image_handle, services) {
        Ok(bytes) => bytes,
        Err(status) => return fail(status),
    };
    #[cfg(feature = "kernel-hash-test")]
    if let Err(status) = kernel_hash::verify(kernel_file::loaded_bytes(loaded_bytes)) {
        return fail(status);
    }
    #[cfg(feature = "kernel-handoff-test")]
    let handoff =
        match kernel_handoff::prepare(table, services, kernel_file::loaded_bytes(loaded_bytes)) {
            Ok(prepared) => prepared,
            Err(status) => return fail(status),
        };
    #[cfg(all(feature = "kernel-file-test", not(feature = "kernel-hash-test")))]
    let _ = loaded_bytes;
    let first = match get_owned_map(services) {
        Ok(map) => map,
        Err(status) => return fail(status),
    };
    serial("TURKOS/UEFI-MAP-COPY: bytes=0x");
    serial_hex(first.bytes as u64);
    serial(" descriptor=0x");
    serial_hex(first.descriptor_bytes as u64);
    serial(" version=0x");
    serial_hex(u64::from(first.descriptor_version));
    serial(" digest=0x");
    serial_hex(first.digest);
    serial("\n");

    // The copy has its own static storage. Refresh the map immediately before
    // ExitBootServices because prior firmware activity may have changed its key.
    let mut retry = ExitBootServicesRetry::new();
    while let Some(attempt) = retry.next_attempt() {
        let current = match get_owned_map(services) {
            Ok(map) => map,
            Err(status) => return fail(status),
        };
        #[cfg(feature = "stale-map-key-test")]
        if attempt == 1 {
            let mut invalidating_page = 0u64;
            // Deliberate negative fixture: mutate the firmware map after
            // obtaining its key, then prove the first exit is rejected.
            let injected = unsafe {
                ((*services).allocate_pages)(
                    efi::ALLOCATE_ANY_PAGES,
                    efi::LOADER_DATA,
                    1,
                    &mut invalidating_page,
                )
            };
            if injected != efi::Status::SUCCESS {
                return fail(injected);
            }
            serial("TURKOS/UEFI-KEY-FAULT: attempt=0x");
            serial_hex(u64::from(attempt));
            serial(" page=0x");
            serial_hex(invalidating_page);
            serial(" ALLOCATED\n");
        }
        // SAFETY: `key` came from the immediately preceding successful map
        // query, and no boot service is called before the exit attempt.
        let status = unsafe { ((*services).exit_boot_services)(_image_handle, current.key) };
        if status == efi::Status::SUCCESS {
            let owned = core::ptr::addr_of!(OWNED_MAP).cast::<u8>();
            let after = digest(owned, current.bytes);
            serial("TURKOS/UEFI-EXIT: attempt=0x");
            serial_hex(u64::from(attempt));
            serial(" bytes=0x");
            serial_hex(current.bytes as u64);
            serial(" before=0x");
            serial_hex(current.digest);
            serial(" after=0x");
            serial_hex(after);
            serial(if after == current.digest {
                " PASS\n"
            } else {
                " FAIL\n"
            });
            #[cfg(feature = "kernel-handoff-test")]
            unsafe {
                let handoff = match kernel_handoff::finalize(handoff, current) {
                    Ok(ready) => ready,
                    Err(_) => {
                        serial("TURKOS/UEFI-BOOTINFO-ERROR\n");
                        loop {
                            core::hint::spin_loop();
                        }
                    }
                };
                kernel_handoff::jump(handoff);
            }
            #[cfg(not(feature = "kernel-handoff-test"))]
            loop {
                core::hint::spin_loop();
            }
        }
        if status != efi::Status::INVALID_PARAMETER {
            return fail(status);
        }
        #[cfg(feature = "stale-map-key-test")]
        {
            serial("TURKOS/UEFI-KEY-RETRY: attempt=0x");
            serial_hex(u64::from(attempt));
            serial(" INVALID-KEY\n");
        }
    }
    fail(efi::Status::INVALID_PARAMETER)
}