/**
 * HELLO, the HOS EFI Lightweight LOader.
 */
module hello.hello;

import uefi;
import hello.console;
import hello.scratch;
import hulk.bootinfo;
import hulk.header;
import hos.page_table;
import hos.cpu;
import hos.memory;
import ldc.llvmasm;

__gshared EFI_SYSTEM_TABLE * st;
extern extern(C) __gshared ubyte _hulk_bin_start;
extern extern(C) __gshared ubyte _hulk_bin_end;

private HulkHeader * hulk_header()
{
    return cast(HulkHeader *)&_hulk_bin_start;
}

private BootInfo * bootinfo()
{
    return &hulk_header().bootinfo;
}

private ulong hulk_bin_phys()
{
    return cast(ulong)&_hulk_bin_start;
}

private ulong hulk_bin_size()
{
    return cast(ulong)&_hulk_bin_end - cast(ulong)&_hulk_bin_start;
}

private ulong hulk_total_size()
{
    return cast(ulong)hulk_header().total_size;
}

private ulong hulk_bss_size()
{
    return hulk_total_size() - hulk_bin_size();
}

private ulong hulk_stack_size()
{
    return hulk_header().stack_size;
}

private ulong hulk_virt_base_address()
{
    return hulk_header().virt_base;
}

private ulong hulk_virt_stack_top()
{
    return hulk_header().virt_stack_top;
}

private ulong hulk_virt_framebuffer_address()
{
    return hulk_header().virt_fb_buffer;
}

/**
 * Detect if we're running in QEMU.
 */
private bool in_qemu()
{
    ulong * firmware_vendor = cast(ulong *) st.FirmwareVendor;
    ulong fv1 = firmware_vendor[0];
    return fv1 ==
        ((cast(ulong)'E') |
         (cast(ulong)'D' << 16) |
         (cast(ulong)'K' << 32) |
         (cast(ulong)' ' << 48));
}

/**
 * Set a graphics mode.
 */
private bool set_graphics_mode()
{
    uint max_horizontal_resolution = in_qemu() ? 1920u : 0xFFFFFFFFu;
    UINTN buffer_size = scratch.free();
    EFI_GUID gop_guid = EFI_GRAPHICS_OUTPUT_PROTOCOL_GUID;
    EFI_HANDLE * handles = cast(EFI_HANDLE *)scratch.current();
    EFI_STATUS status = st.BootServices.LocateHandle(ByProtocol,
            &gop_guid, null, &buffer_size, handles);
    if (status != EFI_SUCCESS)
    {
        console.writeln("LocateHandle: error %x", status);
        return false;
    }
    EFI_HANDLE gop_handle = handles[0];
    EFI_HANDLE gop_interface;
    status = st.BootServices.HandleProtocol(gop_handle,
            &gop_guid, &gop_interface);
    if (status != EFI_SUCCESS)
    {
        console.writeln("HandleProtocol: error %x", status);
        return false;
    }
    if (gop_interface == null)
    {
        console.writeln("null interface from HandleProtocol");
        return false;
    }
    EFI_GRAPHICS_OUTPUT_PROTOCOL * gop = cast(EFI_GRAPHICS_OUTPUT_PROTOCOL *)gop_interface;
    UINTN info_size;
    EFI_GRAPHICS_OUTPUT_MODE_INFORMATION * info;
    uint best_width;
    uint best_mode_number;
    for (uint mode_number = 0u;
         (status = gop.QueryMode(gop, mode_number, &info_size, &info)) == EFI_SUCCESS;
         mode_number++)
    {
        if ((info.HorizontalResolution > best_width) &&
            (info.HorizontalResolution <= max_horizontal_resolution))
        {
            best_width = info.HorizontalResolution;
            best_mode_number = mode_number;
        }
        st.BootServices.FreePool(info);
    }
    if ((status = gop.SetMode(gop, best_mode_number)) != EFI_SUCCESS)
    {
        console.writeln("SetMode: Error %x\n", status);
        return false;
    }
    bootinfo().fb.buffer = cast(uint *)gop.Mode.FrameBufferBase;
    bootinfo().fb.width = gop.Mode.Info.HorizontalResolution;
    bootinfo().fb.height = gop.Mode.Info.VerticalResolution;
    bootinfo().fb.stride = gop.Mode.Info.PixelsPerScanLine;
    bootinfo().fb.format = gop.Mode.Info.PixelFormat;
    return true;
}

