#include <acpi.h>

#include <types.h>

/* ACPI 2.0 RSDP: eb9d2d30-2d88-11d3-9a16-0090273fc14d */
static const uint8_t guid_rsdp_v2[16] = {
  0x30, 0x2D, 0x9D, 0xEB, 0x88, 0x2D, 0xD3, 0x11,
  0x9A, 0x16, 0x00, 0x90, 0x27, 0x3F, 0xC1, 0x4D
};

/* ACPI 1.0 RSDP: 8868e871-e4f1-11d3-bc22-0080c73c8881 */
static const uint8_t guid_rsdp_v1[16] = {
  0x71, 0xE8, 0x68, 0x88, 0xF1, 0xE4, 0xD3, 0x11,
  0xBC, 0x22, 0x00, 0x80, 0xC7, 0x3C, 0x88, 0x81
};

#define MADT_TYPE_LOCAL_APIC 0
#define MADT_TYPE_IO_APIC 1
/* x2APIC: replaces type 0 when the firmware enabled x2APIC mode */
#define MADT_TYPE_X2APIC 9

/* signatures read little-endian via read_u32(): the first four ASCII bytes. */
#define ACPI_SIG_RSDP 0x20445352u /* "RSD " */
#define ACPI_SIG_RSDT 0x54445352u /* "RSDT" */
#define ACPI_SIG_MADT 0x43495041u /* "APIC" */

#define ACPI_SDT_HEADER_SIZE 36

/*
 * the MADT header is 8 bytes longer than a plain SDT header: it appends
 * LocalAPICAddress (4), PCATCompatibility (1), PCATCompatCount (1) and
 * Reserved (2). Structure entries therefore begin at 44, not 36.
 */
#define ACPI_MADT_HEADER_SIZE 44

#define ACPI_RSDT_ENTRY_SIZE 8

/* Local APIC structure: Type, Length, ACPI ID, APIC ID, Flags(4). */
#define MADT_LOCAL_APIC_FLAGS_OFFSET 4
#define MADT_LOCAL_APIC_ENABLED 0x1

/*
 * x2APIC structure: Type, Length, Reserved(2), ACPI UID(4), APIC ID(4),
 * Flags(4), ... so the APIC ID is a full 32-bit field at offset 8 and the
 * enabled flag sits at offset 12. the entry is 16 bytes minimum.
 */
#define MADT_X2APIC_ID_OFFSET 8
#define MADT_X2APIC_FLAGS_OFFSET 12
#define MADT_X2APIC_ENABLED 0x1
#define MADT_X2APIC_MIN_LEN 16

static uint32_t read_u32(const uint8_t *p) {
  return (uint32_t)p[0] | ((uint32_t)p[1] << 8) | ((uint32_t)p[2] << 16) | ((uint32_t)p[3] << 24);
}

/* EFI_GUID is { u32, u16, u16, u8[8] } with no padding, so its memory image is
   exactly the 16-byte wire form we compare against. */
static bool guid_equal(const EFI_GUID *guid, const uint8_t *bytes) {
  const uint8_t *p = (const uint8_t *)guid;
  int i;

  for (i = 0; i < 16; i++) {
    if (p[i] != bytes[i]) {
      return false;
    }
  }
  return true;
}

/* an SDT's checksum is the sum of all its bytes, modulo 256, taken to zero. */
static bool checksum_ok(const uint8_t *table, uint32_t len) {
  uint8_t sum = 0;
  uint32_t i;

  for (i = 0; i < len; i++) {
    sum = (uint8_t)(sum + table[i]);
  }
  return sum == 0;
}

static void madt_parse(const uint8_t *madt, uint32_t madt_len, struct acpi_madt_info *out) {
  const uint8_t *entry = madt + ACPI_MADT_HEADER_SIZE;
  const uint8_t *end = madt + madt_len;

  while (entry + 2 <= end) {
    uint8_t type = entry[0];
    uint8_t len = entry[1];

    /* length 0 would not advance us and could not be a valid structure. */
    if (len < 2 || entry + len > end) {
      break;
    }

    if ((type == MADT_TYPE_LOCAL_APIC && len >= 8) ||
        (type == MADT_TYPE_X2APIC && len >= MADT_X2APIC_MIN_LEN)) {
      bool is_x2apic = (type == MADT_TYPE_X2APIC);
      uint32_t flags = read_u32(entry + (is_x2apic ? MADT_X2APIC_FLAGS_OFFSET
                                                  : MADT_LOCAL_APIC_FLAGS_OFFSET));
      uint32_t enabled_bit = is_x2apic ? MADT_X2APIC_ENABLED : MADT_LOCAL_APIC_ENABLED;

      /*
       * bit 0 of flags is "enabled"; skip CPUs the firmware marked unusable.
       *
       * the two entry types are handled together because a MADT uses one or the
       * other, never both, and reading a type 0 entry as type 9 would pick up a
       * 32-bit field spanning the flags and the following entry's header.
       */
      if ((flags & enabled_bit) != 0 && out->cpu_count < ACPI_MAX_APIC_IDS) {
        out->apic_ids[out->cpu_count] =
            is_x2apic ? read_u32(entry + MADT_X2APIC_ID_OFFSET) : (uint32_t)entry[3];

        /*
         * the slot is dense rather than the APIC ID, because an x2APIC ID may
         * exceed the per-CPU array we reserve. the BSP passes the slot to the
         * trampoline, so the AP never has to index anything by ID itself.
         */
        out->slots[out->cpu_count] = (uint8_t)out->cpu_count;
        out->cpu_count++;
      }
    } else if (type == MADT_TYPE_IO_APIC && len >= 12) {
      out->ioapic_count++;
    }

    entry += len;
  }
}

