/**
 * PROJECT:         ExectOS
 * COPYRIGHT:       See COPYING.md in the top level directory
 * FILE:            xtldr/modules/framebuf/framebuf.cc
 * DESCRIPTION:     EFI framebuffer support module for XTLDR
 * DEVELOPERS:      Rafal Kupiec <belliash@codingworkshop.eu.org>
 */

#include <framebuf.hh>


/* PE/COFF_O module information */
MODULE_AUTHOR(L"Rafal Kupiec <belliash@codingworkshop.eu.org>");
MODULE_DESCRIPTION(L"EFI FB (FrameBuffer) support");
MODULE_LICENSE(L"GPLv3");
MODULE_VERSION(L"0.2");


/**
 * Finds a PCI Display Adapter and returns its framebuffer address.
 *
 * @param Address
 *        Supplies a pointer to the memory area where framebuffer address will be stored.
 *
 * @return This routine returns a status code.
 *
 * @since XT 1.0
 */
XTCDECL
EFI_STATUS
FrameBuffer::FindFramebufferAddress(OUT PEFI_PHYSICAL_ADDRESS Address)
{
    EFI_GUID PciIoGuid = EFI_PCI_IO_PROTOCOL_GUID;
    PEFI_ACPI_ADDRESS_SPACE_DESCRIPTOR BarInfo;
    PEFI_PCI_IO_PROTOCOL IoProtocol;
    ULONGLONG FramebufAddressLength;
    PCI_TYPE0_DEVICE PciDevice;
    PVOID FramebufAddress;
    UINT_PTR HandlesCount;
    EFI_HANDLE *Handles;
    EFI_STATUS Status;
    UINT Index;

    /* Initialize variables */
    FramebufAddressLength = 0;
    Handles = NULLPTR;

    /* Locate EFI_PCI_IO_PROTOCOL handles */
    Status = XtLdrProtocol->Protocol.LocateHandles(&Handles, &HandlesCount, &PciIoGuid);
    if(Status != STATUS_EFI_SUCCESS)
    {
        /* Failed to get handles, return error code */
        XtLdrProtocol->Debug.Print(L"ERROR: Failed to get handles (Status Code: 0x%zX)\n", Status);
        return Status;
    }

    /* Iterate through handles */
    for(Index = 0; Index < HandlesCount; Index++)
    {
        /* Open EFI_PCI_IO_PROTOCOL handle */
        Status = XtLdrProtocol->Protocol.OpenHandle(Handles[Index], (PVOID *)&IoProtocol, &PciIoGuid);
        if(Status != STATUS_EFI_SUCCESS)
        {
            /* Failed to open protocol, continue with next handle */
            XtLdrProtocol->Debug.Print(L"ERROR: Failed to open protocol (Status Code: 0x%zX)\n", Status);
            continue;
        }

        /* Read PCI controller registers from PCI configuration space */
        Status = IoProtocol->Pci.Read(IoProtocol, EfiPciIoWidthUint32, 0, sizeof(PciDevice) / sizeof(UINT), &PciDevice);
        if(Status != STATUS_EFI_SUCCESS)
        {
            /* Failed to read PCI device class */
            XtLdrProtocol->Debug.Print(L"ERROR: Failed to read class (Status Code: 0x%zX)\n", Status);

            /* Close protocol and continue with next handle */
            XtLdrProtocol->Protocol.Close(&Handles[Index], &PciIoGuid);
            continue;
        }

        /* Check if device is a graphics adapter */
        if(PciDevice.Hdr.ClassCode[2] != 0x03)
        {
            /* Not a graphics adapter, close protocol and continue with next handle */
            XtLdrProtocol->Protocol.Close(&Handles[Index], &PciIoGuid);
            continue;
        }

        /* Iterate through all PCI device's Base Address Registers (BARs) */
        for(UINT Bars = 0; Bars < 6; Bars++)
        {
            /* Get BAR attributes */
            Status = IoProtocol->GetBarAttributes(IoProtocol, Bars, NULLPTR, (VOID **)&BarInfo);
            if(Status != STATUS_EFI_SUCCESS)
            {
                /* Failed to get BAR attributes, continue with next BAR */
                continue;
            }

            /* Check if this BAR is 'Memory Range' of 'ACPI QWORD Address Space' */
            if(BarInfo->SpaceDescriptor == EFI_ACPI_ADDRESS64_SPACE_DESCRIPTOR &&
               BarInfo->ResourceType == EFI_ACPI_ADDRESS_SPACE_TYPE_MEMORY)
            {
                /* Check if this BAR is the biggest we've seen so far */
                if(BarInfo->AddressLength > FramebufAddressLength)
                {
                    /* The biggest BAR should be the framebuffer; save its address and length */
                    FramebufAddress = (PVOID)(ULONG_PTR)(BarInfo->AddressRangeMin << 16);
                    FramebufAddressLength = BarInfo->AddressLength;
                }
            }
        }

        /* Close handle and continue with next one */
        XtLdrProtocol->Protocol.Close(&Handles[Index], &PciIoGuid);
    }

    /* Set framebuffer address and return success */
    *Address = (EFI_PHYSICAL_ADDRESS)FramebufAddress;
    return STATUS_EFI_SUCCESS;
}

