// TODO fix all TODOs
use core::{mem::size_of, sync::atomic::Ordering};

use abi::error::Error;
use acpi_lib::{mcfg::Mcfg, AcpiTables, InterruptModel};
use alloc::boxed::Box;
use device_api::{
    input::KeyboardProducer, interrupt::ExternalInterruptController,
    timer::MonotonicTimestampProviderDevice, Device,
};
use git_version::git_version;
use kernel_util::util::OneTimeInit;
use memtables::FixedTables;
use yboot_proto::{v1::AvailableMemoryRegion, LoadProtocolV1};

mod acpi;
mod apic;
mod boot;
pub mod context;
pub mod cpu;
mod cpuid;
mod exception;
mod gdt;
mod intrinsics;
pub mod mem;
mod peripherals;
mod registers;
mod smp;
mod syscall;

use crate::{
    arch::x86_64::{
        intrinsics::{IoPort, IoPortAccess},
        mem::{map_heap_block, table::L2, HEAP_MAPPING_OFFSET},
    },
    debug::{self, LogLevel},
    device::{
        self,
        bus::pci::PciBusManager,
        display::{console, fb_console::FramebufferConsole, linear_fb::LinearFramebuffer},
        tty::CombinedTerminal,
    },
    fs::{
        devfs::{self, CharDeviceType},
        Initrd, INITRD_DATA,
    },
    mem::{
        address::{FromRaw, IntoRaw},
        device::RawDeviceMemoryMapping,
        heap,
        phys::{self, reserved::reserve_region, PhysicalMemoryRegion},
        table::EntryLevel,
        PhysicalAddress,
    },
    sync::SpinFence,
};

use self::{
    acpi::{AcpiAllocator, AcpiHandlerImpl},
    apic::{ioapic::IoApic, local::LocalApic},
    boot::BootData,
    cpu::Cpu,
    cpuid::{ProcessorFeatures, PROCESSOR_FEATURES},
    mem::{
        init_fixed_tables,
        table::{PageAttributes, PageEntry, L1, L3},
        EarlyMapping, KERNEL_TABLES, MEMORY_LIMIT, RAM_MAPPING_L1, RAM_MAPPING_OFFSET,
    },
    peripherals::{i8253::I8253, ps2::PS2Controller, serial::ComPort},
    smp::CPU_COUNT,
};

use super::{Architecture, CpuMessage};

#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash)]
pub enum IrqNumber {
    Isa(u8),
    Gsi(u8),
}

pub struct X86_64 {
    boot_data: OneTimeInit<BootData>,
    acpi: OneTimeInit<AcpiTables<AcpiHandlerImpl>>,

    // Display
    framebuffer: OneTimeInit<LinearFramebuffer>,
    fbconsole: OneTimeInit<FramebufferConsole>,
    tty: OneTimeInit<CombinedTerminal>,

    ioapic: OneTimeInit<IoApic>,
    timer: OneTimeInit<I8253>,
}

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

pub static ARCHITECTURE: X86_64 = X86_64 {
    boot_data: OneTimeInit::new(),
    acpi: OneTimeInit::new(),

    framebuffer: OneTimeInit::new(),
    fbconsole: OneTimeInit::new(),
    tty: OneTimeInit::new(),

    ioapic: OneTimeInit::new(),
    timer: OneTimeInit::new(),
};

impl Architecture for X86_64 {
    const KERNEL_VIRT_OFFSET: usize = 0xFFFFFF8000000000;
    type IrqNumber = IrqNumber;

    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);
        }
    }

    fn cpu_count() -> usize {
        CPU_COUNT.load(Ordering::Acquire)
    }

    unsafe fn set_interrupt_mask(mask: bool) {
        if mask {
            core::arch::asm!("cli");
        } else {
            core::arch::asm!("sti");
        }
    }

    fn interrupt_mask() -> bool {
        let mut flags: u64;
        unsafe {
            core::arch::asm!("pushfq; pop {0}", out(reg) flags, options(att_syntax));
        }
        // If IF is zero, interrupts are disabled (masked)
        flags & (1 << 9) == 0
    }

    fn wait_for_interrupt() {
        unsafe {
            core::arch::asm!("hlt");
        }
    }

    #[inline]
    unsafe fn map_device_memory(
        &self,
        base: PhysicalAddress,
        size: usize,
    ) -> Result<RawDeviceMemoryMapping, Error> {
        mem::map_device_memory(base, size)
    }

    #[inline]
    unsafe fn unmap_device_memory(&self, map: &RawDeviceMemoryMapping) {
        mem::unmap_device_memory(map)
    }

    fn map_physical_memory<I: Iterator<Item = PhysicalMemoryRegion> + Clone>(
        &self,
        it: I,
        _memory_start: PhysicalAddress,
        memory_end: PhysicalAddress,
    ) -> Result<(), Error> {
        let end_l1i = (IntoRaw::<usize>::into_raw(memory_end) + (1 << 30) - 1) >> 30;

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

        MEMORY_LIMIT.init(memory_end.into_raw());

        // 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_raw(l1i << 30),
                        PageAttributes::WRITABLE,
                    );
                }
            }

            Ok(())
        } else {
            todo!();
        }
    }

    #[inline]
    fn virtualize(address: PhysicalAddress) -> Result<usize, Error> {
        let raw: usize = address.into_raw();
        if raw < *mem::MEMORY_LIMIT.get() {
            Ok(raw + RAM_MAPPING_OFFSET)
        } else {
            errorln!("Invalid physical address: {:#x}", address);
            Err(Error::InvalidMemoryOperation)
        }
    }

    #[inline]
    fn physicalize(address: usize) -> Result<PhysicalAddress, Error> {
        if address < RAM_MAPPING_OFFSET || address - RAM_MAPPING_OFFSET >= *mem::MEMORY_LIMIT.get()
        {
            errorln!("Not a virtualized physical address: {:#x}", address);
            return Err(Error::InvalidMemoryOperation);
        }

        Ok(PhysicalAddress::from_raw(address - RAM_MAPPING_OFFSET))
    }

    fn external_interrupt_controller(
        &'static self,
    ) -> &'static dyn ExternalInterruptController<IrqNumber = Self::IrqNumber> {
        self.ioapic.get()
    }

    fn monotonic_timer(&'static self) -> &'static dyn MonotonicTimestampProviderDevice {
        self.timer.get()
    }
}

