pub mod dmar;
pub mod madt;
pub mod mcfg;

use x86_64::structures::paging::OffsetPageTable;

use crate::mem::addr;
use crate::mem::phys::BitmapFrameAllocator;
use crate::mem::virt;
use crate::static_vec::StaticVec;

const RSDP_SIGNATURE: &[u8; 8] = b"RSD PTR ";

#[repr(C, packed)]
struct Rsdp {
    signature: [u8; 8],
    checksum: u8,
    oem_id: [u8; 6],
    revision: u8,
    rsdt_address: u32,
}

#[repr(C, packed)]
struct Xsdp {
    rsdp: Rsdp,
    length: u32,
    xsdt_address: u64,
    extended_checksum: u8,
    _reserved: [u8; 3],
}

#[repr(C, packed)]
pub(crate) struct AcpiSdtHeader {
    pub signature: [u8; 4],
    pub length: u32,
    pub revision: u8,
    pub checksum: u8,
    pub oem_id: [u8; 6],
    pub oem_table_id: [u8; 8],
    pub oem_revision: u32,
    pub creator_id: u32,
    pub creator_revision: u32,
}

pub struct AcpiInfo {
    pub ioapic_addr: u64,
    #[allow(dead_code)]
    pub ioapic_gsi_base: u32,
    pub irq_overrides: StaticVec<IrqOverride, 16>,
    pub mcfg_entries: StaticVec<mcfg::McfgEntry, 4>,
    pub dmar_info: Option<dmar::DmarInfo>,
}

#[derive(Clone, Copy)]
pub struct IrqOverride {
    pub source_irq: u8,
    pub gsi: crate::arch::ioapic::Gsi,
    #[allow(dead_code)]
    pub flags: u16,
}

impl AcpiInfo {
    pub fn gsi_for_irq(&self, irq: u8) -> crate::arch::ioapic::Gsi {
        self.irq_overrides
            .iter()
            .find(|o| o.source_irq == irq)
            .map(|o| o.gsi)
            .unwrap_or(crate::arch::ioapic::Gsi::new(irq as u32))
    }
}

fn ensure_mapped(
    phys: u64,
    mapper: &mut OffsetPageTable,
    allocator: &mut BitmapFrameAllocator,
    hhdm_offset: u64,
) -> Result<(), crate::error::KernelError> {
    let page_aligned = phys & !0xFFF;
    virt::map_mmio(mapper, allocator, page_aligned, hhdm_offset)
}

struct XsdtEntry {
    table_phys: u64,
    sig: [u8; 4],
    length: u32,
    virt_addr: u64,
}