/**
 * Calculates color mask and shift based upon pixel bit mask.
 *
 * @param PixelBitMask
 *        Provides a pixel bit mask.
 *
 * @param ColorSize
 *        Supplies a pointer to the memory area where the color size will be stored.
 *
 * @param ColorShift
 *        Supplies a pointer to the memory area where the color shift (position) will be stored.
 *
 * @return This routine does not return any value.
 *
 * @since XT 1.0
 */
XTCDECL
VOID
FrameBuffer::GetColorMask(IN UINT PixelBitMask,
                          OUT PUSHORT ColorSize,
                          OUT PUSHORT ColorShift)
{
    UINT Shift, Size;

    /* Initialize variables */
    Shift = 0;
    Size = 0;

    /* Make sure EfiMask is not zero */
    if(PixelBitMask)
    {
        /* Get color shift */
        while((PixelBitMask & 1) == 0)
        {
            Shift++;
            PixelBitMask >>= 1;
        }

        /* Get color size */
        while((PixelBitMask & 1) == 1)
        {
            Size++;
            PixelBitMask >>= 1;
        }

    }

    /* Set color mask and shift */
    *ColorShift = Shift;
    *ColorSize = Size;
}

/**
 * Provides an EFI Frame Buffer protocol driver name used for initialization.
 *
 * @param Protocol
 *        Supplies a pointer to the memory area where framebuffer driver information will be stored.
 *
 * @return This routine returns a status code.
 *
 * @since XT 1.0
 */
XTCDECL
EFI_STATUS
FrameBuffer::GetDisplayDriver(OUT PEFI_GRAPHICS_PROTOCOL Protocol)
{
        /* Check if framebuffer is initialized */
    if(!DisplayInfo.Initialized)
    {
        /* Return error if framebuffer is not initialized */
        return STATUS_EFI_NOT_READY;
    }

    /* Copy framebuffer driver information */
    *Protocol = DisplayInfo.Protocol;

    /* Return success */
    return STATUS_EFI_SUCCESS;
}

/**
 * Returns information about EFI Frame Buffer.
 *
 * @param FbInfo
 *        Supplies a pointer to the memory area where framebuffer information will be stored.
 *
 * @return This routine returns a status code.
 *
 * @since XT 1.0
 */
XTCDECL
EFI_STATUS
FrameBuffer::GetDisplayInformation(OUT PEFI_PHYSICAL_ADDRESS FrameBufferBase,
                                   OUT PULONG_PTR FrameBufferSize,
                                   OUT PXTBL_FRAMEBUFFER_MODE_INFORMATION ModeInfo)
{
    /* Check if framebuffer is initialized */
    if(!DisplayInfo.Initialized)
    {
        /* Return error if framebuffer is not initialized */
        return STATUS_EFI_NOT_READY;
    }

    /* Set basic framebuffer information */
    *FrameBufferBase = DisplayInfo.FrameBufferBase;
    *FrameBufferSize = DisplayInfo.FrameBufferSize;

    /* Set framebuffer mode information */
    ModeInfo->Width = DisplayInfo.ModeInfo.Width;
    ModeInfo->Height = DisplayInfo.ModeInfo.Height;
    ModeInfo->Depth = DisplayInfo.ModeInfo.Depth;
    ModeInfo->RefreshRate = DisplayInfo.ModeInfo.RefreshRate;
    ModeInfo->BitsPerPixel = DisplayInfo.ModeInfo.BitsPerPixel;
    ModeInfo->BytesPerPixel = DisplayInfo.ModeInfo.BytesPerPixel;
    ModeInfo->PixelsPerScanLine = DisplayInfo.ModeInfo.PixelsPerScanLine;
    ModeInfo->Pitch = DisplayInfo.ModeInfo.Pitch;
    ModeInfo->PixelFormat = DisplayInfo.ModeInfo.PixelFormat;
    ModeInfo->PixelInformation = DisplayInfo.ModeInfo.PixelInformation;

    /* Return success */
    return STATUS_EFI_SUCCESS;
}


