/** @file
  Copyright (C) 2016 - 2017, The HermitCrabs Lab. 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 <Uefi.h>

#include <IndustryStandard/Pci.h>

#include <Protocol/PciIo.h>

#include <Library/BaseMemoryLib.h>
#include <Library/DebugLib.h>
#include <Library/IoLib.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/OcDeviceMiscLib.h>

#include "PciExtInternal.h"

/**
  Release XHCI USB controllers ownership.

  @param[in]  PciIo  PCI I/O protocol for the device.

  @retval EFI_NOT_FOUND  No XHCI controllers had ownership incorrectly set.
  @retval EFI_SUCCESS    Corrected XHCI controllers ownership to OS.
**/
STATIC
EFI_STATUS
XhciReleaseOwnership (
  IN  EFI_PCI_IO_PROTOCOL  *PciIo
  )
{
  EFI_STATUS  Status;

  UINT32  HcCapParams;
  UINT32  ExtendCap;
  UINT32  Value;
  INT32   TimeOut;

  //
  // XHCI controller, then disable legacy support, if enabled.
  //

  Status = PciIo->Mem.Read (
                        PciIo,
                        EfiPciIoWidthUint32,
                        XHC_USBCMD_OFFSET,
                        XHC_HCCPARAMS_OFFSET,
                        1,
                        &HcCapParams
                        );

  ExtendCap = EFI_ERROR (Status) ? 0 : ((HcCapParams >> 14U) & 0x3FFFCU);

  while (ExtendCap) {
    Status = PciIo->Mem.Read (
                          PciIo,
                          EfiPciIoWidthUint32,
                          XHC_USBCMD_OFFSET,
                          ExtendCap,
                          1,
                          &Value
                          );

    if (EFI_ERROR (Status)) {
      break;
    }

    if ((Value & XHC_CAPABILITY_ID_MASK) == 1) {
      //
      // Do nothing if BIOS ownership is cleared.
      //
      if (!(Value & BIT16)) {
        break;
      }

      Value |= BIT24;

      PciIo->Mem.Write (
                   PciIo,
                   EfiPciIoWidthUint32,
                   XHC_USBCMD_OFFSET,
                   ExtendCap,
                   1,
                   &Value
                   );

      TimeOut = 40;
      while (TimeOut--) {
        gBS->Stall (500);

        Status = PciIo->Mem.Read (
                              PciIo,
                              EfiPciIoWidthUint32,
                              XHC_USBCMD_OFFSET,
                              ExtendCap,
                              1,
                              &Value
                              );

        if (EFI_ERROR (Status) || !(Value & BIT16)) {
          break;
        }
      }

      //
      // Disable all SMI in USBLEGCTLSTS
      //
      Status = PciIo->Mem.Read (
                            PciIo,
                            EfiPciIoWidthUint32,
                            XHC_USBCMD_OFFSET,
                            ExtendCap + 4,
                            1,
                            &Value
                            );

      if (EFI_ERROR (Status)) {
        break;
      }

      Value &= 0x1F1FEEU;
      Value |= 0xE0000000U;

      PciIo->Mem.Write (
                   PciIo,
                   EfiPciIoWidthUint32,
                   XHC_USBCMD_OFFSET,
                   ExtendCap + 4,
                   1,
                   &Value
                   );

      //
      // Clear all ownership
      //
      Status = PciIo->Mem.Read (
                            PciIo,
                            EfiPciIoWidthUint32,
                            XHC_USBCMD_OFFSET,
                            ExtendCap,
                            1,
                            &Value
                            );

      if (EFI_ERROR (Status)) {
        break;
      }

      Value &= ~(BIT24 | BIT16);
      PciIo->Mem.Write (
                   PciIo,
                   EfiPciIoWidthUint32,
                   XHC_USBCMD_OFFSET,
                   ExtendCap,
                   1,
                   &Value
                   );

      break;
    }

    if (!(Value & XHC_NEXT_CAPABILITY_MASK)) {
      break;
    }

    ExtendCap += ((Value >> 6U) & 0x3FCU);
  }

  return Status;
}

/**
 Release EHCI USB controllers ownership.

 @param[in]  PciIo  PCI I/O protocol for the device.

 @retval EFI_NOT_FOUND  No EHCI controllers had ownership incorrectly set.
 @retval EFI_SUCCESS    Corrected EHCI controllers ownership to OS.
 **/
