//! x86-64 architecture implementation

use core::{ptr::null_mut, sync::atomic::Ordering};

use abi::error::Error;
use acpi::{mcfg::Mcfg, AcpiTables, HpetInfo, InterruptModel};
use alloc::sync::Arc;
use device_api::{
    device::Device,
    interrupt::{IpiDeliveryTarget, IpiMessage},
};
use kernel_arch::{Architecture, ArchitectureImpl};
use kernel_arch_x86::{
    cpuid::{self, CpuFeatures, EcxFeatures, EdxFeatures, ExtEdxFeatures},
    gdt,
};
use kernel_arch_x86_64::{mem, LocalApicInterface, PerCpuData, CPU_COUNT};
use libk::{
    arch::Cpu,
    config, debug,
    device::{
        display::{fb_console::FramebufferConsole, DriverFlags, PixelFormat},
        manager::DEVICE_REGISTRY,
        register_external_interrupt_controller,
    },
    fs::devfs,
};
use libk_mm::{
    address::PhysicalAddress,
    phys::{self, reserved::reserve_region, PhysicalMemoryRegion},
    pointer::PhysicalRef,
    table::EntryLevel,
};
use libk_util::{
    sync::{IrqGuard, SpinFence},
    OneTimeInit,
};
use yboot_proto::{
    v1::{self, AvailableMemoryRegion},
    LoadProtocolV1,
};
use ygg_driver_acpi::{AcpiAllocator, AcpiHandlerImpl, EventAction, FixedEvent};
use ygg_driver_pci::PciBusManager;

use crate::{
    arch::{
        x86::{self, peripherals::hpet::Hpet, InitrdSource},
        x86_64::{
            apic::{ioapic::IoApic, local::LocalApic},
            boot::BootData,
        },
        Platform,
    },
    device::display::linear_fb::LinearFramebuffer,
    panic,
    util::call_init_array,
};

/// Alias for terminal-level page tables
pub type L3 = mem::table::L3;

mod apic;
mod boot;
mod exception;
mod smp;
mod syscall;

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

    fbconsole: OneTimeInit<Arc<FramebufferConsole>>,
}

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

    fbconsole: OneTimeInit::new(),
};

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

impl Platform for X86_64 {
    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);
        }
    }

    unsafe fn power_off(&self) -> Result<!, Error> {
        let _guard = IrqGuard::acquire();

        let ap_count = CPU_COUNT.load(Ordering::Acquire) - 1;
        if ap_count > 0 {
            self.send_ipi(IpiDeliveryTarget::OtherCpus, IpiMessage::Shutdown)?;
            SHUTDOWN_FENCE.wait_all(ap_count);
        }

        ygg_driver_acpi::power_off()
    }
}

impl X86_64 {
    fn handle_ipi(ipi: Option<IpiMessage>) {
        let Some(ipi) = ipi else {
            log::warn!("Spurious IPI received, assuming panic");
            panic::panic_secondary();
        };

        match ipi {
            IpiMessage::Shutdown => {
                log::info!("Core halted");
                SHUTDOWN_FENCE.signal();
                ArchitectureImpl::halt()
            }
            IpiMessage::Panic => panic::panic_secondary(),
        }
    }

    fn send_ipi(&self, target: IpiDeliveryTarget, message: IpiMessage) -> Result<(), Error> {
        let cpu = Cpu::local();
        cpu.local_apic.send_ipi(target, message)
    }