/**
 * Gets information about the current display mode and stores it in internal structure.
 *
 * @return This routine returns a status code.
 *
 * @since XT 1.0
 */
XTCDECL
EFI_STATUS
FrameBuffer::GetModeInformation()
{
    PEFI_GRAPHICS_OUTPUT_MODE_INFORMATION ModeInfo;
    EFI_PIXEL_BITMASK PixelBitMask;
    XTSTATUS Status;
    UINT_PTR Size;

    switch(DisplayInfo.Protocol)
    {
        case GOP:
            /* Query GOP mode information */
            Status = DisplayInfo.Driver.Gop->QueryMode(DisplayInfo.Driver.Gop,
                                                          DisplayInfo.Driver.Gop->Mode->Mode,
                                                          &Size, &ModeInfo);
            if(Status != STATUS_EFI_SUCCESS)
            {
                /* Failed to get GOP mode information, return error */
                return Status;
            }

            /* Get pixel bit mask information */
            GetPixelInformation(&DisplayInfo.Driver.Gop->Mode->Info->PixelInformation);

            /* Store GOP framebuffer information */
            DisplayInfo.ModeInfo.Width = DisplayInfo.Driver.Gop->Mode->Info->HorizontalResolution;
            DisplayInfo.ModeInfo.Height = DisplayInfo.Driver.Gop->Mode->Info->VerticalResolution;
            DisplayInfo.ModeInfo.Depth = DisplayInfo.ModeInfo.BitsPerPixel;
            DisplayInfo.ModeInfo.PixelsPerScanLine = DisplayInfo.Driver.Gop->Mode->Info->PixelsPerScanLine;
            DisplayInfo.ModeInfo.Pitch = DisplayInfo.ModeInfo.PixelsPerScanLine *
                                            (DisplayInfo.ModeInfo.BitsPerPixel / 8);
            DisplayInfo.ModeInfo.RefreshRate = 0;

            /* Store pixel format information and frame buffer size */
            DisplayInfo.ModeInfo.PixelFormat = DisplayInfo.Driver.Gop->Mode->Info->PixelFormat;
            DisplayInfo.FrameBufferSize = DisplayInfo.Driver.Gop->Mode->FrameBufferSize;
            break;
        case UGA:
            /* Query UGA mode information */
            Status = DisplayInfo.Driver.Uga->GetMode(DisplayInfo.Driver.Uga, &DisplayInfo.ModeInfo.Width,
                                                        &DisplayInfo.ModeInfo.Height, &DisplayInfo.ModeInfo.Depth,
                                                        &DisplayInfo.ModeInfo.RefreshRate);
            if(Status != STATUS_EFI_SUCCESS)
            {
                /* Failed to get UGA mode information, return error */
                return Status;
            }

            /* Get pixel bit mask information */
            PixelBitMask = (EFI_PIXEL_BITMASK){0, 0, 0, 0};
            GetPixelInformation(&PixelBitMask);

            /* Store UGA framebuffer information */
            DisplayInfo.ModeInfo.PixelsPerScanLine = DisplayInfo.ModeInfo.Width;
            DisplayInfo.ModeInfo.Pitch = DisplayInfo.ModeInfo.PixelsPerScanLine *
                                            (DisplayInfo.ModeInfo.BitsPerPixel / 8);

            /* Store pixel format information and recalculate frame buffer size */
            DisplayInfo.ModeInfo.PixelFormat = PixelBlueGreenRedReserved8BitPerColor;
            DisplayInfo.FrameBufferSize = DisplayInfo.ModeInfo.Width *
                                             DisplayInfo.ModeInfo.Height *
                                             DisplayInfo.ModeInfo.BytesPerPixel + 1024;
            break;
        default:
            /* This should never be reached as no other display driver is supported */
            break;
    }

    /* Return success */
    return STATUS_EFI_SUCCESS;
}

