//! x86-64 architecture implementation
use core::{mem::size_of, ops::DerefMut, ptr::null_mut, sync::atomic::Ordering};

use abi::error::Error;
use acpi_lib::{mcfg::Mcfg, AcpiTables, InterruptModel};
use alloc::boxed::Box;
use device_api::{interrupt::Irq, Device};
use git_version::git_version;
use kernel_arch_x86_64::{
    mem::{
        init_fixed_tables,
        table::{PageAttributes, PageEntry, PageTable, L1, L2, L3},
        EarlyMapping, MEMORY_LIMIT, RAM_MAPPING_L1,
    },
    PerCpuData,
};
use kernel_fs::devfs;
use libk::{arch::Cpu, device::register_external_interrupt_controller};
use libk_device::register_monotonic_timestamp_provider;
use libk_mm::{
    address::{PhysicalAddress, Virtualize},
    phys::{self, reserved::reserve_region, PhysicalMemoryRegion},
    table::{EntryLevel, EntryLevelExt},
};
use libk_util::{sync::SpinFence, OneTimeInit};
use yboot_proto::{v1::AvailableMemoryRegion, LoadProtocolV1};
use ygg_driver_pci::PciBusManager;

mod acpi;
mod apic;
mod boot;
mod cpuid;
mod exception;
mod smp;
mod syscall;

use crate::{
    arch::x86::peripherals::{i8253::I8253, i8259::I8259, ps2::PS2, serial::ComPort},
    debug::{self, LogLevel},
    device::{
        self,
        display::{console, fb_console::FramebufferConsole, linear_fb::LinearFramebuffer},
    },
    fs::{Initrd, INITRD_DATA},
};

use self::{
    acpi::{AcpiAllocator, AcpiHandlerImpl},
    apic::{ioapic::IoApic, local::LocalApic},
    boot::BootData,
    cpuid::{ProcessorFeatures, PROCESSOR_FEATURES},
};

use super::{
    x86::{gdt, intrinsics},
    IpiMessage, Platform,
};

/// Offset where legacy ISA IRQs are remapped
pub const ISA_IRQ_OFFSET: u32 = apic::ioapic::ISA_IRQ_OFFSET;

/// x86-64 architecture implementation
pub struct X86_64 {
    boot_data: OneTimeInit<BootData>,
    acpi: OneTimeInit<AcpiTables<AcpiHandlerImpl>>,

    // Display
    framebuffer: OneTimeInit<LinearFramebuffer>,
    fbconsole: OneTimeInit<FramebufferConsole>,
}

static SHUTDOWN_FENCE: SpinFence = SpinFence::new();

/// Global x86-64 architecture value
pub static PLATFORM: X86_64 = X86_64 {
    boot_data: OneTimeInit::new(),
    acpi: OneTimeInit::new(),

    framebuffer: OneTimeInit::new(),
    fbconsole: OneTimeInit::new(),
};

impl Platform for X86_64 {
    const KERNEL_VIRT_OFFSET: usize = 0xFFFFFF8000000000;
    type L3 = kernel_arch_x86_64::mem::table::L3;

    unsafe fn start_application_processors(&self) {
        if let Some(acpi) = self.acpi.try_get() {
            let Some(pinfo) = acpi
                .platform_info_in(AcpiAllocator)
                .ok()
                .and_then(|p| p.processor_info)
            else {
                return;
            };

            smp::start_ap_cores(&pinfo);
        }
    }
}

impl X86_64 {
    unsafe fn handle_ipi(&self, _msg: IpiMessage) {
        warnln!("Received an IPI");
        todo!();
    }

    fn set_boot_data(&self, data: BootData) {
        match data {
            BootData::YBoot(data) => {
                // Reserve the memory map
                unsafe {
                    reserve_region(
                        "mmap",
                        PhysicalMemoryRegion {
                            base: PhysicalAddress::from_u64(data.memory_map.address),
                            size: data.memory_map.len as usize * size_of::<AvailableMemoryRegion>(),
                        },
                    );
                }

                // Reserve initrd, if not NULL
                if data.initrd_address != 0 && data.initrd_size != 0 {
                    let aligned_start = data.initrd_address & !0xFFF;
                    let aligned_end = (data.initrd_address + data.initrd_size + 0xFFF) & !0xFFF;

                    unsafe {
                        reserve_region(
                            "initrd",
                            PhysicalMemoryRegion {
                                base: PhysicalAddress::from_u64(aligned_start),
                                size: (aligned_end - aligned_start) as usize,
                            },
                        );
                    }
                }
            }
        }

        self.boot_data.init(data);
    }

