#include "ioapic.h"
#include "apic.h"
#include "mm/vmm.h"
#include "mm/memory.h"
#include "libk/stdio.h"
#include <uacpi/tables.h>   /* uacpi_table_find_by_signature / uacpi_table_unref */
#include "pic.h"


/* ══════════════════════════════════════════════════════════════════════════
 *  ACPI MADT structures
 *  (Defined locally so we don't depend on uACPI's internal acpi.h layout.)
 * ══════════════════════════════════════════════════════════════════════════ */

extern const PICDriver* g_Driver; // old pic driver

bool g_IOAPIC = false; // IOAPIC enabled

typedef struct {
    uint8_t  signature[4];  /* "APIC" */
    uint32_t length;
    uint8_t  revision;
    uint8_t  checksum;
    uint8_t  oem_id[6];
    uint8_t  oem_table_id[8];
    uint32_t oem_revision;
    uint32_t creator_id;
    uint32_t creator_revision;
    /* --- MADT-specific -------------------------------------------------- */
    uint32_t local_apic_addr;  /* Default LAPIC physical address (32-bit) */
    uint32_t flags;            /* bit 0 = dual 8259 PICs present */
    /* followed by variable-length Interrupt Controller Structure entries  */
} __attribute__((packed)) madt_t;

typedef struct {
    uint8_t type;
    uint8_t length;
} __attribute__((packed)) madt_entry_hdr_t;

/* Type 0: Processor Local APIC */
typedef struct {
    madt_entry_hdr_t hdr;
    uint8_t  uid;
    uint8_t  apic_id;
    uint32_t flags;  /* bit 0 = enabled */
} __attribute__((packed)) madt_lapic_t;

/* Type 1: I/O APIC */
typedef struct {
    madt_entry_hdr_t hdr;
    uint8_t  id;
    uint8_t  reserved;
    uint32_t address;   /* Physical MMIO base */
    uint32_t gsi_base;  /* First GSI handled by this IOAPIC */
} __attribute__((packed)) madt_ioapic_t;

/* Type 2: Interrupt Source Override */
typedef struct {
    madt_entry_hdr_t hdr;
    uint8_t  bus;     /* 0 = ISA */
    uint8_t  source;  /* ISA IRQ number */
    uint32_t gsi;     /* Remapped GSI */
    uint16_t flags;   /* bits [1:0] = polarity, bits [3:2] = trigger mode */
} __attribute__((packed)) madt_iso_t;

/* Type 4: Local APIC NMI */
typedef struct {
    madt_entry_hdr_t hdr;
    uint8_t  uid;    /* 0xFF = all processors */
    uint16_t flags;
    uint8_t  lint;   /* 0 or 1 */
} __attribute__((packed)) madt_nmi_t;

/* ══════════════════════════════════════════════════════════════════════════
 *  Internal state
 * ══════════════════════════════════════════════════════════════════════════ */

typedef struct {
    volatile uint32_t *base;    /* Virtual address of IOAPIC MMIO */
    uint32_t           gsi_base;
    uint32_t           gsi_count;   /* Number of redirection entries */
} ioapic_t;

static ioapic_t g_ioapics[IOAPIC_MAX];
static int      g_ioapic_count = 0;

/* ISA interrupt source overrides (max 16 ISA IRQs) */
typedef struct {
    bool     present;
    uint32_t gsi;
    bool     active_low;
    bool     level;
} iso_t;
static iso_t g_iso[16];   /* indexed by ISA IRQ (source) */

/* ══════════════════════════════════════════════════════════════════════════
 *  Low-level IOAPIC MMIO access
 *
 *  The IOAPIC has two MMIO registers:
 *    base+0x00  IOREGSEL  – index register (write which register to access)
 *    base+0x10  IOWIN     – data window  (read/write the selected register)
 * ══════════════════════════════════════════════════════════════════════════ */

static uint32_t ioapic_read(const ioapic_t *io, uint8_t reg) {
    *((volatile uint32_t *)(io->base + (0x00 >> 2))) = reg;
    return *((volatile uint32_t *)(io->base + (0x10 >> 2)));
}

static void ioapic_write(const ioapic_t *io, uint8_t reg, uint32_t val) {
    *((volatile uint32_t *)(io->base + (0x00 >> 2))) = reg;
    *((volatile uint32_t *)(io->base + (0x10 >> 2))) = val;
}

/* ── Find which IOAPIC owns a given GSI ─────────────────────────────────── */

static ioapic_t *ioapic_for_gsi(uint32_t gsi) {
    for (int i = 0; i < g_ioapic_count; i++) {
        ioapic_t *io = &g_ioapics[i];
        if (gsi >= io->gsi_base && gsi < io->gsi_base + io->gsi_count)
            return io;
    }
    return NULL;
}

/* ══════════════════════════════════════════════════════════════════════════
 *  Public API
 * ══════════════════════════════════════════════════════════════════════════ */