/** 
 * Gets pixel information based on the reported pixel format.
 *
 * @param FrameBufferInfo
 *        Supplies a pointer to the framebuffer information structure.
 *
 * @param PixelsBitMask
 *        Supplies a pointer to the pixel bit mask information provided by EFI graphics protocol.
 *
 * @return This routine does not return any value.
 *
 * @since XT 1.0
 */
XTCDECL
VOID
FrameBuffer::GetPixelInformation(IN PEFI_PIXEL_BITMASK PixelsBitMask)
{
    UINT CompoundMask;

    /* Check reported pixel format */
    switch(DisplayInfo.ModeInfo.PixelFormat)
    {
        case PixelBlueGreenRedReserved8BitPerColor:
            /* BGRR, 32 bits per pixel */
            DisplayInfo.ModeInfo.BitsPerPixel = 32;
            DisplayInfo.ModeInfo.PixelInformation.BlueShift = 0;
            DisplayInfo.ModeInfo.PixelInformation.BlueSize = 8;
            DisplayInfo.ModeInfo.PixelInformation.GreenShift = 8;
            DisplayInfo.ModeInfo.PixelInformation.GreenSize = 8;
            DisplayInfo.ModeInfo.PixelInformation.RedShift = 16;
            DisplayInfo.ModeInfo.PixelInformation.RedSize = 8;
            DisplayInfo.ModeInfo.PixelInformation.ReservedShift = 24;
            DisplayInfo.ModeInfo.PixelInformation.ReservedSize = 8;
            break;
        case PixelRedGreenBlueReserved8BitPerColor:
            /* RGBR, 32 bits per pixel */
            DisplayInfo.ModeInfo.BitsPerPixel = 32;
            DisplayInfo.ModeInfo.PixelInformation.BlueShift = 16;
            DisplayInfo.ModeInfo.PixelInformation.BlueSize = 8;
            DisplayInfo.ModeInfo.PixelInformation.GreenShift = 8;
            DisplayInfo.ModeInfo.PixelInformation.GreenSize = 8;
            DisplayInfo.ModeInfo.PixelInformation.RedShift = 0;
            DisplayInfo.ModeInfo.PixelInformation.RedSize = 8;
            DisplayInfo.ModeInfo.PixelInformation.ReservedShift = 24;
            DisplayInfo.ModeInfo.PixelInformation.ReservedSize = 8;
            break;
        case PixelBitMask:
            /* Assume 32 bits per pixel */
            DisplayInfo.ModeInfo.BitsPerPixel = 32;

            /* Calculate compound mask */
            CompoundMask = PixelsBitMask->RedMask |
                           PixelsBitMask->GreenMask |
                           PixelsBitMask->BlueMask |
                           PixelsBitMask->ReservedMask;

            /* Recalculate bits per pixel */
            while((CompoundMask & (1 << 31)) == 0)
            {
                DisplayInfo.ModeInfo.BitsPerPixel--;
                CompoundMask <<= 1;
            }

            /* Set pixel information */
            GetColorMask(PixelsBitMask->RedMask, &DisplayInfo.ModeInfo.PixelInformation.RedSize,
                         &DisplayInfo.ModeInfo.PixelInformation.RedShift);
            GetColorMask(PixelsBitMask->GreenMask, &DisplayInfo.ModeInfo.PixelInformation.GreenSize,
                         &DisplayInfo.ModeInfo.PixelInformation.GreenShift);
            GetColorMask(PixelsBitMask->BlueMask, &DisplayInfo.ModeInfo.PixelInformation.BlueSize,
                         &DisplayInfo.ModeInfo.PixelInformation.BlueShift);
            GetColorMask(PixelsBitMask->ReservedMask, &DisplayInfo.ModeInfo.PixelInformation.ReservedSize,
                         &DisplayInfo.ModeInfo.PixelInformation.ReservedShift);
            break;
        default:
            /* Unknown pixel format */
            DisplayInfo.ModeInfo.BitsPerPixel = 0;
            DisplayInfo.ModeInfo.PixelInformation.BlueShift = 0;
            DisplayInfo.ModeInfo.PixelInformation.BlueSize = 0;
            DisplayInfo.ModeInfo.PixelInformation.GreenShift = 0;
            DisplayInfo.ModeInfo.PixelInformation.GreenSize = 0;
            DisplayInfo.ModeInfo.PixelInformation.RedShift = 0;
            DisplayInfo.ModeInfo.PixelInformation.RedSize = 0;
            DisplayInfo.ModeInfo.PixelInformation.ReservedShift = 0;
            DisplayInfo.ModeInfo.PixelInformation.ReservedSize = 0;
            break;
    }

    /* Calculate bytes per pixel based on bits per pixel */
    DisplayInfo.ModeInfo.BytesPerPixel = DisplayInfo.ModeInfo.BitsPerPixel >> 3;
}

