// SPDX-License-Identifier: GPL-3.0-or-later
#include "elilong.h"

#define EFI_GLOBAL_VARIABLE_GUID {0x8BE4DF61, 0x93CA, 0x11d2, {0xAA, 0x0D, 0x00, 0xE0, 0x98, 0x03, 0x2B, 0x8C}}

// function to check if secureboot is on
BOOLEAN is_secure_boot_enabled(void) {
    EFI_GUID global_guid = EFI_GLOBAL_VARIABLE_GUID;
    UINT8 sb_value = 0;
    UINT8 sm_value = 0;
    UINTN size = sizeof(UINT8);
    EFI_STATUS status;

    // Check if SecureBoot variable is set to 1
    status = uefi_call_wrapper(RT->GetVariable, 5,
                               L"SecureBoot", &global_guid,
                               NULL, &size, &sb_value);
    if (EFI_ERROR(status) || sb_value == 0) {
        return FALSE;
    }

    // Check if system is in SetupMode (if SetupMode == 1, Secure Boot is active but not strictly enforced)
    size = sizeof(UINT8);
    status = uefi_call_wrapper(RT->GetVariable, 5,
                               L"SetupMode", &global_guid,
                               NULL, &size, &sm_value);
    if (!EFI_ERROR(status) && sm_value == 1) {
        return FALSE; // Keys are being edited/cleared, not strictly enforcing
    }

    return TRUE;
}

VOID display_message_file(EFI_FILE_HANDLE root, CHAR16 *path) {
  EFI_STATUS status;
  EFI_FILE_HANDLE file;
  UINTN buffer_size = 2000; // max charaters of message to display
  CHAR8 buffer[2001];       // CHAR8 for ASCII content

  if (path[0] == L'\0') return;  //return silently if no path specified

  status = uefi_call_wrapper(root->Open, 5, root, &file, path, EFI_FILE_MODE_READ, 0);
  if (EFI_ERROR(status)) {
    Print(L"Message file '%s' not found.\n", path);
    return;
  }

  // Read the file (up to 2000 (ascii) bytes)
  status = uefi_call_wrapper(file->Read, 3, file, &buffer_size, buffer);
  if (!EFI_ERROR(status) && buffer_size > 0) {
    buffer[buffer_size] = '\0'; // Null terminate the ASCII buffer
    Print(L"%a\n", buffer);     // use %a to safely print ascii
  }

  uefi_call_wrapper(file->Close, 1, file);
}

// we use this for handling long strings in menu display. if too long to fit in, remove the middle
// and replace with ... it doesn't make any sense to use this if maxlen is 4 or less.
VOID FormatMiddleTruncate(CONST CHAR16 *InputStr, CHAR16 *OutputBuffer, UINTN  MaxLen) {
  UINTN Len = StrLen(InputStr);
  UINTN i;
  if (Len <= MaxLen) {  // we pad it to maxLen with spaces because uefi print function is buggy
    // and doesn't handle print(L"%-*.*s",var1,var1,string1); properly and right justifies instead of
    // left justifying the string.
    StrCpy(OutputBuffer, InputStr);
    for (i = Len; i < MaxLen; i++) {
      OutputBuffer[i] = L' ';
    }
    OutputBuffer[MaxLen] = L'\0';
    return;
  }
  if (MaxLen <= 3) { // If they ask for 3 or fewer characters, we can only fit the dots (or a subset of them)
    StrnCpy(OutputBuffer, L"...", MaxLen);
    OutputBuffer[MaxLen] = L'\0';
    return;
  }
  UINTN AvailableSpace = MaxLen - 3; // Space left after accounting for "..."
  UINTN FrontLen = AvailableSpace / 2;
  UINTN BackLen  = AvailableSpace - FrontLen;
  if (FrontLen > 0) {
    StrnCpy(OutputBuffer, InputStr, FrontLen);
  }
  OutputBuffer[FrontLen] = L'\0';
  StrCat(OutputBuffer, L"...");
  CONST CHAR16 *BackStart = InputStr + (Len - BackLen);
  StrCat(OutputBuffer, BackStart);
}