uint32_t ioapic_gsi_for_isa_irq(uint8_t isa_irq) {
    if (isa_irq < 16 && g_iso[isa_irq].present)
        return g_iso[isa_irq].gsi;
    return isa_irq;
}

void ioapic_redirect(uint32_t gsi, uint8_t vector, uint8_t dest_lapic,
                     bool active_low, bool level, bool masked) {
    ioapic_t *io = ioapic_for_gsi(gsi);
    if (!io) {
        printf("[IOAPIC] No IOAPIC for GSI %u\n", gsi);
        return;
    }

    uint8_t  idx   = (uint8_t)(gsi - io->gsi_base);
    uint32_t lo    = (uint32_t)vector
                   | IOAPIC_RTE_DM_FIXED
                   | (active_low ? IOAPIC_RTE_ACTIVE_LOW : 0)
                   | (level      ? IOAPIC_RTE_LEVEL       : 0)
                   | (masked     ? IOAPIC_RTE_MASKED       : 0);
    uint32_t hi    = (uint32_t)dest_lapic << 24;

    /* Write high word first, then low (avoids momentary spurious delivery) */
    ioapic_write(io, IOAPIC_REDTBL_HI(idx), hi);
    ioapic_write(io, IOAPIC_REDTBL_LO(idx), lo);
}

void ioapic_mask_gsi(uint32_t gsi) {
    ioapic_t *io = ioapic_for_gsi(gsi);
    if (!io) return;
    uint8_t  idx = (uint8_t)(gsi - io->gsi_base);
    uint32_t lo  = ioapic_read(io, IOAPIC_REDTBL_LO(idx));
    ioapic_write(io, IOAPIC_REDTBL_LO(idx), lo | IOAPIC_RTE_MASKED);
}

void ioapic_unmask_gsi(uint32_t gsi) {
    ioapic_t *io = ioapic_for_gsi(gsi);
    if (!io) return;
    uint8_t  idx = (uint8_t)(gsi - io->gsi_base);
    uint32_t lo  = ioapic_read(io, IOAPIC_REDTBL_LO(idx));
    ioapic_write(io, IOAPIC_REDTBL_LO(idx), lo & ~IOAPIC_RTE_MASKED);
}

/* ══════════════════════════════════════════════════════════════════════════
 *  Initialisation
 * ══════════════════════════════════════════════════════════════════════════ */