/**
 * Determines the preferred (native) screen resolution from EDID. This works only with GOP.
 *
 * @param PreferredWidth
 *        Supplies a pointer to the memory area where preferred screen width will be stored.
 *
 * @param PreferredHeight
 *        Supplies a pointer to the memory area where preferred screen height will be stored.
 *
 * @return This routine returns a status code.
 *
 * @since XT 1.0
 */
XTCDECL
EFI_STATUS
FrameBuffer::GetPreferredScreenResolution(OUT PUINT PreferredWidth,
                                          OUT PUINT PreferredHeight)
{
    EFI_GUID GopGuid = EFI_GRAPHICS_OUTPUT_PROTOCOL_GUID;
    EFI_GUID EdidGuid = EFI_EDID_ACTIVE_PROTOCOL_GUID;
    PEFI_EDID_ACTIVE_PROTOCOL ActiveEdid;
    EFI_STATUS Status;

    /* Check if framebuffer is initialized */
    if(!DisplayInfo.Initialized)
    {
        /* Framebuffer not ready to use EDID protocol */
        return STATUS_EFI_NOT_READY;
    }

    /* Check if GOP device driver is used */
    if(DisplayInfo.Protocol != GOP)
    {
        /* Unsupported device driver */
        return STATUS_EFI_UNSUPPORTED;
    }

    /* Open EDID protocol */
    Status = XtLdrProtocol->Protocol.OpenHandle(DisplayInfo.Handle, (PVOID *)&ActiveEdid, &EdidGuid);
    if(Status != STATUS_EFI_SUCCESS)
    {
        /* Failed to open EDID protocol, close GOP protocol and return */
        XtLdrProtocol->Protocol.Close(&DisplayInfo.Handle, &GopGuid);
        return Status;
    }

    /* Return preferred screen resolution */
    *PreferredWidth = ActiveEdid->Edid[0x38] | ((ActiveEdid->Edid[0x3A] & 0xF0) << 4);
    *PreferredHeight = ActiveEdid->Edid[0x3B] | ((ActiveEdid->Edid[0x3D] & 0xF0) << 4);

    /* Close EDID & GOP protocols */
    XtLdrProtocol->Protocol.Close(&DisplayInfo.Handle, &EdidGuid);

    /* Return success */
    return STATUS_EFI_SUCCESS;
}

/**
 * Initializes FrameBuffer device on GOP and UGA compatible adapters.
 *
 * @return This routine returns a status code.
 *
 * @since XT 1.0
 */