    fn map_physical_memory<I: Iterator<Item = PhysicalMemoryRegion> + Clone>(
        it: I,
        _memory_start: PhysicalAddress,
        memory_end: PhysicalAddress,
    ) -> Result<(), Error> {
        let end_l1i = memory_end
            .into_usize()
            .page_align_up::<L1>()
            .page_index::<L1>();

        if end_l1i > 512 {
            todo!(
                "Cannot handle {}GiB of RAM",
                end_l1i * L1::SIZE / (1024 * 1024 * 1024)
            );
        }

        MEMORY_LIMIT.store(memory_end.into_usize(), Ordering::Release);

        // Check if 1GiB pages are supported
        if PROCESSOR_FEATURES
            .get()
            .contains(ProcessorFeatures::PDPE1GB)
        {
            // Just map gigabytes of RAM
            for l1i in 0..end_l1i {
                // TODO NX
                unsafe {
                    RAM_MAPPING_L1[l1i] = PageEntry::<L1>::block(
                        PhysicalAddress::from_usize(l1i * L1::SIZE),
                        PageAttributes::WRITABLE,
                    );
                }
            }
        } else {
            // Allocate the intermediate tables first
            let l2_tables_start = phys::find_contiguous_region(it, end_l1i)
                .expect("Could not allocate the memory for RAM mapping L2 tables");

            unsafe {
                reserve_region(
                    "ram-l2-tables",
                    PhysicalMemoryRegion {
                        base: l2_tables_start,
                        size: end_l1i * L3::SIZE,
                    },
                );
            }

            // Fill in the tables
            for l1i in 0..end_l1i {
                let l2_phys_addr = l2_tables_start.add(l1i * L3::SIZE);

                // Safety: ok, the mapping is done to the memory obtained from
                // find_contiguous_region()
                let mut l2_data =
                    unsafe { EarlyMapping::<[PageEntry<L2>; 512]>::map(l2_phys_addr)? };
                // Safety: ok, the slice comes from EarlyMapping of a page-aligned region
                let l2 = unsafe { PageTable::from_raw_slice_mut(l2_data.deref_mut()) };

                for l2i in 0..512 {
                    // TODO NX
                    l2[l2i] = PageEntry::<L2>::block(
                        PhysicalAddress::from_usize((l1i * L1::SIZE) | (l2i * L2::SIZE)),
                        PageAttributes::WRITABLE,
                    );
                }

                // Point the L1 entry to the L2 table
                unsafe {
                    RAM_MAPPING_L1[l1i] =
                        PageEntry::<L1>::table(l2_phys_addr, PageAttributes::WRITABLE)
                };

                intrinsics::flush_cpu_cache();
                // The EarlyMapping is then dropped
            }
        }

        Ok(())
    }

    unsafe fn init_physical_memory_from_yboot(data: &LoadProtocolV1) -> Result<(), Error> {
        let mmap = EarlyMapping::<AvailableMemoryRegion>::map_slice(
            PhysicalAddress::from_u64(data.memory_map.address),
            data.memory_map.len as usize,
        )?;

        phys::init_from_iter(
            mmap.as_ref().iter().map(|reg| PhysicalMemoryRegion {
                base: PhysicalAddress::from_u64(reg.start_address),
                size: reg.page_count as usize * L3::SIZE,
            }),
            Self::map_physical_memory,
        )
    }

    unsafe fn init_memory_management(&self) -> Result<(), Error> {
        init_fixed_tables();

        // Reserve lower 4MiB just in case
        reserve_region(
            "lowmem",
            PhysicalMemoryRegion {
                base: PhysicalAddress::ZERO,
                size: 4 * 1024 * 1024,
            },
        );

        match self.boot_data.get() {
            &BootData::YBoot(data) => Self::init_physical_memory_from_yboot(data)?,
        }

        Ok(())
    }

