/*
  StarForth — Steady-State Virtual Machine Runtime

  Copyright (c) 2023–2025 Robert A. James
  All rights reserved.

  This file is part of the StarForth project.

  Licensed under the StarForth License, Version 1.0 (the "License");
  you may not use this file except in compliance with the License.
*/

/**
 * smp.c (amd64) - secondary-core bring-up, FABRIC-3.10.8 MC-1a.
 *
 * Cores come from the MADT's enabled Local APIC entries. Each is started
 * with INIT, 10 ms, STARTUP, 200 us, STARTUP into ap_trampoline.S, copied
 * into a conventional-memory page below 1 MiB taken from the UEFI memory
 * map. The trampoline enters long mode on the boot core's own control
 * state and calls sk_smp_secondary_main() with its GS base already set.
 *
 * Limit (recorded for SER5): the trampoline loads CR3 in 32-bit mode, so
 * the page tables must lie below 4 GiB; a start is refused otherwise.
 */

#include "starkernel/smp.h"
#include "starkernel/apic.h"
#include "starkernel/timer.h"
#include "starkernel/console.h"
#include "uefi.h"

extern char sk_ap_tramp_start[], sk_ap_tramp_end[];
extern char sk_ap_tr_pm_jmp[], sk_ap_tr_lm_jmp[], sk_ap_tr_pm32[], sk_ap_tr_lm64[];
extern char sk_ap_tr_gdt[], sk_ap_tr_gdtr[], sk_ap_tr_data[];

/* The trampoline's data block (ap_trampoline.S, D_* offsets). */
typedef struct __attribute__((packed)) {
    uint64_t cr0, cr3, cr4_boot, cr4, efer, pat, xcr0, xcr0_valid;
    uint64_t percpu, stack, entry;
    uint16_t kgdt_limit; uint64_t kgdt_base;
    uint16_t kidt_limit; uint64_t kidt_base;
} ap_tramp_data_t;

#define CR4_PCIDE    (1ull << 17)
#define CR4_OSXSAVE  (1ull << 18)

#define ICR_INIT     0x00004500u  /* INIT, level assert */
#define ICR_STARTUP  0x00004600u  /* STARTUP; low byte = vector (page number) */

/* ---- MADT ---------------------------------------------------------------- */

typedef struct __attribute__((packed)) {
    uint8_t  signature[8];
    uint8_t  checksum;
    uint8_t  oem_id[6];
    uint8_t  revision;
    uint32_t rsdt_address;
    uint32_t length;
    uint64_t xsdt_address;
} smp_rsdp_t;

typedef struct __attribute__((packed)) {
    uint8_t  signature[4];
    uint32_t length;
    uint8_t  rest[28];
} smp_sdt_t;

static int sig_eq(const uint8_t *p, const char *s, int n)
{
    int i;
    for (i = 0; i < n; i++) if (p[i] != (uint8_t)s[i]) return 0;
    return 1;
}

static const uint8_t *find_madt(void *rsdp_ptr)
{
    const smp_rsdp_t *r = (const smp_rsdp_t *)rsdp_ptr;
    const smp_sdt_t  *xsdt;
    const uint64_t   *ent;
    uint32_t n, i;

    if (!r || !sig_eq(r->signature, "RSD PTR ", 8) || r->revision < 2 || !r->xsdt_address) return 0;
    xsdt = (const smp_sdt_t *)(uintptr_t)r->xsdt_address;
    if (!sig_eq(xsdt->signature, "XSDT", 4) || xsdt->length <= 36u) return 0;
    n   = (xsdt->length - 36u) / 8u;
    ent = (const uint64_t *)((const uint8_t *)xsdt + 36u);
    for (i = 0; i < n; i++) {
        const smp_sdt_t *sdt = (const smp_sdt_t *)(uintptr_t)ent[i];
        if (sdt && sig_eq(sdt->signature, "APIC", 4)) return (const uint8_t *)sdt;
    }
    return 0;
}

static int id_seen(const uint64_t *ids, uint32_t n, uint64_t id)
{
    uint32_t i;
    for (i = 0; i < n; i++) if (ids[i] == id) return 1;
    return 0;
}

