/** @file
  Linux boot driver, supporting Boot Loader Specification, GRUB2 blscfg, and autodetect.

  Copyright (c) 2021, Mike Beaton. All rights reserved.<BR>
  SPDX-License-Identifier: BSD-3-Clause
**/

#include "LinuxBootInternal.h"

#include <Uefi.h>
#include <Guid/Gpt.h>
#include <Library/BaseLib.h>
#include <Library/DevicePathLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/OcBootManagementLib.h>
#include <Library/OcDebugLogLib.h>
#include <Library/OcFileLib.h>
#include <Library/OcFlexArrayLib.h>
#include <Library/OcStringLib.h>
#include <Library/UefiBootServicesTableLib.h>

#include <Protocol/OcBootEntry.h>

UINTN gLinuxBootFlags = LINUX_BOOT_ALL & ~LINUX_BOOT_ADD_DEBUG_INFO;

OC_PICKER_CONTEXT *gPickerContext;
OC_FLEX_ARRAY     *gParsedLoadOptions;
OC_FLEX_ARRAY     *gNamedLoaderEntries;
EFI_GUID          gPartuuid;
CHAR8             *gFileSystemType;

VOID
InternalFreePickerEntry (
  IN   OC_PICKER_ENTRY          *Entry
  )
{
  ASSERT (Entry != NULL);

  if (Entry == NULL) {
    return;
  }

  //
  // TODO: Is this un-CONST casting okay?
  // (Are they CONST because they are not supposed to be freed when used as before?)
  //
  if (Entry->Id != NULL) {
    FreePool ((CHAR8 *)Entry->Id);
  }
  if (Entry->Name != NULL) {
    FreePool ((CHAR8 *)Entry->Name);
  }
  if (Entry->Path != NULL) {
    FreePool ((CHAR8 *)Entry->Path);
  }
  if (Entry->Arguments != NULL) {
    FreePool ((CHAR8 *)Entry->Arguments);
  }
  if (Entry->Flavour != NULL) {
    FreePool ((CHAR8 *)Entry->Flavour);
  }
}

STATIC
VOID
EFIAPI
OcFreeLinuxBootEntries (
  IN   OC_PICKER_ENTRY          **Entries,
  IN   UINTN                    NumEntries
  )
{
  UINTN Index;

  ASSERT (Entries   != NULL);
  ASSERT (*Entries  != NULL);
  if (Entries == NULL || *Entries == NULL) {
    return;
  }

  for (Index = 0; Index < NumEntries; Index++) {
    InternalFreePickerEntry (&(*Entries)[Index]);
  }

  FreePool (*Entries);
  *Entries = NULL;
}