STATIC
EFI_STATUS
EhciReleaseOwnership (
  IN  EFI_PCI_IO_PROTOCOL  *PciIo
  )
{
  EFI_STATUS  Status;
  UINT32      Value;
  UINT32      Base;
  UINT32      OpAddr;
  UINT32      ExtendCap;
  UINT32      UsbCmd;
  UINT32      UsbLegSup;
  UINT32      UsbLegCtlSts;
  UINTN       IsOsOwned;
  UINTN       IsBiosOwned;
  BOOLEAN     IsOwnershipConflict;
  UINT32      HcCapParams;
  INT32       TimeOut;

  Value  = 0x0002;
  Status = EFI_SUCCESS;

  PciIo->Pci.Write (
               PciIo,
               EfiPciIoWidthUint16,
               0x04,
               1,
               &Value
               );

  Base = 0;

  PciIo->Pci.Read (
               PciIo,
               EfiPciIoWidthUint32,
               0x10,
               1,
               &Base
               );

  if (MmioRead8 (Base) < 0x0C) {
    //
    // Config space too small: no legacy implementation.
    //
    return EFI_NOT_FOUND;
  }

  //
  // Operational Registers = capaddr + offset (8bit CAPLENGTH in Capability Registers + offset 0).
  //
  OpAddr = Base + MmioRead8 (Base);

  PciIo->Mem.Read (
               PciIo,
               EfiPciIoWidthUint32,
               EHC_BAR_INDEX,
               EHC_HCCPARAMS_OFFSET,
               1,
               &HcCapParams
               );

  ExtendCap = (HcCapParams >> 8U) & 0xFFU;

  //
  // Read PCI Config 32bit USBLEGSUP (eecp+0).
  //
  PciIo->Pci.Read (
               PciIo,
               EfiPciIoWidthUint32,
               ExtendCap,
               1,
               &UsbLegSup
               );

  IsBiosOwned = (UsbLegSup & BIT16) != 0;
  if (!IsBiosOwned) {
    //
    // No BIOS ownership, ignore.
    //
    return EFI_NOT_FOUND;
  }

  //
  // Read PCI Config 32bit USBLEGCTLSTS (eecp+4).
  //
  PciIo->Pci.Read (
               PciIo,
               EfiPciIoWidthUint32,
               ExtendCap + 0x4,
               1,
               &UsbLegCtlSts
               );

  //
  // Disable the SMI in USBLEGCTLSTS firstly.
  //
  UsbLegCtlSts &= 0xFFFF0000U;
  PciIo->Pci.Write (
               PciIo,
               EfiPciIoWidthUint32,
               ExtendCap + 0x4,
               1,
               &UsbLegCtlSts
               );

  UsbCmd = MmioRead32 (OpAddr + EHC_USBCMD_OFFSET);

  //
  // Clear registers to default.
  //
  UsbCmd = UsbCmd & 0xFFFFFF00U;
  MmioWrite32 (OpAddr + EHC_USBCMD_OFFSET, UsbCmd);
  MmioWrite32 (OpAddr + EHC_USBINT_OFFSET, 0);
  MmioWrite32 (OpAddr + EHC_USBSTS_OFFSET, 0x1000);

  Value = 1;
  PciIo->Pci.Write (
               PciIo,
               EfiPciIoWidthUint32,
               ExtendCap,
               1,
               &Value
               );

  //
  // Read 32bit USBLEGSUP (eecp+0).
  //
  PciIo->Pci.Read (
               PciIo,
               EfiPciIoWidthUint32,
               ExtendCap,
               1,
               &UsbLegSup
               );

  IsBiosOwned = (UsbLegSup & BIT16) != 0;
  IsOsOwned   = (UsbLegSup & BIT24) != 0;

  //
  // Read 32bit USBLEGCTLSTS (eecp+4).
  //
  PciIo->Pci.Read (
               PciIo,
               EfiPciIoWidthUint32,
               ExtendCap + 0x4,
               1,
               &UsbLegCtlSts
               );

  //
  // Get EHCI Ownership from legacy bios.
  //
  PciIo->Pci.Read (
               PciIo,
               EfiPciIoWidthUint32,
               ExtendCap,
               1,
               &UsbLegSup
               );

  IsOwnershipConflict = IsBiosOwned && IsOsOwned;

  if (IsOwnershipConflict) {
    //
    // EHCI - Ownership conflict - attempting soft reset.
    //
    Value = 0;
    PciIo->Pci.Write (
                 PciIo,
                 EfiPciIoWidthUint8,
                 ExtendCap + 3,
                 1,
                 &Value
                 );

    TimeOut = 40;
    while (TimeOut--) {
      gBS->Stall (500);

      PciIo->Pci.Read (
                   PciIo,
                   EfiPciIoWidthUint32,
                   ExtendCap,
                   1,
                   &Value
                   );

      if ((Value & BIT24) == 0x0) {
        break;
      }
    }
  }

  PciIo->Pci.Read (
               PciIo,
               EfiPciIoWidthUint32,
               ExtendCap,
               1,
               &Value
               );

  Value |= BIT24;
  PciIo->Pci.Write (
               PciIo,
               EfiPciIoWidthUint32,
               ExtendCap,
               1,
               &Value
               );

  TimeOut = 40;
  while (TimeOut--) {
    gBS->Stall (500);

    PciIo->Pci.Read (
                 PciIo,
                 EfiPciIoWidthUint32,
                 ExtendCap,
                 1,
                 &Value
                 );

    if ((Value & BIT16) == 0x0) {
      break;
    }
  }

  IsOwnershipConflict = (Value & BIT16) != 0x0;
  if (IsOwnershipConflict) {
    //
    // Soft reset has failed. Assume SMI being ignored and do hard reset.
    //
    Value = 0;
    PciIo->Pci.Write (
                 PciIo,
                 EfiPciIoWidthUint8,
                 ExtendCap + 2,
                 1,
                 &Value
                 );

    TimeOut = 40;
    while (TimeOut--) {
      gBS->Stall (500);

      PciIo->Pci.Read (
                   PciIo,
                   EfiPciIoWidthUint32,
                   ExtendCap,
                   1,
                   &Value
                   );

      if ((Value & BIT16) == 0x0) {
        break;
      }
    }

    //
    // Disable further SMI events.
    //
    PciIo->Pci.Read (
                 PciIo,
                 EfiPciIoWidthUint32,
                 ExtendCap + 0x4,
                 1,
                 &UsbLegCtlSts
                 );

    UsbLegCtlSts &= 0xFFFF0000U;
    PciIo->Pci.Write (
                 PciIo,
                 EfiPciIoWidthUint32,
                 ExtendCap + 0x4,
                 1,
                 &UsbLegCtlSts
                 );
  }

  if (Value & BIT16) {
    //
    // EHCI controller unable to take control from BIOS.
    //
    Status = EFI_NOT_FOUND;
  }

  return Status;
}