/**
 * Walk the EFI memory map and translate it to the HULK bootinfo format.
 */
private void get_memory_map(ulong * bss_phys, ulong * stack_phys, ulong * max_physical_address, UINTN * memory_map_key)
{
    immutable static ubyte[] efi_to_hulk_memory_map_type = [
        BootInfo.MemoryRegion.Type.Reserved, // EfiReservedMemoryType
        BootInfo.MemoryRegion.Type.Conventional, // EfiLoaderCode
        BootInfo.MemoryRegion.Type.Conventional, // EfiLoaderData
        BootInfo.MemoryRegion.Type.Conventional, // EfiBootServicesCode
        BootInfo.MemoryRegion.Type.Conventional, // EfiBootServicesData
        BootInfo.MemoryRegion.Type.Reserved, // EfiRuntimeServicesCode
        BootInfo.MemoryRegion.Type.Reserved, // EfiRuntimeServicesData
        BootInfo.MemoryRegion.Type.Conventional, // EfiConventionalMemory
        BootInfo.MemoryRegion.Type.Unusable, // EfiUnusableMemory
        BootInfo.MemoryRegion.Type.ACPIReclaim, // EfiACPIReclaimMemory
        BootInfo.MemoryRegion.Type.ACPINVS, // EfiACPIMemoryNVS
        BootInfo.MemoryRegion.Type.MemoryMappedIO, // EfiMemoryMappedIO
        BootInfo.MemoryRegion.Type.MemoryMappedIOPortSpace, // EfiMemoryMappedIOPortSpace
        BootInfo.MemoryRegion.Type.PalCode, // EfiPalCode
    ];
    *max_physical_address = 0u;
    UINTN memory_map_size = scratch.free();
    UINTN descriptor_size;
    UINT32 descriptor_version;
    ubyte * scratch_base = scratch.current();
    UINTN status = st.BootServices.GetMemoryMap(
            &memory_map_size,
            cast(EFI_MEMORY_DESCRIPTOR *)scratch_base,
            memory_map_key,
            &descriptor_size,
            &descriptor_version);
    if (status != EFI_SUCCESS)
    {
        console.writeln("GetMemoryMap: Error %x", status);
        for (;;) {}
    }
    size_t n_entries = memory_map_size / descriptor_size;
    size_t count;
    bool found_bss;
    bool found_stack;
    for (count = 0u; count < n_entries; count++)
    {
        if (count > bootinfo().memory_map.length)
        {
            console.writeln("Memory map too large");
            for (;;) {}
        }
        EFI_MEMORY_DESCRIPTOR * descriptor = cast(EFI_MEMORY_DESCRIPTOR *)&scratch_base[count * descriptor_size];
        ulong end_address = descriptor.PhysicalStart + descriptor.NumberOfPages * 4096u;
        if ((end_address > *max_physical_address) &&
            ((descriptor.Type == EfiLoaderCode) ||
             (descriptor.Type == EfiLoaderData) ||
             (descriptor.Type == EfiBootServicesCode) ||
             (descriptor.Type == EfiBootServicesData) ||
             (descriptor.Type == EfiRuntimeServicesCode) ||
             (descriptor.Type == EfiRuntimeServicesData) ||
             (descriptor.Type == EfiConventionalMemory)))
        {
            *max_physical_address = end_address;
        }
        if (descriptor.Type >= efi_to_hulk_memory_map_type.length)
        {
            continue;
        }
        bootinfo().memory_map[count].base = descriptor.PhysicalStart;
        bootinfo().memory_map[count].size = descriptor.NumberOfPages * 4096u;
        bootinfo().memory_map[count].type = efi_to_hulk_memory_map_type[descriptor.Type];
        if ((!found_bss) &&
            (descriptor.Type == EfiConventionalMemory) &&
            (bootinfo().memory_map[count].size >= hulk_bss_size()))
        {
            *bss_phys = bootinfo().memory_map[count].base;
            bootinfo().memory_map[count].base += hulk_bss_size();
            bootinfo().memory_map[count].size -= hulk_bss_size();
            found_bss = true;
        }
        if ((!found_stack) &&
            (descriptor.Type == EfiConventionalMemory) &&
            (bootinfo().memory_map[count].size >= hulk_stack_size()))
        {
            *stack_phys = bootinfo().memory_map[count].base;
            bootinfo().memory_map[count].base += hulk_stack_size();
            bootinfo().memory_map[count].size -= hulk_stack_size();
            found_stack = true;
        }
    }
    bootinfo().memory_map_count = count;
    if ((!found_bss) && (!found_stack))
    {
        for (;;) {}
    }
}