    unsafe fn init_platform(&'static self, cpu_id: usize) -> Result<(), Error> {
        let local_apic = Box::leak(Box::new(LocalApic::new()));
        let tss_address = gdt::init();

        let cpu_data = PerCpuData {
            this: null_mut(),
            tss_address,
            tmp_address: 0,
            local_apic,
        };
        cpuid::enable_features();
        Cpu::init_local(Some(cpu_id as _), cpu_data);
        syscall::init_syscall();

        if cpu_id == 0 {
            // Register the PCI drivers
            //  TODO make this implicit init
            ygg_driver_pci::register_class_driver(
                "NVMe Host Controller",
                0x01,
                Some(0x08),
                Some(0x02),
                ygg_driver_nvme::probe,
            );
            ygg_driver_pci::register_class_driver(
                "AHCI SATA Controller",
                0x01,
                Some(0x06),
                Some(0x01),
                ygg_driver_ahci::probe,
            );
            ygg_driver_pci::register_class_driver(
                "USB xHCI",
                0x0C,
                Some(0x03),
                Some(0x30),
                ygg_driver_usb_xhci::probe,
            );
            ygg_driver_pci::register_vendor_driver(
                "Virtio PCI Network Device",
                0x1AF4,
                0x1000,
                ygg_driver_virtio_net::probe,
            );

            match self.boot_data.get() {
                &BootData::YBoot(data) => {
                    let start = PhysicalAddress::from_u64(data.initrd_address);
                    Self::init_initrd(start, start.add(data.initrd_size as usize));
                }
            }

            self.init_acpi_from_boot_data()?;

            I8259.disable();

            register_monotonic_timestamp_provider(&I8253);

            let com1_3 = Box::leak(Box::new(ComPort::new(
                0x3F8,
                0x3E8,
                Irq::External(ISA_IRQ_OFFSET + 4),
            )));
            debug::add_sink(com1_3.port_a(), LogLevel::Debug);

            self.init_framebuffer()?;

            debug::init();

            infoln!(
                "Yggdrasil v{} ({})",
                env!("CARGO_PKG_VERSION"),
                git_version!()
            );

            PS2.init()?;

            if let Some(acpi) = self.acpi.try_get() {
                self.init_platform_from_acpi(acpi)?;
            }

            I8253.init_irq()?;
            // ps2.connect(self.tty.get());
            PS2.init_irq()?;

            device::register_device(&PS2);

            PciBusManager::setup_bus_devices()?;
        }

        Ok(())
    }

    unsafe fn init_acpi_from_boot_data(&self) -> Result<(), Error> {
        match self.boot_data.get() {
            &BootData::YBoot(data) => self.init_acpi_from_rsdp(data.rsdp_address as usize),
        }
    }

    unsafe fn init_acpi_from_rsdp(&self, rsdp: usize) -> Result<(), Error> {
        let acpi_tables = AcpiTables::from_rsdp(AcpiHandlerImpl, rsdp).map_err(|err| {
            errorln!("Could not initialize ACPI tables: {:?}", err);
            Error::InvalidArgument
        })?;
        self.acpi.init(acpi_tables);
        Ok(())
    }

    unsafe fn init_platform_from_acpi(
        &self,
        acpi: &'static AcpiTables<AcpiHandlerImpl>,
    ) -> Result<(), Error> {
        let platform_info = acpi.platform_info_in(AcpiAllocator).map_err(|err| {
            errorln!("Could not get ACPI platform info: {:?}", err);
            Error::InvalidArgument
        })?;

        let InterruptModel::Apic(apic_info) = platform_info.interrupt_model else {
            panic!("The processor does not support APIC");
        };

        let ioapic = IoApic::from_acpi(&apic_info)?;
        let ioapic = Box::leak(Box::new(ioapic));
        register_external_interrupt_controller(ioapic);

        // acpi::init_acpi(acpi).unwrap();

        if let Ok(mcfg) = acpi.find_table::<Mcfg>() {
            for entry in mcfg.entries() {
                PciBusManager::add_segment_from_mcfg(entry)?;
            }
        }

        Ok(())
    }

    unsafe fn init_framebuffer(&'static self) -> Result<(), Error> {
        match self.boot_data.get() {
            &BootData::YBoot(data) => {
                let info = &data.opt_framebuffer;

                self.framebuffer.init(LinearFramebuffer::from_physical_bits(
                    PhysicalAddress::from_u64(info.res_address),
                    info.res_size as usize,
                    info.res_stride as usize,
                    info.res_width,
                    info.res_height,
                )?);
            }
        }

        self.fbconsole.init(FramebufferConsole::from_framebuffer(
            self.framebuffer.get(),
            None,
        )?);
        debug::add_sink(self.fbconsole.get(), LogLevel::Info);

        // self.tty.init(CombinedTerminal::new(self.fbconsole.get()));

        devfs::add_named_block_device(self.framebuffer.get(), "fb0")?;
        // devfs::add_char_device(self.tty.get(), CharDeviceType::TtyRegular)?;
        console::add_console_autoflush(self.fbconsole.get());

        Ok(())
    }

    fn init_initrd(initrd_start: PhysicalAddress, initrd_end: PhysicalAddress) {
        if initrd_start.is_zero() || initrd_end <= initrd_start {
            infoln!("No initrd loaded");
            return;
        }

        let start_aligned = initrd_start.page_align_down::<L3>();
        let end_aligned = initrd_start.page_align_up::<L3>();

        let data = unsafe {
            core::slice::from_raw_parts(
                start_aligned.virtualize() as *const u8,
                initrd_end - initrd_start,
            )
        };

        let initrd = Initrd {
            phys_page_start: start_aligned,
            phys_page_len: end_aligned - start_aligned,
            data,
        };

        INITRD_DATA.init(initrd);
    }
}