/* Enabled type-0 (Local APIC) entries, then type-9 (x2APIC) entries whose
 * id fits xAPIC (<= 255) and was not already listed. */
uint32_t arch_smp_enumerate(BootInfo *boot_info, uint64_t *hw_ids, uint32_t cap)
{
    const uint8_t *madt = find_madt(boot_info ? boot_info->acpi_table : (void *)0);
    const uint8_t *p, *end;
    uint32_t n = 0, pass;

    if (!madt) {
        if (cap) hw_ids[0] = apic_id();
        return 1;
    }
    end = madt + ((const smp_sdt_t *)madt)->length;
    for (pass = 0; pass < 2; pass++) {
        for (p = madt + 44; p + 2 <= end; p += p[1]) {
            uint64_t id;
            uint32_t flags;
            if (p[1] == 0 || p + p[1] > end) break;
            if (pass == 0 && p[0] == 0 && p[1] >= 8) {
                id    = p[3];
                flags = (uint32_t)p[4] | ((uint32_t)p[5] << 8) | ((uint32_t)p[6] << 16) | ((uint32_t)p[7] << 24);
            } else if (pass == 1 && p[0] == 9 && p[1] >= 16) {
                id    = (uint64_t)p[4] | ((uint64_t)p[5] << 8) | ((uint64_t)p[6] << 16) | ((uint64_t)p[7] << 24);
                flags = (uint32_t)p[8] | ((uint32_t)p[9] << 8) | ((uint32_t)p[10] << 16) | ((uint32_t)p[11] << 24);
                if (id > 255u) continue;
            } else {
                continue;
            }
            if (!(flags & 1u)) continue; /* not enabled */
            if (cap) {
                if (id_seen(hw_ids, n < cap ? n : cap, id)) continue;
                if (n < cap) hw_ids[n] = id;
            }
            n++;
        }
    }
    if (n == 0) {
        if (cap) hw_ids[0] = apic_id();
        return 1;
    }
    return n;
}

uint64_t arch_smp_boot_hw_id(void) { return apic_id(); }

void arch_smp_park_once(void)
{
    __asm__ volatile ("cli; hlt" ::: "memory");
}

/* ---- the trampoline page ------------------------------------------------- */

static uint64_t g_tramp_page; /* 0 = not chosen yet */

/* The lowest conventional page in [0x1000, 1 MiB) whose whole 4 KiB is
 * free in the UEFI memory map. */
static uint64_t pick_tramp_page(BootInfo *bi)
{
    uint64_t best = 0;
    UINTN    n, i;
    if (!bi || !bi->memory_map || !bi->memory_map_descriptor_size) return 0;
    n = bi->memory_map_size / bi->memory_map_descriptor_size;
    for (i = 0; i < n; i++) {
        const EFI_MEMORY_DESCRIPTOR *d = (const EFI_MEMORY_DESCRIPTOR *)
            ((const uint8_t *)bi->memory_map + i * bi->memory_map_descriptor_size);
        uint64_t start, end;
        if (d->Type != EfiConventionalMemory) continue;
        start = d->PhysicalStart;
        end   = start + d->NumberOfPages * 4096ull;
        if (start < 0x1000ull) start = 0x1000ull;
        start = (start + 0xFFFull) & ~0xFFFull;
        if (start + 4096ull > end || start + 4096ull > 0x100000ull) continue;
        if (!best || start < best) best = start;
    }
    return best;
}

static uint64_t rdmsr64(uint32_t msr)
{
    uint32_t lo, hi;
    __asm__ volatile ("rdmsr" : "=a"(lo), "=d"(hi) : "c"(msr));
    return ((uint64_t)hi << 32) | lo;
}

static void wait_us(uint64_t us)
{
    uint64_t hz = timer_tsc_hz(), t0 = heartbeat_read_counter();
    uint64_t budget = hz ? (hz / 1000000u) * us : 0;
    uint64_t spin = us * 1000u; /* only if the counter's rate is unknown */
    while (hz ? (heartbeat_read_counter() - t0 < budget) : (spin-- > 0)) {
        __asm__ volatile ("pause");
    }
}