/**
 * Allocate a new page table.
 */
private PageTableEntry * new_page_table()
{
    PageTableEntry * pt = cast(PageTableEntry *)scratch.alloc(1u);
    memset64(pt, 0u, 512);
    return pt;
}

/**
 * Map a virtual address to a physical address using 4KB pages.
 *
 * @param source_page Source page address.
 * @param dest_page Destination page address.
 * @param pt_base Page table base address.
 */
private void map4k(ulong source_page, ulong dest_page, PageTableEntry * pt_base)
{
    PageTableEntry * pt = pt_base;
    PageTableEntry * next_pt;
    for (size_t level = 0; level < 4u; level++)
    {
        size_t pt_index = PageTableEntry.page_table_index(source_page, level);
        if (pt[pt_index].present())
        {
            pt = pt[pt_index].next();
        }
        else
        {
            ulong addr;
            if (level < 3u)
            {
                next_pt = new_page_table();
                addr = cast(ulong)next_pt;
            }
            else
            {
                addr = dest_page;
            }
            pt[pt_index] = PageTableEntry(addr,
                    0u, 0u, 0u, 0u, 0u, 0u, 1u, 1u);
            pt = next_pt;
        }
    }
}

/**
 * Map a virtual region to a physical region using 4KB pages.
 *
 * @param source Source address.
 * @param dest Destination address.
 * @param size Region size.
 * @param pt_base Page table base address.
 */
private void map4kregion(ulong source, ulong dest, size_t size, PageTableEntry * pt_base)
{
    ulong end = source + size;
    while (source < end)
    {
        map4k(source, dest, pt_base);
        source += 4096u;
        dest += 4096u;
    }
}

/**
 * Map a virtual address to a physical address using 2MB pages.
 *
 * @param source_page Source page address.
 * @param dest_page Destination page address.
 * @param pt_base Page table base address.
 */
private void map2m(ulong source_page, ulong dest_page, PageTableEntry * pt_base)
{
    PageTableEntry * pt = pt_base;
    PageTableEntry * next_pt;
    for (size_t level = 0; level < 3u; level++)
    {
        size_t pt_index = PageTableEntry.page_table_index(source_page, level);
        if (pt[pt_index].present())
        {
            pt = pt[pt_index].next();
        }
        else
        {
            if (level < 2u)
            {
                next_pt = new_page_table();
                pt[pt_index] = PageTableEntry(cast(ulong)next_pt,
                        0u, 0u, 0u, 0u, 0u, 0u, 1u, 1u);
            }
            else
            {
                pt[pt_index] = PageTableEntry(dest_page,
                        0u, 0u, 1u, 0u, 0u, 0u, 1u, 1u);
            }
            pt = next_pt;
        }
    }
}