impl X86_64 {
    unsafe fn handle_ipi(&self, msg: CpuMessage) {
        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_raw(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_raw(aligned_start),
                                size: (aligned_end - aligned_start) as usize,
                            },
                        );
                    }
                }
            }
        }

        self.boot_data.init(data);
    }

    unsafe fn init_physical_memory_from_yboot(data: &LoadProtocolV1) {
        let mmap = EarlyMapping::<AvailableMemoryRegion>::map_slice(
            PhysicalAddress::from_raw(data.memory_map.address),
            data.memory_map.len as usize,
        )
        .unwrap();

        phys::init_from_iter(mmap.as_ref().iter().map(|reg| PhysicalMemoryRegion {
            base: PhysicalAddress::from_raw(reg.start_address),
            size: reg.page_count as usize * 0x1000,
        }));
    }

    unsafe fn init_memory_management(&self) {
        const HEAP_PAGES: usize = 16;

        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),
        }

        // Setup heap
        for i in 0..HEAP_PAGES {
            // Allocate in 2MiB chunks
            let l2_page = phys::alloc_2m_page().unwrap();

            map_heap_block(i, l2_page);
        }

        heap::init_heap(HEAP_MAPPING_OFFSET, HEAP_PAGES * L2::SIZE);
    }

    unsafe fn init_platform(&'static self, cpu_id: usize) {
        Cpu::init_local(LocalApic::new(), cpu_id as _);

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

            self.init_acpi_from_boot_data();

            Self::disable_8259();

            self.timer.init(I8253::new());

            let com1_3 = Box::leak(Box::new(ComPort::new(0x3F8, 0x3E8, IrqNumber::Isa(4))));
            debug::add_sink(com1_3.port_a(), LogLevel::Debug);

            self.init_framebuffer();

            debug::init();

            let ps2 = Box::leak(Box::new(PS2Controller::new(
                IrqNumber::Isa(1),
                IrqNumber::Isa(12),
                0x64,
                0x60,
            )));
            ps2.init().unwrap();

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

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

            self.timer.get().init_irq().unwrap();
            ps2.connect(self.tty.get());
            ps2.init_irq().unwrap();

            device::register_device(self.ioapic.get());
            device::register_device(ps2);

            PciBusManager::setup_bus_devices().unwrap();
        }
    }

    unsafe fn init_acpi_from_boot_data(&self) {
        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) {
        let acpi_tables = AcpiTables::from_rsdp(AcpiHandlerImpl, rsdp).unwrap();
        self.acpi.init(acpi_tables);
    }

    unsafe fn init_platform_from_acpi(&self, acpi: &'static AcpiTables<AcpiHandlerImpl>) {
        let platform_info = acpi.platform_info_in(AcpiAllocator).unwrap();

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

        self.ioapic.init(IoApic::from_acpi(&apic_info).unwrap());

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

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

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

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

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

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

        devfs::add_char_device(self.tty.get(), CharDeviceType::TtyRegular).unwrap();
        console::add_console_autoflush(self.fbconsole.get());
    }

    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_raw() 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);
    }

    unsafe fn disable_8259() {
        infoln!("Disabling i8259 PIC");
        // TODO should I make a module for 8259 if I don't even use it?
        let pic_master_cmd = IoPort::<u8>::new(0x20);
        let pic_master_data = IoPort::<u8>::new(0x21);
        let pic_slave_cmd = IoPort::<u8>::new(0xA0);
        let pic_slave_data = IoPort::<u8>::new(0xA1);

        // Remap PIC IRQ vectors to 32..
        pic_master_cmd.write(0x11);
        pic_slave_cmd.write(0x11);

        pic_master_data.write(32);
        pic_slave_data.write(32 + 8);

        pic_slave_data.write(0xFF);
        pic_master_data.write(0xFF);

        pic_master_cmd.write(0x20);
        pic_slave_cmd.write(0x20);
    }
}