/*
 * locate a table by the signature of the table itself rather than by the
 * Type byte recorded in the RSDT entry. it costs one extra read and needs no
 * trust in the entry metadata, which some firmware images get wrong.
 */
static const uint8_t *rsdt_find_table(const uint8_t *rsdt, uint32_t rsdt_len, uint32_t sig) {
  const uint8_t *entry = rsdt + ACPI_SDT_HEADER_SIZE;
  const uint8_t *end = rsdt + rsdt_len;

  while (entry + ACPI_RSDT_ENTRY_SIZE <= end) {
    const uint8_t *t = (const uint8_t *)(uint64_t)read_u32(entry + 4);

    if (t != NULL && read_u32(t) == sig) {
      return t;
    }
    entry += ACPI_RSDT_ENTRY_SIZE;
  }
  return NULL;
}

bool acpi_madt_parse(const EFI_SYSTEM_TABLE *st, struct acpi_madt_info *out) {
  const uint8_t *rsdp = NULL;
  const uint8_t *rsdt;
  const uint8_t *madt;
  uint32_t rsdt_len;
  uint32_t madt_len;
  UINTN i;

  out->cpu_count = 0;
  out->ioapic_count = 0;

  /* some firmware installs the MADT directly in the UEFI configuration table. */
  for (i = 0; i < st->NumberOfTableEntries; i++) {
    const uint8_t *t = (const uint8_t *)st->ConfigurationTable[i].VendorTable;

    if (t != NULL && read_u32(t) == ACPI_SIG_MADT) {
      madt_len = read_u32(t + 4);
      if (madt_len >= ACPI_MADT_HEADER_SIZE && checksum_ok(t, madt_len)) {
        madt_parse(t, madt_len, out);
        return out->cpu_count > 0;
      }
    }
  }

  for (i = 0; i < st->NumberOfTableEntries; i++) {
    const EFI_CONFIGURATION_TABLE *ct = &st->ConfigurationTable[i];

    if (guid_equal(&ct->VendorGuid, guid_rsdp_v2) || guid_equal(&ct->VendorGuid, guid_rsdp_v1)) {
      rsdp = (const uint8_t *)ct->VendorTable;
      break;
    }
  }

  if (rsdp == NULL || read_u32(rsdp) != ACPI_SIG_RSDP) {
    return false;
  }
  /* signature is 8 bytes: "RSD PTR " -> R(0) S(1) D(2) ' '(3) P(4) T(5) R(6) ' '(7). */
  if (rsdp[4] != 'P' || rsdp[5] != 'T' || rsdp[6] != 'R' || rsdp[7] != ' ') {
    return false;
  }

  /* RsdtPhysicalAddress is a 32-bit field at offset 16. */
  rsdt = (const uint8_t *)(uint64_t)read_u32(rsdp + 16);
  if (rsdt == NULL || read_u32(rsdt) != ACPI_SIG_RSDT) {
    return false;
  }

  rsdt_len = read_u32(rsdt + 4);
  if (rsdt_len < ACPI_SDT_HEADER_SIZE || !checksum_ok(rsdt, rsdt_len)) {
    return false;
  }

  madt = rsdt_find_table(rsdt, rsdt_len, ACPI_SIG_MADT);
  if (madt == NULL) {
    return false;
  }

  madt_len = read_u32(madt + 4);
  if (madt_len < ACPI_MADT_HEADER_SIZE || !checksum_ok(madt, madt_len)) {
    return false;
  }

  madt_parse(madt, madt_len, out);
  return out->cpu_count > 0;
}

uint32_t acpi_madt_ap_count(const struct acpi_madt_info *madt) {
  /*
   * one of the listed processors is the BSP, which is already running. the
   * subtraction is guarded rather than written as `cpu_count - 1`: a MADT with
   * no enabled processor entries would otherwise wrap to UINT32_MAX and send
   * the bring-up loop looking for four billion APs.
   */
  return (madt->cpu_count > 0) ? (madt->cpu_count - 1) : 0;
}