void ioapic_init(void) {
    /*
     * ── Step 1: Find the MADT via uACPI ──────────────────────────────────
     *
     * The MADT signature in ACPI is "APIC" (not "MADT").
     * uacpi_table_find_by_signature returns a handle whose .ptr field points
     * to the raw table data in (HHDM-mapped) physical memory.
     */
    struct uacpi_table madt_table;
    uacpi_status st = uacpi_table_find_by_signature("APIC", &madt_table);
    if (uacpi_unlikely_error(st)) {
        printf("[IOAPIC] Could not find MADT: %s\n", uacpi_status_to_string(st));
        return;
    }

    madt_t *madt = (madt_t *)madt_table.ptr;
    printf("[IOAPIC] MADT @ virt 0x%lx  len=%u  lapic_addr=0x%08x\n",
           (uint64_t)madt, madt->length, madt->local_apic_addr);

    /*
     * ── Step 2: Walk the MADT entry list ─────────────────────────────────
     *
     * Entries start immediately after the fixed 44-byte MADT header and run
     * until madt->length bytes from the table base.
     */
    const uint8_t *entry_ptr = (const uint8_t *)madt + sizeof(madt_t);
    const uint8_t *madt_end  = (const uint8_t *)madt + madt->length;

    while (entry_ptr < madt_end) {
        const madt_entry_hdr_t *hdr = (const madt_entry_hdr_t *)entry_ptr;

        if (hdr->length < 2) {
            printf("[IOAPIC] MADT entry with length < 2, stopping parse\n");
            break;
        }

        switch (hdr->type) {

        /* ── Type 0: Processor Local APIC ──────────────────────────────── */
        case 0: {
            const madt_lapic_t *e = (const madt_lapic_t *)entry_ptr;
            printf("[IOAPIC] MADT[0] Processor UID=%u  LAPIC ID=%u  flags=0x%x%s\n",
                   e->uid, e->apic_id, e->flags,
                   (e->flags & 1) ? "" : " (disabled)");
            break;
        }

        /* ── Type 1: I/O APIC ───────────────────────────────────────────── */
        case 1: {
            const madt_ioapic_t *e = (const madt_ioapic_t *)entry_ptr;
            if (g_ioapic_count >= IOAPIC_MAX) {
                printf("[IOAPIC] Too many IOAPICs, skipping ID=%u\n", e->id);
                break;
            }

            ioapic_t *io  = &g_ioapics[g_ioapic_count];
            io->gsi_base  = e->gsi_base;

            /*
             * Map the IOAPIC MMIO page.  Like the LAPIC, the HHDM covers the
             * IOAPIC's physical address (typically 0xFEC00000) so we just add
             * MEM_PHYS_OFFSET.  Two MMIO registers are accessed (offsets 0 and
             * 0x10) so one 4 KiB page is sufficient.
             *
             * TODO: Mark the page UC (cache-disable) in the PTE when VMM
             *       gains support for PAT / PCD flags.
             */
            uint64_t phys = (uint64_t)e->address;
            uintptr_t virt = (uintptr_t)phys + MEM_PHYS_OFFSET;
            /* The HHDM loop in vmm_init already covered this range; if not,
             * uncomment the explicit map call below:                         */
            /* vmm_map_page(kernel_pagemap, virt, phys,
                            PAGE_READ | PAGE_WRITE | PAGE_NO_EXECUTE, Size4KiB); */
            io->base = (volatile uint32_t *)virt;

            /* Read version register to discover number of redirection entries */
            uint32_t ver = ioapic_read(io, IOAPIC_REG_VER);
            io->gsi_count = ((ver >> 16) & 0xFF) + 1;  /* bits [23:16] = max entry index */

            printf("[IOAPIC] MADT[1] ID=%u  phys=0x%08x  GSI base=%u  entries=%u\n",
                   e->id, e->address, io->gsi_base, io->gsi_count);

            /*
             * ── Mask every redirection entry ─────────────────────────────
             *
             * We keep all entries masked at boot time.  Legacy ISA IRQs are
             * handled by the i8259 → LAPIC LINT0 ExtINT path, not via the
             * IOAPIC.  PCI devices can be connected later with ioapic_redirect().
             */
            for (uint32_t n = 0; n < io->gsi_count; n++) {
                ioapic_write(io, IOAPIC_REDTBL_HI(n), 0);
                ioapic_write(io, IOAPIC_REDTBL_LO(n), IOAPIC_RTE_MASKED | 0xFF);
            }

            g_ioapic_count++;
            break;
        }

        /* ── Type 2: Interrupt Source Override ──────────────────────────── */
        case 2: {
            const madt_iso_t *e = (const madt_iso_t *)entry_ptr;

            /* flags bits [1:0]: 00/11 = conforms to bus (ISA = active high edge)
             *                   01    = active high,  11 = active low
             * flags bits [3:2]: 00/11 = conforms to bus (ISA = edge)
             *                   01    = edge,          11 = level             */
            bool active_low = ((e->flags & 0x3) == 3);
            bool level      = ((e->flags >> 2) & 0x3) == 3;

            printf("[IOAPIC] MADT[2] ISA IRQ %u -> GSI %u  %s  %s\n",
                   e->source, e->gsi,
                   active_low ? "active-low"  : "active-high",
                   level      ? "level"        : "edge");

            if (e->source < 16) {
                g_iso[e->source].present    = true;
                g_iso[e->source].gsi        = e->gsi;
                g_iso[e->source].active_low = active_low;
                g_iso[e->source].level      = level;
            }
            break;
        }

        /* ── Type 4: Local APIC NMI ─────────────────────────────────────── */
        case 4: {
            const madt_nmi_t *e = (const madt_nmi_t *)entry_ptr;
            printf("[IOAPIC] MADT[4] NMI  UID=0x%02x  LINT%u\n",
                   e->uid, e->lint);
            break;
        }

        default:
            printf("[IOAPIC] MADT entry type=%u length=%u (skipped)\n",
                   hdr->type, hdr->length);
            break;
        }

        entry_ptr += hdr->length;
    }

    uacpi_table_unref(&madt_table);

    printf("[IOAPIC] Init done: %d IOAPIC(s) found\n", g_ioapic_count);
}
void irq_redirect_to_apic(uint8_t isa_irq, uint8_t vector,
                          uint8_t dest_lapic, bool masked)
{
    if (isa_irq >= 16) {
        printf("[IRQ] irq_redirect_to_apic: ISA IRQ %u out of range\n", isa_irq);
        return;
    }

    uint32_t gsi = ioapic_gsi_for_isa_irq(isa_irq);

    /* Get polarity/trigger from MADT ISO or use ISA defaults */
    bool active_low = false;
    bool level      = false;

    if (isa_irq < 16 && g_iso[isa_irq].present) {
        active_low = g_iso[isa_irq].active_low;
        level      = g_iso[isa_irq].level;
    } else {
        /* Standard ISA: active-high, edge-triggered (except some like IRQ0/8) */
        if (isa_irq == 0 || isa_irq == 8) {
            level = true;           /* often level on real hardware */
        }
    }

    /* Mask in the 8259 so it stops firing through LINT0 */
    if (g_Driver) {
        g_Driver->Mask(isa_irq);
    }

    /* Programme IOAPIC redirection entry */
    ioapic_redirect(gsi, vector, dest_lapic,
                    active_low, level, masked);

    printf("[IRQ] Redirected ISA IRQ %u -> GSI %u  vector 0x%02x  LAPIC %u  %s %s\n",
           isa_irq, gsi, vector, dest_lapic,
           active_low ? "active-low" : "active-high",
           level      ? "level" : "edge");

    g_IOAPIC = true;
}