XTCDECL
EFI_STATUS
FrameBuffer::InitializeDisplay()
{
    EFI_GUID GopGuid = EFI_GRAPHICS_OUTPUT_PROTOCOL_GUID;
    EFI_GUID UgaGuid = EFI_UNIVERSAL_GRAPHICS_ADAPTER_PROTOCOL_GUID;
    PEFI_GRAPHICS_OUTPUT_MODE_INFORMATION GopModeInfo;
    UINT Depth, QueryMode, Refresh;
    UINT_PTR InfoSize;
    EFI_STATUS Status;

    /* Check if framebuffer already initialized */
    if(!DisplayInfo.Initialized)
    {
        /* Print debug message */
        XtLdrProtocol->Debug.Print(L"Initializing framebuffer device\n");

        /* Attempt to open EFI GOP protocol */
        Status = XtLdrProtocol->Protocol.Open(&DisplayInfo.Handle, (PVOID*)&DisplayInfo.Driver.Gop, &GopGuid);

        /* Check if Graphics Output Protocol (GOP) is available */
        if(Status == STATUS_EFI_SUCCESS)
        {
            /* Check if there are any video modes available */
            if(DisplayInfo.Driver.Gop->Mode->MaxMode == 0)
            {
                /* No video modes available */
                XtLdrProtocol->Debug.Print(L"ERROR: No GOP video mode available\n");

                /* Close GOP protocol and return error */
                XtLdrProtocol->Protocol.Close(&DisplayInfo.Handle, &GopGuid);
                return STATUS_EFI_UNSUPPORTED;
            }

            /* Query current graphics mode */
            QueryMode = DisplayInfo.Driver.Gop->Mode == NULLPTR ? 0 : DisplayInfo.Driver.Gop->Mode->Mode;
            Status = DisplayInfo.Driver.Gop->QueryMode(DisplayInfo.Driver.Gop, QueryMode, &InfoSize, &GopModeInfo);
            if(Status == STATUS_EFI_NOT_STARTED)
            {
                /* Set the mode to circumvent buggy UEFI firmware */
                Status = DisplayInfo.Driver.Gop->SetMode(DisplayInfo.Driver.Gop, 0);
            }
            if(Status != STATUS_EFI_SUCCESS)
            {
                /* Unable to query GOP modes */
                XtLdrProtocol->Debug.Print(L"ERROR: Failed to get GOP native mode (Status Code: 0x%zX)\n");

                /* Close GOP protocol and return error */
                XtLdrProtocol->Protocol.Close(&DisplayInfo.Handle, &GopGuid);
                return STATUS_EFI_UNSUPPORTED;
            }

            /* Store frame buffer base address and protocol used */
            DisplayInfo.FrameBufferBase = DisplayInfo.Driver.Gop->Mode->FrameBufferBase;
            DisplayInfo.DefaultMode = DisplayInfo.Driver.Gop->Mode->Mode;
            DisplayInfo.Protocol = GOP;

            /* Get current mode information */
            Status = GetModeInformation();
            if(Status != STATUS_EFI_SUCCESS)
            {
                /* Unable to get mode information */
                XtLdrProtocol->Debug.Print(L"ERROR: Failed to get GOP mode information (Status Code: 0x%zX)\n");

                /* Close GOP protocol and return error */
                XtLdrProtocol->Protocol.Close(&DisplayInfo.Handle, &GopGuid);
                return STATUS_EFI_UNSUPPORTED;
            }

            /* Found GOP */
            XtLdrProtocol->Debug.Print(L"Found EFI-GOP compatible display adapter @ %P (%zu bytes)\n",
                                       DisplayInfo.FrameBufferBase, DisplayInfo.FrameBufferSize);

            /* Close GOP protocol */
            Status = XtLdrProtocol->Protocol.Close(&DisplayInfo.Handle, &GopGuid);
        }
        else
        {
            /* GOP is unavailable, attempt to open UGA protocol */
            Status = XtLdrProtocol->Protocol.Open(&DisplayInfo.Handle, (PVOID*)&DisplayInfo.Driver.Uga, &UgaGuid);

            /* Check if Universal Graphics Adapter (UGA) is available */
            if(Status == STATUS_EFI_SUCCESS)
            {
                /* Get current video mode */
                Status = DisplayInfo.Driver.Uga->GetMode(DisplayInfo.Driver.Uga, &DisplayInfo.ModeInfo.Width,
                                                            &DisplayInfo.ModeInfo.Height, &Depth, &Refresh);
                if(Status != STATUS_EFI_SUCCESS)
                {
                    /* Unable to get current UGA mode */
                    XtLdrProtocol->Debug.Print(L"ERROR: Failed to get current UGA mode (Status Code: 0x%zX)\n", Status);

                    /* Close UGA protocol and return error */
                    XtLdrProtocol->Protocol.Close(&DisplayInfo.Handle, &UgaGuid);
                    return STATUS_EFI_DEVICE_ERROR;
                }

                /* Find framebuffer address */
                Status = FindFramebufferAddress(&DisplayInfo.FrameBufferBase);
                if(Status != STATUS_EFI_SUCCESS)
                {
                    /* Unable to find framebuffer address */
                    XtLdrProtocol->Debug.Print(L"ERROR: Failed to get EFI FB address (Status Code: 0x%zX)\n", Status);

                    /* Close UGA protocol and return error */
                    XtLdrProtocol->Protocol.Close(&DisplayInfo.Handle, &UgaGuid);
                    return STATUS_EFI_DEVICE_ERROR;
                }

                /* Store framebuffer protocol information */
                DisplayInfo.DefaultMode = 0;
                DisplayInfo.Protocol = UGA;

                /* Get mode information */
                Status = GetModeInformation();
                if(Status != STATUS_EFI_SUCCESS)
                {
                    /* Unable to get mode information */
                    XtLdrProtocol->Debug.Print(L"ERROR: Failed to get UGA mode information (Status Code: 0x%zX)\n");
                    return STATUS_EFI_UNSUPPORTED;
                }

                /* Found UGA */
                XtLdrProtocol->Debug.Print(L"Found EFI-UGA compatible display adapter @ %P (%zu bytes)\n",
                                           DisplayInfo.FrameBufferBase, DisplayInfo.FrameBufferSize);

                /* Close UGA protocol */
                XtLdrProtocol->Protocol.Close(&DisplayInfo.Handle, &UgaGuid);
            }
        }

        /* Make sure framebuffer initialized properly */
        if(DisplayInfo.Protocol == NONE)
        {
            /* GOP and UGA unavailable */
            XtLdrProtocol->Debug.Print(L"WARNING: No display adapter found!\n");
            return STATUS_EFI_NOT_FOUND;
        }

        XtLdrProtocol->Debug.Print(L"Current screen resolution is %ux%ux%u\n", DisplayInfo.ModeInfo.Width,
                                   DisplayInfo.ModeInfo.Height, DisplayInfo.ModeInfo.BitsPerPixel);

        /* Set framebuffer initialization flag */
        DisplayInfo.Initialized = TRUE;
    }

    /* Return success */
    return STATUS_SUCCESS;
}