STATIC
EFI_STATUS
EFIAPI
OcGetLinuxBootEntries (
  IN           OC_PICKER_CONTEXT        *PickerContext,
  IN     CONST EFI_HANDLE               Device,
     OUT       OC_PICKER_ENTRY          **Entries,
     OUT       UINTN                    *NumEntries
  )
{
  EFI_STATUS                      Status;
  EFI_SIMPLE_FILE_SYSTEM_PROTOCOL *FileSystem;
  EFI_FILE_PROTOCOL               *RootDirectory;
  UINT32                          FileSystemPolicy;
  CONST EFI_PARTITION_ENTRY       *PartitionEntry;

  ASSERT (PickerContext != NULL);
  ASSERT (Entries     != NULL);
  ASSERT (NumEntries  != NULL);

  gPickerContext  = PickerContext;
  *Entries        = NULL;
  *NumEntries     = 0;

  //
  // No custom entries.
  //
  if (Device == NULL) {
    return EFI_NOT_FOUND;
  }

  //
  // Open partition file system.
  //
  Status = gBS->HandleProtocol (
    Device,
    &gEfiSimpleFileSystemProtocolGuid,
    (VOID **) &FileSystem
    );
  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_WARN, "LNX: Missing filesystem - %r\n", Status));
    return Status;
  }

  //
  // Get handle to partiton root directory.
  //
  Status = FileSystem->OpenVolume (
    FileSystem,
    &RootDirectory
    );
  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_WARN, "LNX: Invalid root volume - %r\n", Status));
    return Status;
  }

  gFileSystemType   = NULL;

  FileSystemPolicy = OcGetFileSystemPolicyType (Device);

  //
  // Disallow Apple filesystems, mainly to avoid needlessly
  // scanning multiple APFS partitions.
  //
  if ((FileSystemPolicy & OC_SCAN_ALLOW_FS_APFS) != 0) {
    gFileSystemType = "APFS";
    DEBUG ((DEBUG_INFO, "LNX: %a - not scanning\n", gFileSystemType));
    Status = EFI_NOT_FOUND;
  } else if ((FileSystemPolicy & OC_SCAN_ALLOW_FS_HFS) != 0) {
    gFileSystemType = "HFS";
    DEBUG ((DEBUG_INFO, "LNX: %a - not scanning\n", gFileSystemType));
    Status = EFI_NOT_FOUND;
  } else if ((FileSystemPolicy & OC_SCAN_ALLOW_FS_ESP) != 0) {
    gFileSystemType = "ESP";
    if ((gLinuxBootFlags & LINUX_BOOT_SCAN_ESP) == 0) {
      DEBUG ((DEBUG_INFO, "LNX: %a - requested not to scan\n", gFileSystemType));
      Status = EFI_NOT_FOUND;
    }
  } else if ((FileSystemPolicy & OC_SCAN_ALLOW_FS_XBOOTLDR) != 0) {
    gFileSystemType = "XBOOTLDR";
    if ((gLinuxBootFlags & LINUX_BOOT_SCAN_XBOOTLDR) == 0) {
      DEBUG ((DEBUG_INFO, "LNX: %a - requested not to scan\n", gFileSystemType));
      Status = EFI_NOT_FOUND;
    }
  } else if ((FileSystemPolicy & OC_SCAN_ALLOW_FS_LINUX_ROOT) != 0) {
    gFileSystemType = "ROOT";
    if ((gLinuxBootFlags & LINUX_BOOT_SCAN_LINUX_ROOT) == 0) {
      DEBUG ((DEBUG_INFO, "LNX: %a - requested not to scan\n", gFileSystemType));
      Status = EFI_NOT_FOUND;
    }
  } else if ((FileSystemPolicy & OC_SCAN_ALLOW_FS_LINUX_DATA) != 0) {
    gFileSystemType = "DATA";
    if ((gLinuxBootFlags & LINUX_BOOT_SCAN_LINUX_DATA) == 0) {
      DEBUG ((DEBUG_INFO, "LNX: %a - requested not to scan\n", gFileSystemType));
      Status = EFI_NOT_FOUND;
    }
  } else {
    if ((FileSystemPolicy & OC_SCAN_ALLOW_FS_NTFS) != 0) {
      //
      // This is not just NTFS, Msft Basic Data part type GUID is used for non-ESP FAT too.
      //  
      gFileSystemType = "NTFS/FAT";
    } else {
      gFileSystemType = "OTHER";
    }

    if ((gLinuxBootFlags & LINUX_BOOT_SCAN_OTHER) == 0) {
      DEBUG ((DEBUG_INFO, "LNX: %a - requested not to scan\n", gFileSystemType));
      Status = EFI_NOT_FOUND;
    }
  }

  if (EFI_ERROR (Status)) {
    RootDirectory->Close (RootDirectory);
    return Status;
  }

  //
  // Save PARTUUID for autodetect.
  //
  PartitionEntry = OcGetGptPartitionEntry (Device);
  if (PartitionEntry == NULL) {
    gPartuuid = gEfiPartTypeUnusedGuid;
  } else {
    gPartuuid = PartitionEntry->UniquePartitionGUID;
  }

  //
  // Log TypeGUID and PARTUUID of the drive we're in.
  //
  DEBUG ((
    DEBUG_INFO,
    "LNX: TypeGUID: %g (%a) PARTUUID: %g\n",
    PartitionEntry->PartitionTypeGUID,
    gFileSystemType,
    PartitionEntry->UniquePartitionGUID
    ));

  //
  // Scan for boot loader spec & blscfg entries (Fedora-like).
  //
  Status = ScanLoaderEntries (
    RootDirectory,
    Entries,
    NumEntries
    );

  //
  // Note: As currently structured, will fall through to autodetect
  // if no /loader/entries/*.conf files are present, but also if there
  // are only unusable files in there.
  //
  if (EFI_ERROR (Status)) {
    if (Status != EFI_NOT_FOUND) {
      DEBUG ((DEBUG_WARN, "LNX: ScanLoaderEntries - %r\n", Status));
    }
    
    //
    // Auto-detect vmlinuz and initrd files on own root filesystem (Debian-like).
    //
    if ((gLinuxBootFlags & LINUX_BOOT_ALLOW_AUTODETECT) != 0) {
      Status = AutodetectLinux (
        RootDirectory,
        Entries,
        NumEntries
        );

      if (EFI_ERROR (Status) && Status != EFI_NOT_FOUND) {
        DEBUG ((DEBUG_WARN, "LNX: AutodetectLinux - %r\n", Status));
      }
    }
  }

  if (Status == EFI_NOT_FOUND) {
    DEBUG ((DEBUG_INFO, "LNX: Nothing found\n"));
  }

  RootDirectory->Close (RootDirectory);

  return Status;
}

STATIC
OC_BOOT_ENTRY_PROTOCOL
mLinuxBootEntryProtocol = {
  OC_BOOT_ENTRY_PROTOCOL_REVISION,
  OcGetLinuxBootEntries,
  OcFreeLinuxBootEntries
};

EFI_STATUS
EFIAPI
UefiMain (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE  *SystemTable
  )
{
  EFI_STATUS        Status;

  Status = OcParseLoadOptions (ImageHandle, &gParsedLoadOptions);
  if (!EFI_ERROR (Status)) {
    OcParsedVarsGetInt (gParsedLoadOptions, L"flags", &gLinuxBootFlags, TRUE);
  } else if (Status != EFI_NOT_FOUND) {
    return Status;
  }

  Status = gBS->InstallMultipleProtocolInterfaces (
    &ImageHandle,
    &gOcBootEntryProtocolGuid,
    &mLinuxBootEntryProtocol,
    NULL
    );

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

  return EFI_SUCCESS;
}