/** @file
  Copyright (C) 2020, vit9696. All rights reserved.

  All rights reserved.

  This program and the accompanying materials
  are licensed and made available under the terms and conditions of the BSD License
  which accompanies this distribution.  The full text of the license may be found at
  http://opensource.org/licenses/bsd-license.php

  THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
  WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
**/

#include "OcConsoleLibInternal.h"

#include <Library/DebugLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/UefiBootServicesTableLib.h>

#include <Library/OcMiscLib.h>

STATIC
EFI_STATUS
EFIAPI
OcUgaDrawGetMode (
  IN  EFI_UGA_DRAW_PROTOCOL  *This,
  OUT UINT32                 *HorizontalResolution,
  OUT UINT32                 *VerticalResolution,
  OUT UINT32                 *ColorDepth,
  OUT UINT32                 *RefreshRate
  )
{
  OC_UGA_PROTOCOL                       *OcUgaDraw;
  EFI_GRAPHICS_OUTPUT_MODE_INFORMATION  *Info;

  DEBUG ((DEBUG_INFO, "OCC: OcUgaDrawGetMode %p\n", This));

  if (  (This == NULL)
     || (HorizontalResolution == NULL)
     || (VerticalResolution == NULL)
     || (ColorDepth == NULL)
     || (RefreshRate == NULL))
  {
    return EFI_INVALID_PARAMETER;
  }

  OcUgaDraw = BASE_CR (This, OC_UGA_PROTOCOL, Uga);
  Info      = OcUgaDraw->GraphicsOutput->Mode->Info;

  DEBUG ((
    DEBUG_INFO,
    "OCC: OcUgaDrawGetMode Info is %ux%u (%u)\n",
    Info->HorizontalResolution,
    Info->VerticalResolution,
    Info->PixelFormat
    ));

  if ((Info->HorizontalResolution == 0) || (Info->VerticalResolution == 0)) {
    return EFI_NOT_STARTED;
  }

  *HorizontalResolution = Info->HorizontalResolution;
  *VerticalResolution   = Info->VerticalResolution;
  *ColorDepth           = DEFAULT_COLOUR_DEPTH;
  *RefreshRate          = DEFAULT_REFRESH_RATE;

  return EFI_SUCCESS;
}

STATIC
EFI_STATUS
EFIAPI
OcUgaDrawSetMode (
  IN  EFI_UGA_DRAW_PROTOCOL  *This,
  IN  UINT32                 HorizontalResolution,
  IN  UINT32                 VerticalResolution,
  IN  UINT32                 ColorDepth,
  IN  UINT32                 RefreshRate
  )
{
  EFI_STATUS       Status;
  OC_UGA_PROTOCOL  *OcUgaDraw;

  DEBUG ((
    DEBUG_INFO,
    "OCC: OcUgaDrawSetMode %p %ux%u@%u/%u\n",
    This,
    HorizontalResolution,
    VerticalResolution,
    ColorDepth,
    RefreshRate
    ));

  OcUgaDraw = BASE_CR (This, OC_UGA_PROTOCOL, Uga);

  Status = OcSetConsoleResolutionForProtocol (
             OcUgaDraw->GraphicsOutput,
             HorizontalResolution,
             VerticalResolution,
             ColorDepth
             );

  DEBUG ((DEBUG_INFO, "OCC: UGA SetConsoleResolutionOnHandle attempt - %r\n", Status));

  if (EFI_ERROR (Status)) {
    Status = OcSetConsoleResolutionForProtocol (
               OcUgaDraw->GraphicsOutput,
               0,
               0,
               0
               );
    DEBUG ((DEBUG_INFO, "OCC: UGA SetConsoleResolutionOnHandle max attempt - %r\n", Status));
  }

  return Status;
}