pub fn parse(
    rsdp_addr: u64,
    mapper: &mut OffsetPageTable,
    allocator: &mut BitmapFrameAllocator,
    hhdm_offset: u64,
) -> Option<AcpiInfo> {
    if ensure_mapped(rsdp_addr, mapper, allocator, hhdm_offset).is_err() {
        crate::kprintln!("  ACPI: failed to map RSDP page");
        return None;
    }

    let rsdp_virt = addr::phys_to_virt(x86_64::PhysAddr::new(rsdp_addr));
    let rsdp_ptr = rsdp_virt.as_ptr::<Rsdp>();

    let signature = unsafe { core::ptr::addr_of!((*rsdp_ptr).signature).read_unaligned() };
    if signature != *RSDP_SIGNATURE {
        crate::kprintln!("  ACPI: invalid RSDP signature");
        return None;
    }

    let rsdp_checksum = (0u8..20).fold(0u8, |sum, i| {
        sum.wrapping_add(unsafe { *(rsdp_virt.as_ptr::<u8>()).add(i as usize) })
    });
    if rsdp_checksum != 0 {
        crate::kprintln!("  ACPI: RSDP checksum invalid");
        return None;
    }

    let revision = unsafe { core::ptr::addr_of!((*rsdp_ptr).revision).read_unaligned() };
    if revision < 2 {
        crate::kprintln!("  ACPI: RSDP revision 0 (no XSDT)");
        return None;
    }

    let xsdp_checksum = (0u8..36).fold(0u8, |sum, i| {
        sum.wrapping_add(unsafe { *(rsdp_virt.as_ptr::<u8>()).add(i as usize) })
    });
    if xsdp_checksum != 0 {
        crate::kprintln!("  ACPI: XSDP extended checksum invalid");
        return None;
    }

    let xsdt_phys = {
        let xsdp_ptr = rsdp_virt.as_ptr::<Xsdp>();
        unsafe { core::ptr::addr_of!((*xsdp_ptr).xsdt_address).read_unaligned() }
    };

    if ensure_mapped(xsdt_phys, mapper, allocator, hhdm_offset).is_err() {
        crate::kprintln!("  ACPI: failed to map XSDT page");
        return None;
    }

    let xsdt_virt = addr::phys_to_virt(x86_64::PhysAddr::new(xsdt_phys));
    let xsdt_hdr_ptr = xsdt_virt.as_ptr::<AcpiSdtHeader>();
    let entries_offset = core::mem::size_of::<AcpiSdtHeader>();
    let xsdt_length =
        unsafe { core::ptr::addr_of!((*xsdt_hdr_ptr).length).read_unaligned() } as usize;

    if xsdt_length < entries_offset {
        crate::kprintln!("  ACPI: XSDT length too small");
        return None;
    }

    let xsdt_end_phys = xsdt_phys + xsdt_length as u64;
    let xsdt_start_page = xsdt_phys & !0xFFF;
    let xsdt_end_page = (xsdt_end_phys.saturating_sub(1)) & !0xFFF;
    let xsdt_map_ok = (xsdt_start_page..=xsdt_end_page)
        .step_by(4096)
        .skip(1)
        .try_fold((), |(), page| {
            ensure_mapped(page, mapper, allocator, hhdm_offset)
        });
    if xsdt_map_ok.is_err() {
        crate::kprintln!("  ACPI: failed to map XSDT span");
        return None;
    }

    let xsdt_checksum = (0..xsdt_length).fold(0u8, |sum, i| {
        sum.wrapping_add(unsafe { *(xsdt_virt.as_ptr::<u8>()).add(i) })
    });
    if xsdt_checksum != 0 {
        crate::kprintln!("  ACPI: XSDT checksum invalid");
        return None;
    }

    let entry_count = (xsdt_length - entries_offset) / 8;

    crate::kprintln!("  ACPI: XSDT at {:#x}, {} entries", xsdt_phys, entry_count);

    let entries_ptr = unsafe { (xsdt_virt.as_ptr::<u8>()).add(entries_offset) as *const u64 };

    let seed = AcpiInfo {
        ioapic_addr: 0,
        ioapic_gsi_base: 0,
        irq_overrides: StaticVec::new(),
        mcfg_entries: StaticVec::new(),
        dmar_info: None,
    };

    let result = (0..entry_count)
        .map(|i| unsafe { core::ptr::read_unaligned(entries_ptr.add(i)) })
        .filter_map(|table_phys| {
            ensure_mapped(table_phys, mapper, allocator, hhdm_offset).ok()?;
            let table_virt = addr::phys_to_virt(x86_64::PhysAddr::new(table_phys));
            let hdr_ptr = table_virt.as_ptr::<AcpiSdtHeader>();
            let sig = unsafe { core::ptr::addr_of!((*hdr_ptr).signature).read_unaligned() };
            let length = unsafe { core::ptr::addr_of!((*hdr_ptr).length).read_unaligned() };

            let table_end = table_phys + length as u64;
            let start_page = table_phys & !0xFFF;
            let end_page = (table_end.saturating_sub(1)) & !0xFFF;
            let span_ok = (start_page..=end_page)
                .step_by(4096)
                .skip(1)
                .try_fold((), |(), page| {
                    ensure_mapped(page, mapper, allocator, hhdm_offset)
                });
            span_ok.ok()?;

            Some(XsdtEntry {
                table_phys,
                sig,
                length,
                virt_addr: table_virt.as_u64(),
            })
        })
        .fold(seed, |acc, entry| match &entry.sig {
            b"APIC" => {
                crate::kprintln!("  ACPI: found MADT at {:#x}", entry.table_phys);
                let madt = madt::parse_madt(entry.virt_addr, entry.length);
                AcpiInfo {
                    ioapic_addr: madt.ioapic_addr,
                    ioapic_gsi_base: madt.ioapic_gsi_base,
                    irq_overrides: madt.irq_overrides,
                    ..acc
                }
            }
            b"MCFG" => {
                crate::kprintln!("  ACPI: found MCFG at {:#x}", entry.table_phys);
                let mcfg_entries = mcfg::parse_mcfg(entry.virt_addr, entry.length);
                AcpiInfo {
                    mcfg_entries,
                    ..acc
                }
            }
            b"DMAR" => {
                crate::kprintln!("  ACPI: found DMAR at {:#x}", entry.table_phys);
                let dmar = dmar::parse_dmar(entry.virt_addr, entry.length);
                AcpiInfo {
                    dmar_info: Some(dmar),
                    ..acc
                }
            }
            _ => acc,
        });

    match result.ioapic_addr {
        0 => {
            crate::kprintln!("  ACPI: no MADT found");
            None
        }
        _ => Some(result),
    }
}