/**
 * Initializes FRAMEBUF module by opening XTLDR protocol and installing FRAMEBUF protocol.
 *
 * @param ImageHandle
 *        Firmware-allocated handle that identifies the image.
 *
 * @param SystemTable
 *        Provides the EFI system table.
 *
 * @return This routine returns status code.
 *
 * @since XT 1.0
 */
XTCDECL
EFI_STATUS
FrameBuffer::InitializeModule(IN EFI_HANDLE ImageHandle,
                              IN PEFI_SYSTEM_TABLE SystemTable)
{
    EFI_GUID Guid = XT_FRAMEBUFFER_PROTOCOL_GUID;
    EFI_STATUS Status;

    /* Open the XTLDR protocol */
    Status = BlGetXtLdrProtocol(SystemTable, ImageHandle, &XtLdrProtocol);
    if(Status != STATUS_EFI_SUCCESS)
    {
        /* Failed to open loader protocol */
        return STATUS_EFI_PROTOCOL_ERROR;
    }

    /* Set initial framebuffer state */
    DisplayInfo.Protocol = NONE;
    DisplayInfo.Initialized = FALSE;

    /* Set routines available via XTLDR framebuffer protocol */
    FbProtocol.GetDisplayDriver = GetDisplayDriver;
    FbProtocol.GetDisplayInformation = GetDisplayInformation;
    FbProtocol.GetPreferredScreenResolution = GetPreferredScreenResolution;
    FbProtocol.Initialize = InitializeDisplay;
    FbProtocol.SetScreenResolution = SetScreenResolution;

    /* Register XTOS boot protocol */
    return XtLdrProtocol->Protocol.Install(&FbProtocol, &Guid);
}

/**
 * Sets custom screen resolution, based on the provided width and height.
 *
 * @param Width
 *        Supplies the width of the screen.
 *
 * @param Height
 *        Supplies the height of the screen.
 *
 * @return This routine returns a status code.
 *
 * @since XT 1.0
 */