BOOT_ENTRY *display_menu(BOOT_ENTRY *head, INTN config_timeout, CHAR16 *default_label, UINTN scrwidth, UINTN scrheight) {
  EFI_INPUT_KEY key;
  EFI_STATUS status;
  UINTN index = 1;
  UINTN selection = 0;
  BOOT_ENTRY *search = head;
  INTN timeout_countdown = config_timeout;

  UINTN max_label_found = 0;
  UINTN max_desc_found = 0;
  UINTN total_entries = 0;
  UINTN default_index = 1; //if default doesn't match a label, or is not set use first entry.

  BOOLEAN sb_active = is_secure_boot_enabled();

  // if prompt is not specified, and neither was delay or timeout, then skip showing menu
  // and boot default, or first entry if not specified
  if (config_timeout == -2) {
    search = head;
    while (search != NULL) {
      if (default_label[0] != L'\0' && StriCmp(search->label, default_label) == 0) {
        return search;
      }
      search = search->next;
    }
    return head;
  }
  // we should show the menu if we get this far.
  CHAR16 truncated_label[MAX_LABEL]; // Buffers to store processed/truncated text strings cleanly
  CHAR16 truncated_desc[MAX_DESC];
  UINTN start_row = ST->ConOut->Mode->CursorRow;  //get current row.
  //lets parse list for max widths and default index.
  while (search != NULL) {
    UINTN l_len = StrLen(search->label);
    UINTN d_len = StrLen(search->description);
    if (l_len > max_label_found) max_label_found = l_len;
    if (d_len > max_desc_found)  max_desc_found = d_len;
    if (default_label[0] != L'\0' && StriCmp(search->label, default_label) == 0) {
      default_index = total_entries + 1;
    }
    total_entries++;
    search = search->next;
  }
  if (total_entries > MAX_ENTRIES) {  // cap number of boot entries
    total_entries = MAX_ENTRIES;
    // TODO: for now we cap the number of boot entries. so they fit on the screen and don't
    // break the menu display. if we impliment scrollable menu items we could allow for more.
  }

  // now to calculate offset and width of label and description fields to display
  // menu layout is "║YXX label description║"  where Y is default indicator, XX is index,
  // label is 20% of remainining real estate and desc is 80% of remaining real estate
  // Overhead: 1(║) 1("*"|" ")+ 2(index) + 1(space) + 1(space separator) + 1(║) = 7 chars fixed frame space.
  UINTN fixed_frame_overhead = 7;
  UINTN label_width = max_label_found;
  UINTN desc_width = max_desc_found;
  UINTN total_box_width = label_width + desc_width + fixed_frame_overhead;
  UINTN left_margin = 0;
  if (total_box_width < scrwidth) { // Content fits fully! Center using half the remaining real estate
    left_margin = (scrwidth - total_box_width) / 2;
  } else { // Content overflows! Force box to utilize max screen width limits
    total_box_width = scrwidth;
    left_margin = 0;
    //if we have a at least a character for label or description, we need 5 characters for 20/80 split
    UINTN usable_space = (total_box_width > (fixed_frame_overhead + 4)) ? (total_box_width - fixed_frame_overhead) : 10;
    // Split remaining space after fixed frame overhead: 20% Label, 80% Description
    label_width = (usable_space * 20) / 100;
    desc_width = usable_space - label_width;
  }
  // we enforce a hard limit of screen size to 12 characters (MIN_SCR_WIDTH)
  // if any uefi bios supports a smaller screen let us know. we could potentially remove menu borders
  // in that case to allow going down to a max of 10 characters. typically the minimum is 80 characters.
  if (total_box_width < MIN_SCR_WIDTH) {
    total_box_width = MIN_SCR_WIDTH;
    left_margin = 0;
    label_width = 1;
    desc_width = 4;
  }

  // now for some colour stuff. we save the existing screen colour state
  UINTN original_attributes = ST->ConOut->Mode->Attribute;
  // there seems to be macro incompatibililty issues between gnu-efi and edk2 so this doesn't work
  // menu_attributes = EFI_TEXT_ATTR(EFI_LIGHTGRAY, EFI_BACKGROUND_BLUE);

  // lets use the manual method instead
  // 0x0F is Light Gray (Foreground), 0x01 is Blue (Background)
  // This constructs the exact hex attribute byte 0x1F (Light Gray on Blue)
  UINTN menu_attributes = 0x0F | (0x01 << 4);

  // now to display menu
  Print(L"\n");  //move to new line,. (make sure we're not on the same line as any startup message

  // calculate the height of the menu plus countdown and prompt and buffer allowance
  // Box frame: 1 (Top) + total_entries + 1 (Bottom)
  // Status lines: 1 (Spacing) + 1 (Countdown) + 1 (Prompt) + 1 (Buffer space)
  UINTN total_ui_height = 1 + total_entries + 1 + 4;
  start_row = ST->ConOut->Mode->CursorRow; // Establish our baseline text line
  // if the menu doesn't fit below any message text, scroll it up till it will fit
  if (start_row + total_ui_height >= scrheight) {
    UINTN overflow_lines = (start_row + total_ui_height) - scrheight + 1;
    // Move cursor to the bottom line of the screen and add lines to scroll screen up
    uefi_call_wrapper(ST->ConOut->SetCursorPosition, 3, ST->ConOut, 0, scrheight - 1);
    UINTN j;
    for (j = 0; j < overflow_lines; j++) {
      Print(L"\n");
    }
    // set the menu start to it's adjusted position. if it doesn't fit, stick it at the top of the screen
    start_row = (start_row >= overflow_lines) ? (start_row - overflow_lines) : 0;
  }
  // note: TODO handle if the menu doesn't fit virtically on the screen. ie >20 boot entries
  // for now we will probably cap the number of entries to 18.

  // Draw top menu border
  uefi_call_wrapper(ST->ConOut->SetCursorPosition, 3, ST->ConOut, left_margin, start_row);
  uefi_call_wrapper(ST->ConOut->SetAttribute, 2, ST->ConOut, menu_attributes); // Turn on Blue background
  Print(L"%c", BOX_TL);
  UINTN i;
  for (i = 0; i < (total_box_width - 7); i++) {
      Print(L"%c", BOX_H);
  }
  if (sb_active) {
      // Light Gray text on RED background (0x4F)
      uefi_call_wrapper(ST->ConOut->SetAttribute, 2, ST->ConOut, 0x4F);
      Print(L"SB=ON");
      uefi_call_wrapper(ST->ConOut->SetAttribute, 2, ST->ConOut, menu_attributes);
  } else {
      Print(L"%c%c%c%c%c", BOX_H, BOX_H, BOX_H, BOX_H, BOX_H);
  }
  Print(L"%c", BOX_TR);
  uefi_call_wrapper(ST->ConOut->SetAttribute, 2, ST->ConOut, original_attributes); // Turn off Blue before \n
  Print(L"\n");
  start_row++;
  UINTN menu_start_row = start_row;  //row of first menu item
  selection = default_index;
  index = 1;
  search = head;    // Rewind linked list to the beginning to render text items after the max length parsing

  // print main menu entries
  while (search != NULL) {
    // Apply Middle Truncation constraints directly to our working copy buffers
    FormatMiddleTruncate(search->label, truncated_label, label_width);
    FormatMiddleTruncate(search->description, truncated_desc, desc_width);
    // Lock cursor exactly to our calculated horizontal left margin anchor
    uefi_call_wrapper(ST->ConOut->SetCursorPosition, 3, ST->ConOut, left_margin, start_row);
    uefi_call_wrapper(ST->ConOut->SetAttribute, 2, ST->ConOut, menu_attributes); // Turn on Blue
    // Render index number data block inside frame
    Print(L"%c%c", BOX_V, (index == selection) ? L'*' : L' ');
    // Render fixed width components matching format width arguments explicitly
    Print(L"%2d %s %s%c", index,truncated_label, truncated_desc, BOX_V);
    uefi_call_wrapper(ST->ConOut->SetAttribute, 2, ST->ConOut, original_attributes); // Turn off Blue before \n
    Print(L"\n");
    start_row++;
    index++;
    search = search->next;
  }

  // Draw bottom menu border
  uefi_call_wrapper(ST->ConOut->SetCursorPosition, 3, ST->ConOut, left_margin, start_row);
  uefi_call_wrapper(ST->ConOut->SetAttribute, 2, ST->ConOut, menu_attributes); // Turn on Blue
  Print(L"%c%sbit", BOX_BL,get_bitness_str()); // indicate the bitness using 2 characters
  for (i = 0; i < (total_box_width - 7); i++) Print(L"%c", BOX_H);
  Print(L"%c", BOX_BR);
  uefi_call_wrapper(ST->ConOut->SetAttribute, 2, ST->ConOut, original_attributes); // Turn off Blue permanently
  Print(L"\n\n");

  UINTN message_row = ST->ConOut->Mode->CursorRow;
  CHAR16 typed_buffer[3] = L""; // Tracks up to two typed digits + null terminator
  UINTN typed_count = 0;
  BOOLEAN overflow_cleared = FALSE; // Tracks if the wrapped row has been scrubbed

  // we turn off the watchdog timer as we don't want the system rebooting after 5 mins when waiting
  // for the user to make a menu selection
  gBS->SetWatchdogTimer(0, 0, 0, NULL);


  // now for the main menu loop to get keypresses. if there is a delay or timeout set, then do a countdown,
  // defaulting to the default value if nothing is pressed in time. if a key is pressed, stop the
  // countdown. use arrow keys to move up and down the menu items, or enter the menu index and
  // press enter to boot the selected item.

  while (1) {
    uefi_call_wrapper(ST->ConOut->SetCursorPosition, 3, ST->ConOut, 0, message_row);

    if (timeout_countdown > 0) {
      Print(L"Booting selection in %d seconds. \n", timeout_countdown / 10);
      Print(L"Enter selection (1-%d) or use Arrows [Choice: %d, Typed: %s]    ", total_entries, selection, typed_buffer);
    } else {
      Print(L"Enter selection (1-%d) or use Arrows [Choice: %d, Typed: %s]    ", total_entries, selection, typed_buffer);

      if (!overflow_cleared) {
        uefi_call_wrapper(ST->ConOut->SetCursorPosition, 3, ST->ConOut, 0, message_row + 1);
        Print(L"                                                                             ");
        overflow_cleared = TRUE;
      }
    }

    status = uefi_call_wrapper(ST->ConIn->ReadKeyStroke, 2, ST->ConIn, &key);
    if (status == EFI_SUCCESS) {
      timeout_countdown = -1;

      // 1. ARROW CONTROLS
      if (key.UnicodeChar == 0) {
        typed_buffer[0] = L'\0';
        typed_count = 0;

        UINTN old_selection = selection;
        if (key.ScanCode == SCAN_UP && selection > 1) {
          selection--;
        } else if (key.ScanCode == SCAN_DOWN && selection < total_entries) {
          selection++;
        }

        if (selection != old_selection) {
          uefi_call_wrapper(ST->ConOut->SetCursorPosition, 3, ST->ConOut, left_margin + 1, menu_start_row + (old_selection - 1));
          uefi_call_wrapper(ST->ConOut->SetAttribute, 2, ST->ConOut, menu_attributes);
          Print(L" ");
          uefi_call_wrapper(ST->ConOut->SetCursorPosition, 3, ST->ConOut, left_margin + 1, menu_start_row + (selection - 1));
          Print(L"*");
          uefi_call_wrapper(ST->ConOut->SetAttribute, 2, ST->ConOut, original_attributes);
        }
      }
      // 2. CONFIRM AND BOOT
      else if (key.UnicodeChar == L'\r' || key.UnicodeChar == L'\n') {
        break;
      }
      // 3. LIVE NUMERIC INPUT
      else if (key.UnicodeChar >= L'0' && key.UnicodeChar <= L'9') {
        if (typed_count < 2) {
          typed_buffer[typed_count++] = key.UnicodeChar;
          typed_buffer[typed_count] = L'\0';

          UINTN parsed = 0;
          if (typed_count == 1) parsed = typed_buffer[0] - L'0';
          else parsed = ((typed_buffer[0] - L'0') * 10) + (typed_buffer[1] - L'0');

          if (parsed >= 1 && parsed <= total_entries && parsed != selection) {
            UINTN old_selection = selection;
            selection = parsed;

            uefi_call_wrapper(ST->ConOut->SetCursorPosition, 3, ST->ConOut, left_margin + 1, menu_start_row + (old_selection - 1));
            uefi_call_wrapper(ST->ConOut->SetAttribute, 2, ST->ConOut, menu_attributes);
            Print(L" ");
            uefi_call_wrapper(ST->ConOut->SetCursorPosition, 3, ST->ConOut, left_margin + 1, menu_start_row + (selection - 1));
            Print(L"*");
            uefi_call_wrapper(ST->ConOut->SetAttribute, 2, ST->ConOut, original_attributes);
          }
        }
      }
      // 4. DIGIT BACKSPACE REMAPPING
      else if (key.UnicodeChar == 0x08 && typed_count > 0) {
        typed_buffer[--typed_count] = L'\0';
        UINTN old_selection = selection;

        if (typed_count > 0) {
          UINTN parsed = 0;
          if (typed_count == 1) parsed = typed_buffer[0] - L'0';
          else parsed = ((typed_buffer[0] - L'0') * 10) + (typed_buffer[1] - L'0');

          if (parsed >= 1 && parsed <= total_entries) {
            selection = parsed;
          }
        } else {
          selection = 1;
        }

        if (selection != old_selection) {
          uefi_call_wrapper(ST->ConOut->SetCursorPosition, 3, ST->ConOut, left_margin + 1, menu_start_row + (old_selection - 1));
          uefi_call_wrapper(ST->ConOut->SetAttribute, 2, ST->ConOut, menu_attributes);
          Print(L" ");
          uefi_call_wrapper(ST->ConOut->SetCursorPosition, 3, ST->ConOut, left_margin + 1, menu_start_row + (selection - 1));
          Print(L"*");
          uefi_call_wrapper(ST->ConOut->SetAttribute, 2, ST->ConOut, original_attributes);
        }
      }
    }

    if (timeout_countdown > 0) {
      uefi_call_wrapper(ST->BootServices->Stall, 1, 100000);  //wait 0.1seconds
      timeout_countdown--;
      if (timeout_countdown == 0) {  //if we timed out, quit the loop
        break;
      }
    } else {
      uefi_call_wrapper(ST->BootServices->Stall, 1, 20000); // no timeout, so check every 0.02secs for input
    }
  }
  uefi_call_wrapper(ST->ConOut->SetCursorPosition, 3, ST->ConOut, 0, message_row + 2);
  Print(L"\nBooting selection %d...\n", selection);

  // lets get the relevent entry and return it.
  index = 1;
  search = head;
  while (search != NULL) {
    if (index == selection) {
      return search;
    }
    search = search->next;
    index++;
  }
  // lets reinstate the uefi watchdog timer to 5mins now we're ready to load the kernel.
  // 0x00010000 is a standard code used to denote a generic firmware timeout
  gBS->SetWatchdogTimer(300, 0x00010000, 0, NULL);
  return head;
}
// Helper to check if a file exists on the ESP
BOOLEAN file_exists(EFI_FILE_HANDLE root, CHAR16 *path) {
    EFI_FILE_HANDLE file;
    EFI_STATUS status = uefi_call_wrapper(root->Open, 5, root, &file, path, EFI_FILE_MODE_READ, 0);
    if (status == EFI_SUCCESS) {
        uefi_call_wrapper(file->Close, 1, file);
        return TRUE;
    }
    return FALSE;
}
/* *
   * Entry point for the elilong UEFI bootloader.
   */