STATIC
EFI_STATUS
EFIAPI
OcUgaDrawBlt (
  IN  EFI_UGA_DRAW_PROTOCOL  *This,
  IN  EFI_UGA_PIXEL          *BltBuffer OPTIONAL,
  IN  EFI_UGA_BLT_OPERATION  BltOperation,
  IN  UINTN                  SourceX,
  IN  UINTN                  SourceY,
  IN  UINTN                  DestinationX,
  IN  UINTN                  DestinationY,
  IN  UINTN                  Width,
  IN  UINTN                  Height,
  IN  UINTN                  Delta      OPTIONAL
  )
{
  OC_UGA_PROTOCOL  *OcUgaDraw;

  OcUgaDraw = BASE_CR (This, OC_UGA_PROTOCOL, Uga);

  return OcUgaDraw->GraphicsOutput->Blt (
                                      OcUgaDraw->GraphicsOutput,
                                      (EFI_GRAPHICS_OUTPUT_BLT_PIXEL *)BltBuffer,
                                      (EFI_GRAPHICS_OUTPUT_BLT_OPERATION)BltOperation,
                                      SourceX,
                                      SourceY,
                                      DestinationX,
                                      DestinationY,
                                      Width,
                                      Height,
                                      Delta
                                      );
}

EFI_STATUS
OcProvideUgaPassThrough (
  VOID
  )
{
  EFI_STATUS                    Status;
  UINTN                         HandleCount;
  EFI_HANDLE                    *HandleBuffer;
  UINTN                         Index;
  EFI_GRAPHICS_OUTPUT_PROTOCOL  *GraphicsOutput;
  EFI_UGA_DRAW_PROTOCOL         *UgaDraw;
  OC_UGA_PROTOCOL               *OcUgaDraw;

  //
  // For now we do not need this but for launching 10.4, but as a side note:
  // MacPro5,1 has 2 GOP protocols:
  // - for GPU
  // - for ConsoleOutput
  // and 1 UGA protocol:
  // - for unknown handle
  //
  DEBUG_CODE_BEGIN ();
  DEBUG ((DEBUG_INFO, "OCC: Found %u handles with UGA draw\n", (UINT32)OcCountProtocolInstances (&gEfiUgaDrawProtocolGuid)));
  DEBUG_CODE_END ();

  Status = gBS->LocateHandleBuffer (
                  ByProtocol,
                  &gEfiGraphicsOutputProtocolGuid,
                  NULL,
                  &HandleCount,
                  &HandleBuffer
                  );
  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_INFO, "OCC: Failed to find handles with GOP - %r\n", Status));
    return Status;
  }

  DEBUG ((DEBUG_INFO, "OCC: Found %u handles with GOP for UGA check\n", (UINT32)HandleCount));
  for (Index = 0; Index < HandleCount; ++Index) {
    DEBUG ((DEBUG_INFO, "OCC: Trying handle %u - %p\n", (UINT32)Index, HandleBuffer[Index]));

    Status = gBS->HandleProtocol (
                    HandleBuffer[Index],
                    &gEfiGraphicsOutputProtocolGuid,
                    (VOID **)&GraphicsOutput
                    );
    if (EFI_ERROR (Status)) {
      DEBUG ((DEBUG_INFO, "OCC: No GOP protocol - %r\n", Status));
      continue;
    }

    Status = gBS->HandleProtocol (
                    HandleBuffer[Index],
                    &gEfiUgaDrawProtocolGuid,
                    (VOID **)&UgaDraw
                    );
    if (!EFI_ERROR (Status)) {
      DEBUG ((DEBUG_INFO, "OCC: Skipping UGA proxying as it is already present on handle %u - %p\n", (UINT32)Index, HandleBuffer[Index]));
      continue;
    }

    OcUgaDraw = AllocateZeroPool (sizeof (*OcUgaDraw));
    if (OcUgaDraw == NULL) {
      DEBUG ((DEBUG_INFO, "OCC: Failed to allocate UGA protocol\n"));
      continue;
    }

    OcUgaDraw->GraphicsOutput = GraphicsOutput;
    OcUgaDraw->Uga.GetMode    = OcUgaDrawGetMode;
    OcUgaDraw->Uga.SetMode    = OcUgaDrawSetMode;
    OcUgaDraw->Uga.Blt        = OcUgaDrawBlt;

    Status = gBS->InstallMultipleProtocolInterfaces (
                    &HandleBuffer[Index],
                    &gEfiUgaDrawProtocolGuid,
                    &OcUgaDraw->Uga,
                    NULL
                    );
    if (EFI_ERROR (Status)) {
      FreePool (OcUgaDraw);
    }

    DEBUG ((
      DEBUG_INFO,
      "OCC: Installed UGA protocol - %r (Handle %u - %p)\n",
      Status,
      (UINT32)Index,
      HandleBuffer[Index]
      ));
  }

  FreePool (HandleBuffer);

  return Status;
}