    fn set_boot_data(&self, data: BootData) {
        match data {
            BootData::YBoot(data) => {
                // Reserve the memory map
                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;

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

        self.boot_data.init(data);
    }

    unsafe fn init_physical_memory_from_yboot(data: &LoadProtocolV1) -> Result<(), Error> {
        let mmap = PhysicalRef::<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,
        }))
    }

    unsafe fn init_memory_management(
        &self,
        enabled_features: &CpuFeatures,
        is_bsp: bool,
    ) -> Result<(), Error> {
        let have_1gib_pages = enabled_features.ext_edx.contains(ExtEdxFeatures::PDPE1GB);

        mem::init_fixed_tables(have_1gib_pages, is_bsp);

        if !is_bsp {
            return Ok(());
        }

        // 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 (available_features, enabled_features) = self.init_cpu_features();

        PLATFORM
            .init_memory_management(&enabled_features, cpu_id == 0)
            .expect("Could not initialize memory management");
        let local_apic = self.init_local_cpu(cpu_id, available_features, enabled_features);

        if cpu_id == 0 {
            self.setup_from_boot_data()?;

            // TODO yboot doesn't yet pass any command line options
            let cmdline = self.boot_data.get().cmdline();
            let mut early = x86::init_platform_early(cmdline)?;

            if !config::get().x86.disable_boot_fb {
                if let Err(error) = self.init_framebuffer() {
                    log::error!("Could not initialize boot framebuffer: {error:?}");
                }
            }

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

            if !config::get().x86_64.disable_hpet {
                early.add_preferred_clock_source("hpet", Self::setup_hpet);
            } else {
                log::info!("HPET disabled by config");
            }

            call_init_array();

            x86::init_platform_devices(early);

            extern "C" {
                static __kernel_start: u8;
            }

            log::info!("Load base: {:p}", &raw const __kernel_start);
        }

        Ok(())
    }

    fn setup_hpet() -> Result<Arc<dyn Device>, Error> {
        let acpi = PLATFORM.acpi.try_get().ok_or(Error::DoesNotExist)?;
        let hpet = HpetInfo::new(acpi).map_err(|_| Error::DoesNotExist)?;
        let hpet = Hpet::setup_from_acpi(&hpet)?;
        let timer = hpet.setup_timer(0, 1000, true)?;
        Ok(timer)
    }

    unsafe fn init_cpu_features(&self) -> (CpuFeatures, CpuFeatures) {
        let (have_features, will_features) = cpuid::setup_features(
            CpuFeatures {
                ecx: EcxFeatures::SSE3
                    | EcxFeatures::SSE4_1
                    | EcxFeatures::SSE4_2
                    | EcxFeatures::AVX,
                edx: EdxFeatures::empty(),
                ext_edx: ExtEdxFeatures::PDPE1GB,
            },
            CpuFeatures {
                ecx: EcxFeatures::XSAVE,
                edx: EdxFeatures::FXSR
                    | EdxFeatures::SSE
                    | EdxFeatures::PGE
                    | EdxFeatures::PSE
                    | EdxFeatures::FPU,
                ext_edx: ExtEdxFeatures::empty(),
            },
        );
        let will_features = will_features.expect("Could not initialize CPU features");

        (have_features, will_features)
    }

    unsafe fn init_local_cpu(
        &self,
        cpu_id: usize,
        available_features: CpuFeatures,
        enabled_features: CpuFeatures,
    ) -> Arc<LocalApic> {
        let local_apic = Arc::new(LocalApic::new());
        let tss_address = gdt::init();
        exception::init_exceptions(cpu_id);

        let cpu_data = PerCpuData {
            this: null_mut(),
            tss_address,
            tmp_address: 0,
            local_apic: local_apic.clone(),
            available_features,
            enabled_features,
        };

        Cpu::init_local(Some(cpu_id as _), cpu_data);

        syscall::init_syscall();

        local_apic
    }

    unsafe fn setup_from_boot_data(&self) -> Result<(), Error> {
        match self.boot_data.get() {
            &BootData::YBoot(data) => {
                // Already reserved in set_boot_data()
                x86::set_initrd(data, false);

                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| {
            log::error!("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>,
        local_apic: Arc<LocalApic>,
    ) -> Result<(), Error> {
        let platform_info = acpi.platform_info_in(AcpiAllocator).map_err(|err| {
            log::error!("Could not get ACPI platform info: {:?}", err);
            Error::InvalidArgument
        })?;

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

        for _ in 0..10 {
            log::info!("TICK");
        }
        let ioapic = IoApic::from_acpi(&apic_info)?;
        register_external_interrupt_controller(ioapic);

        if let Err(error) = ygg_driver_acpi::switch_to_acpi(acpi) {
            log::error!("ACPI initialization error: {error:?}");
        } else if let Err(error) =
            ygg_driver_acpi::add_event_handler(&FixedEvent::POWER_BUTTON, |_| {
                log::info!("Power button pressed!");
                EventAction::Nothing
            })
        {
            log::error!("Couldn't set ACPI power button handler: {error:?}");
        }

        if let Ok(mcfg) = acpi.find_table::<Mcfg>() {
            for entry in mcfg.entries() {
                if let Err(error) = PciBusManager::add_segment_from_mcfg(entry, local_apic.clone())
                {
                    log::error!("Could not add PCI bus segment: {error:?}");
                }
            }
        }

        Ok(())
    }

    unsafe fn init_framebuffer(&'static self) -> Result<(), Error> {
        DEVICE_REGISTRY.display.set_callback(|device, node| {
            log::info!("Display registered: {:?}", device.display_name());

            if device
                .driver_flags()
                .contains(DriverFlags::REPLACES_BOOT_FRAMEBUFFER)
            {
                // Switch fbconsole to the new framebuffer
                if let Some(fbconsole) = PLATFORM.fbconsole.try_get()
                    && fbconsole.uses_boot_framebuffer()
                {
                    if let Err(error) = fbconsole.set_display(device.clone()) {
                        log::error!("Couldn't set fb console display: {error:?}");
                    }
                }

                // Switch /dev/fb0 to the new node
                if let Some((name, _)) = node {
                    devfs::redirect("fb0", name).ok();
                }
            }
        });

        let display = match self.boot_data.get() {
            &BootData::YBoot(data) => {
                let info = &data.opt_framebuffer;
                let format = match info.res_format {
                    v1::PIXEL_B8G8R8A8 => PixelFormat::B8G8R8A8,
                    v1::PIXEL_R8G8B8A8 => PixelFormat::R8G8B8A8,
                    // Fallback
                    _ => PixelFormat::R8G8B8A8,
                };

                let framebuffer = 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,
                    format,
                )?;

                Some(Arc::new(framebuffer))
            }
        };

        let Some(display) = display else {
            return Ok(());
        };

        // Register boot display
        DEVICE_REGISTRY.display.register(display.clone(), true)?;

        let fbconsole = self
            .fbconsole
            .init(Arc::new(FramebufferConsole::from_framebuffer(
                display.clone(),
                None,
                true,
            )?));

        debug::add_sink(fbconsole.clone(), config::get().debug.display_level);

        Ok(())
    }

    pub(crate) fn cpu_local_apic() -> Arc<dyn LocalApicInterface> {
        Cpu::local().local_apic.clone()
    }
}

impl InitrdSource for LoadProtocolV1 {
    fn start(&self) -> PhysicalAddress {
        PhysicalAddress::from_u64(self.initrd_address)
    }

    fn end(&self) -> PhysicalAddress {
        PhysicalAddress::from_u64(self.initrd_address + self.initrd_size)
    }
}