EFI_STATUS efi_main(EFI_HANDLE image_handle, EFI_SYSTEM_TABLE *system_table) {
  EFI_STATUS status = EFI_SUCCESS;  //initialise status flag
  BOOT_CONFIG config; // this is the structure that will hold the config file information parsed, and
                      // other system-wide info.
  InitializeLib(image_handle, system_table);
  system_table->ConOut->ClearScreen(system_table->ConOut);
  Print(L"%s elilong v%a...\n", DISTRO, VERSION);  // use %a as version is passed as ascii.
  uefi_call_wrapper(BS->Stall, 1, 2000000);  //note hard-coded delay to see this message 2sec.
  //this should give a monitor/uefi bios graphics time to initialise. may remove this or add a quiet option
  //to control it.
  ZeroMem(&config, sizeof(BOOT_CONFIG)); // initialise it to zeros
  status = uefi_call_wrapper(ST->ConOut->QueryMode,4,ST->ConOut,ST->ConOut->Mode->Mode,&config.scrwidth,&config.scrheight);
  // in case there was some failure in the call use sensible defaults
  if (EFI_ERROR(status) || config.scrwidth < MIN_SCR_WIDTH) {
    config.scrwidth = 80;
    config.scrheight = 25;
  }
  EFI_FILE_HANDLE root = get_esp_root(image_handle);
  if (root == NULL) {
      Print(L"Error: Unable to access the EFI system partition (ESP).\n");
      // Give the user 5 seconds to read the error before exiting to shell
      uefi_call_wrapper(ST->BootServices->Stall, 1, 5000000);
      return EFI_LOAD_ERROR;
  }
  CHAR16 *conf_path = L"\\EFI\\" SUBDIR L"\\elilong.conf";
  CHAR16 *legacy_conf = L"\\EFI\\" SUBDIR L"\\elilo.conf";
  CHAR16 *fallback_path = L"\\EFI\\BOOT\\BOOTX64.EFI";
  // Print(L"Searching for configuration at %s or %s\n", conf_path, legacy_conf);
  // Try elilong.conf then elilo.conf using existing variables
  if (load_config(root, conf_path, &config) == EFI_SUCCESS ||
      load_config(root, legacy_conf, &config) == EFI_SUCCESS) {
      //show message file if one is in the config file and exists.
      if (config.message_path[0] != L'\0') {
          display_message_file(root, config.message_path);
          Print(L"\n"); // Clean buffer space between splash text and options
      }
      // Pass the list and the parsed/default timeout to the menu
      BOOT_ENTRY *chosen = display_menu(config.entries, config.timeout, config.default_label, config.scrwidth, config.scrheight);

      if (chosen != NULL) {
          manual_pe_load_and_jump(image_handle, root, chosen);
      }
  } else {
      // Fallback logic using existing fallback_path variable
      if (file_exists(root, fallback_path)) {
          Print(L"No config found. Press 'B' to boot fallback or any other key for Shell...\n");
          if (wait_for_specific_char(L'b')) {
              // Create a temporary boot entry for the fallback
              BOOT_ENTRY fallback_entry;
              ZeroMem(&fallback_entry, sizeof(BOOT_ENTRY));
              StrCpy(fallback_entry.label, L"Emergency Fallback");
              StrCpy(fallback_entry.kernel_path, fallback_path);
              fallback_entry.initrd_path[0] = L'\0';
              fallback_entry.append[0] = L'\0';

              manual_pe_load_and_jump(image_handle, root, &fallback_entry);
          }
      }

      // If we can't find any config files nor a fallback file
      Print(L"Exiting to UEFI Shell...\n");
      return EFI_SUCCESS;
  }

  return EFI_SUCCESS;
}