XTCDECL
EFI_STATUS
FrameBuffer::SetScreenResolution(IN UINT Width,
                                 IN UINT Height)
{
    PEFI_GRAPHICS_OUTPUT_MODE_INFORMATION ModeInfo;
    BOOLEAN ModeChanged;
    EFI_STATUS Status;
    UINT_PTR Size;
    UINT Mode;

    /* Check if framebuffer is initialized */
    if(!DisplayInfo.Initialized)
    {
        /* Framebuffer not ready to change screen mode */
        return STATUS_EFI_NOT_READY;
    }

    ModeChanged = FALSE;

    /* Change screen mode depending on display adapter protocol */
    switch(DisplayInfo.Protocol)
    {
        case GOP:
            /* GOP available, check if user specified screen resolution */
            if(Width == 0 || Height == 0)
            {
                /* No resolution specified, temporarily set lowest supported screen resolution */
                Status = DisplayInfo.Driver.Gop->SetMode(DisplayInfo.Driver.Gop, 1);
                if(Status == STATUS_EFI_SUCCESS)
                {
                    /* Restore default graphics mode */
                    Status = DisplayInfo.Driver.Gop->SetMode(DisplayInfo.Driver.Gop, DisplayInfo.DefaultMode);
                    ModeChanged = (Status == STATUS_EFI_SUCCESS) ? TRUE : FALSE;
                }
            }
            else
            {
                /* User specified screen resolution, find a corresponding mode */
                Mode = 1;
                while(Mode <= DisplayInfo.Driver.Gop->Mode->MaxMode)
                {
                    /* Get mode information */
                    Status = DisplayInfo.Driver.Gop->QueryMode(DisplayInfo.Driver.Gop, Mode, &Size, &ModeInfo);
                    if(Status == STATUS_EFI_SUCCESS && Size >= sizeof(*ModeInfo) && ModeInfo != NULLPTR)
                    {
                        /* Check if match found */
                        if(ModeInfo->HorizontalResolution == Width && ModeInfo->VerticalResolution == Height)
                        {
                            /* Found corresponding mode, attempt to set it */
                            Status = DisplayInfo.Driver.Gop->SetMode(DisplayInfo.Driver.Gop, Mode);
                            if(Status == STATUS_EFI_SUCCESS)
                            {
                                /* New mode set correctly, use it */
                                ModeChanged = TRUE;
                                break;
                            }
                        }
                    }

                    /* Try with next mode */
                    Mode++;
                }
            }
            break;
        case UGA:
            /* Set UGA screen mode, trying to keep current color depth and refresh rate */
            Status = DisplayInfo.Driver.Uga->SetMode(DisplayInfo.Driver.Uga, Width, Height,
                                                        DisplayInfo.ModeInfo.Depth,
                                                        DisplayInfo.ModeInfo.RefreshRate);
            if(Status == STATUS_EFI_SUCCESS)
            {
                /* New mode set correctly, use it */
                ModeChanged = TRUE;
            }
            break;
        default:
            /* This should never be reached */
            break;
    }

    if(!ModeChanged)
    {
        /* Failed to change screen mode */
        XtLdrProtocol->Debug.Print(L"ERROR: Failed to change screen mode to %ux%u (Status Code: 0x%zX)\n",
                                   Width, Height, Status);
        return STATUS_EFI_UNSUPPORTED;
    }

    /* Get new screen mode information */
    Status = GetModeInformation();
    if(Status == STATUS_EFI_SUCCESS)
    {
        XtLdrProtocol->Debug.Print(L"Changed screen resolution to %ux%ux%u\n", DisplayInfo.ModeInfo.Width,
                                   DisplayInfo.ModeInfo.Height, DisplayInfo.ModeInfo.BitsPerPixel);
    }

    /* Return success */
    return STATUS_EFI_SUCCESS;
}

/**
 * This routine is the entry point of the XT EFI boot loader module.
 *
 * @param ImageHandle
 *        Firmware-allocated handle that identifies the image.
 *
 * @param SystemTable
 *        Provides the EFI system table.
 *
 * @return This routine returns status code.
 *
 * @since XT 1.0
 */
XTCDECL
EFI_STATUS
XtLdrModuleMain(IN EFI_HANDLE ImageHandle,
                IN PEFI_SYSTEM_TABLE SystemTable)
{
    /* Initialize FRAMEBUF module */
    return FrameBuffer::InitializeModule(ImageHandle, SystemTable);
}