/**
 * Map HULK virtual addresses to physical kernel location.
 *
 * @param pt_base Page table base address.
 */
private void map_hulk(PageTableEntry * pt_base, ulong bss_phys, ulong stack_phys)
{
    /* Map HULK bin region. */
    ulong virt = hulk_virt_base_address();
    map4kregion(virt, hulk_bin_phys(), hulk_bin_size(), pt_base);
    /* Map HULK bss region. */
    virt += hulk_bin_size();
    map4kregion(virt, bss_phys, hulk_bss_size(), pt_base);
    /* Map HULK stack. */
    virt = hulk_virt_stack_top() - hulk_stack_size();
    map4kregion(virt, stack_phys, hulk_stack_size(), pt_base);
}

/**
 * Build page tables in preparation to jump to HULK.
 *
 * @param max_physical_address Maximum physical address to identity map.
 */
private void build_page_tables(ulong max_physical_address, ulong bss_phys, ulong stack_phys)
{
    PageTableEntry * pt_base = new_page_table();
    /* Map physical RAM. */
    for (size_t addr = 0u; addr < max_physical_address; addr += (2u * 1024u * 1024u))
    {
        map2m(addr, addr, pt_base);
    }
    /* Map any memory regions that are outside physical RAM. */
    for (size_t i = 0u; i < bootinfo().memory_map_count; i++)
    {
        ulong addr = bootinfo().memory_map[i].base;
        if (addr >= max_physical_address)
        {
            map4kregion(addr, addr, bootinfo().memory_map[i].size, pt_base);
        }
    }
    /* Map graphics framebuffer. */
    ulong framebuffer_size = bootinfo().fb.height * bootinfo().fb.stride * 4u;
    ulong fb_phys = cast(ulong)bootinfo().fb.buffer;
    ulong fb_virt = hulk_virt_framebuffer_address();
    map4kregion(fb_virt, fb_phys, framebuffer_size, pt_base);
    /* Map HULK regions. */
    map_hulk(pt_base, bss_phys, stack_phys);
    /* Switch to the new page table. */
    write_cr3(cast(ulong)pt_base);
}

/**
 * Jump to HULK entry point.
 */
private void jump_to_hulk()
{
    __asm(
            "jmpq *$0",
            "r,{rsp}",
            hulk_header().entry, hulk_virt_stack_top());
}

/**
 * EFI application entry point.
 *
 * @param image_handle EFI application handle.
 * @param st Pointer to EFI system table.
 */
extern (C) EFI_STATUS efi_main(EFI_HANDLE image_handle, EFI_SYSTEM_TABLE * st)
{
    .st = st;

    console.clear();

    console.writeln("Welcome to HELLO, HOS EFI Lightweight LOader, v0.1.0");
    console.writeln("Firmware vendor: '%S', version: 0x%x", st.FirmwareVendor, st.FirmwareVendor);

    if (!set_graphics_mode())
    {
        console.wait_key();
        return EFI_SUCCESS;
    }

    ulong bss_phys;
    ulong stack_phys;
    ulong max_physical_address;
    UINTN memory_map_key;
    for (;;)
    {
        get_memory_map(&bss_phys, &stack_phys, &max_physical_address, &memory_map_key);

        EFI_STATUS status = st.BootServices.ExitBootServices(image_handle, memory_map_key);
        if (status == EFI_INVALID_PARAMETER)
        {
            continue;
        }
        if (status == EFI_SUCCESS)
        {
            break;
        }
        console.writeln("ExitBootServices: Error %x", status);
        console.wait_key();
        return EFI_SUCCESS;
    }

    build_page_tables(max_physical_address, bss_phys, stack_phys);

    bootinfo().hulk_phys = hulk_bin_phys();
    bootinfo().bss_phys = bss_phys;
    bootinfo().stack_phys = stack_phys;

    jump_to_hulk();

    return EFI_SUCCESS;
}