int arch_smp_start(BootInfo *boot_info, sk_percpu_t *pc, void *stack_top)
{
    uint8_t         *page;
    size_t           len = (size_t)(sk_ap_tramp_end - sk_ap_tramp_start), k;
    ap_tramp_data_t *d;
    uint64_t         cr0, cr3, cr4;
    struct __attribute__((packed)) { uint16_t limit; uint64_t base; } gdtr, idtr;
    uint32_t         T;

    if (pc->hw_id > 255u || len > 4096u) return -1;
    if (!g_tramp_page) g_tramp_page = pick_tramp_page(boot_info);
    if (!g_tramp_page) {
        console_println("SMP: no conventional page below 1 MiB for the trampoline");
        return -1;
    }
    __asm__ volatile ("mov %%cr0, %0" : "=r"(cr0));
    __asm__ volatile ("mov %%cr3, %0" : "=r"(cr3));
    __asm__ volatile ("mov %%cr4, %0" : "=r"(cr4));
    if (cr3 >> 32) {
        console_println("SMP: page tables above 4 GiB -- the trampoline cannot load CR3");
        return -1;
    }
    __asm__ volatile ("sgdt %0" : "=m"(gdtr));
    __asm__ volatile ("sidt %0" : "=m"(idtr));

    T    = (uint32_t)g_tramp_page;
    page = (uint8_t *)(uintptr_t)g_tramp_page;
    for (k = 0; k < len; k++) page[k] = (uint8_t)sk_ap_tramp_start[k];

    *(uint32_t *)(void *)(page + (sk_ap_tr_pm_jmp - sk_ap_tramp_start)) =
        T + (uint32_t)(sk_ap_tr_pm32 - sk_ap_tramp_start);
    *(uint32_t *)(void *)(page + (sk_ap_tr_lm_jmp - sk_ap_tramp_start)) =
        T + (uint32_t)(sk_ap_tr_lm64 - sk_ap_tramp_start);
    *(uint32_t *)(void *)(page + (sk_ap_tr_gdtr - sk_ap_tramp_start) + 2) =
        T + (uint32_t)(sk_ap_tr_gdt - sk_ap_tramp_start);

    d = (ap_tramp_data_t *)(void *)(page + (sk_ap_tr_data - sk_ap_tramp_start));
    d->cr0        = cr0;
    d->cr3        = cr3;
    d->cr4        = cr4;
    d->cr4_boot   = cr4 & ~CR4_PCIDE;
    d->efer       = rdmsr64(0xC0000080u) & ~(1ull << 10); /* LMA is the CPU's to set */
    d->pat        = rdmsr64(0x277u);
    d->xcr0_valid = (cr4 & CR4_OSXSAVE) ? 1u : 0u;
    d->xcr0       = 0;
    if (d->xcr0_valid) {
        uint32_t lo, hi;
        __asm__ volatile ("xgetbv" : "=a"(lo), "=d"(hi) : "c"(0));
        d->xcr0 = ((uint64_t)hi << 32) | lo;
    }
    d->percpu     = (uint64_t)(uintptr_t)pc;
    d->stack      = (uint64_t)(uintptr_t)stack_top;
    d->entry      = (uint64_t)(uintptr_t)&sk_smp_secondary_main;
    d->kgdt_limit = gdtr.limit;  d->kgdt_base = gdtr.base;
    d->kidt_limit = idtr.limit;  d->kidt_base = idtr.base;
    __asm__ volatile ("mfence" ::: "memory");

    if (apic_send_ipi((uint8_t)pc->hw_id, ICR_INIT) != 0) {
        console_println("SMP: INIT IPI not delivered");
        return -1;
    }
    wait_us(10000u);
    (void)apic_send_ipi((uint8_t)pc->hw_id, ICR_STARTUP | (T >> 12));
    wait_us(200u);
    if (!__atomic_load_n(&pc->online, __ATOMIC_ACQUIRE)) {
        (void)apic_send_ipi((uint8_t)pc->hw_id, ICR_STARTUP | (T >> 12));
    }
    return 0;
}