import uefi;
import output;
import scratch;
import hulk.bootinfo;
import hos.page_table;
import hos.cpu;
import hos.memory;

__gshared EFI_SYSTEM_TABLE * st;
__gshared BootInfo bootinfo;
__gshared UINTN memory_map_key;
extern extern(C) __gshared ubyte hulk_start;
extern extern(C) __gshared ubyte hulk_end;

private void wait_key()
{
    writeln("Press any key...");
    st.ConIn.Reset(st.ConIn, FALSE);
    EFI_INPUT_KEY key;
    while (st.ConIn.ReadKeyStroke(st.ConIn, &key) == EFI_NOT_READY)
    {
    }
}

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

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)
    {
        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)
    {
        writeln("HandleProtocol: error %x", status);
        return false;
    }
    if (gop_interface == null)
    {
        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)
    {
        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;
}

private ulong get_memory_map()
{
    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)
    {
        writeln("GetMemoryMap: Error %x", status);
        return 0u;
    }
    ulong max_physical_address;
    size_t n_entries = memory_map_size / descriptor_size;
    size_t count;
    for (count = 0u; count < n_entries; count++)
    {
        if (count > bootinfo.memory_map.length)
        {
            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;
        }
        bootinfo.memory_map[count].base = descriptor.PhysicalStart;
        bootinfo.memory_map[count].size = descriptor.NumberOfPages * 4096u;
        bootinfo.memory_map[count].type = cast(ubyte)descriptor.Type;
    }
    bootinfo.memory_map_count = count;
    return max_physical_address;
}

private PageTableEntry * new_page_table()
{
    PageTableEntry * pt = cast(PageTableEntry *)scratch_alloc(1u);
    memset64(pt, 0u, 512);
    return pt;
}

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;
        }
    }
}

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;
        }
    }
}

private void map_hulk(PageTableEntry * pt_base)
{
    ulong virt = HULK_VIRTUAL_START;
    ulong phys = cast(ulong)&hulk_start;
    while (phys < cast(ulong)&hulk_end)
    {
        map4k(virt, phys, pt_base);
        virt += 4096u;
        phys += 4096u;
    }
}

private void build_page_tables(ulong max_physical_address)
{
    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++)
    {
        if (bootinfo.memory_map[i].base > max_physical_address)
        {
            for (size_t offset = 0u; offset < bootinfo.memory_map[i].size; offset += 4096u)
            {
                ulong addr = bootinfo.memory_map[i].base + offset;
                map4k(addr, addr, pt_base);
            }
        }
    }
    /* Map graphics framebuffer. */
    ulong framebuffer_size = bootinfo.fb.height * bootinfo.fb.stride;
    for (size_t offset = 0u; offset < framebuffer_size; offset++)
    {
        ulong addr = cast(ulong)bootinfo.fb.buffer + offset;
        map4k(addr, addr, pt_base);
    }
    /* Map HULK to its requested starting virtual address. */
    map_hulk(pt_base);
    /* Switch to the new page table. */
    write_cr3(cast(ulong)pt_base);
}

private void jump_to_hulk()
{
    void function(BootInfo *) hulk_start = cast(void function(BootInfo *))HULK_VIRTUAL_START;
    hulk_start(&bootinfo);
}

extern (C) EFI_STATUS efi_main(EFI_HANDLE image_handle, EFI_SYSTEM_TABLE * st)
{
    .st = st;

    st.ConOut.ClearScreen(st.ConOut);

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

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

    ulong max_physical_address;
    for (;;)
    {
        max_physical_address = get_memory_map();

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

    build_page_tables(max_physical_address);

    jump_to_hulk();

    return EFI_SUCCESS;
}