/**
 Release UHCI USB controllers ownership.

 @param[in]  PciIo  PCI I/O protocol for the device.

 @retval EFI_NOT_FOUND  No UHCI controllers had ownership incorrectly set.
 @retval EFI_SUCCESS    Corrected UHCI controllers ownership to OS.
 **/
STATIC
EFI_STATUS
UhciReleaseOwnership (
  IN  EFI_PCI_IO_PROTOCOL  *PciIo
  )
{
  EFI_STATUS  Status;
  UINT32      Base;
  UINT32      PortBase;
  UINT16      Command;

  Base = 0;

  PciIo->Pci.Read (
               PciIo,
               EfiPciIoWidthUint32,
               0x20,
               1,
               &Base
               );

  PortBase = (Base >> 5) & 0x07ff;

  Command = 0x8f00;

  Status = PciIo->Pci.Write (
                        PciIo,
                        EfiPciIoWidthUint16,
                        0xC0,
                        1,
                        &Command
                        );

  if ((PortBase != 0) && ((PortBase & BIT0) == 0)) {
    IoWrite16 (PortBase, 0x0002);
    gBS->Stall (500);
    IoWrite16 (PortBase + 4, 0);
    gBS->Stall (500);
    IoWrite16 (PortBase, 0);
  }

  return Status;
}

EFI_STATUS
ReleaseUsbOwnership (
  VOID
  )
{
  EFI_STATUS           Result;
  EFI_STATUS           Status;
  EFI_HANDLE           *HandleArray;
  UINTN                HandleArrayCount;
  UINTN                Index;
  EFI_PCI_IO_PROTOCOL  *PciIo;
  PCI_TYPE00           Pci;

  Status = gBS->LocateHandleBuffer (
                  ByProtocol,
                  &gEfiPciIoProtocolGuid,
                  NULL,
                  &HandleArrayCount,
                  &HandleArray
                  );

  if (EFI_ERROR (Status)) {
    return Status;
  }

  Result = EFI_UNSUPPORTED;

  for (Index = 0; Index < HandleArrayCount; ++Index) {
    Status = gBS->HandleProtocol (
                    HandleArray[Index],
                    &gEfiPciIoProtocolGuid,
                    (VOID **)&PciIo
                    );

    if (EFI_ERROR (Status)) {
      continue;
    }

    Status = PciIo->Pci.Read (
                          PciIo,
                          EfiPciIoWidthUint32,
                          0,
                          sizeof (Pci) / sizeof (UINT32),
                          &Pci
                          );

    if (  EFI_ERROR (Status)
       || (Pci.Hdr.ClassCode[1] != PCI_CLASS_SERIAL_USB)
       || (Pci.Hdr.ClassCode[2] != PCI_CLASS_SERIAL))
    {
      continue;
    }

    if (Pci.Hdr.ClassCode[0] == PCI_IF_XHCI) {
      Result = XhciReleaseOwnership (PciIo);
    } else if (Pci.Hdr.ClassCode[0] == PCI_IF_EHCI) {
      Result = EhciReleaseOwnership (PciIo);
    } else if (Pci.Hdr.ClassCode[0] == PCI_IF_UHCI) {
      Result = UhciReleaseOwnership (PciIo);
    }
  }

  gBS->FreePool (HandleArray);
  return Result;
}
