#![no_std]
#![no_main]
#![feature(abi_x86_interrupt)]
#![allow(clippy::single_match)]

mod acpi;
mod arch;
mod cap;
mod console;
mod dashboard;
mod error;
mod flight;
mod iommu;
mod ipc;
mod irq;
mod log;
mod mem;
mod panic;
mod pci;
mod proc;
#[cfg(lancer_test)]
mod qemu;
mod ring;
mod sched;
mod static_vec;
mod sync;
mod syscall;
#[cfg(lancer_test)]
mod test_harness;
#[cfg(lancer_test)]
mod tests;
mod types;
mod wcet;

use wcet::region::{WcetGuard, WcetTracker};
#[cfg(not(lancer_test))]
use wcet::tsc;

#[unsafe(no_mangle)]
#[unsafe(naked)]
extern "C" fn kmain() -> ! {
    core::arch::naked_asm!(
            "mov rax, cr0",
            "and ax, 0xFFFB",
            "or ax, 0x0002",
            "mov cr0, rax",
            "mov rax, cr4",
            "or ax, 0x0600",
            "mov cr4, rax",
            "jmp {entry}",
            entry = sym kmain_inner,
    )
}

extern "C" fn kmain_inner() -> ! {
    arch::serial::init();
    arch::syscall::init_bsp_percpu();

    dashboard::dash().add_section("kern", 3);
    dashboard::dash().add_section("mem", 3);
    dashboard::dash().add_section("gdt", 1);
    dashboard::dash().add_section("idt", 1);
    dashboard::dash().add_section("phys", 2);
    dashboard::dash().add_section("virt", 1);
    dashboard::dash().add_section("apic", 2);
    dashboard::dash().add_section("acpi", 1);
    dashboard::dash().add_section("ioapic", 2);
    dashboard::dash().add_section("pci", 1);
    dashboard::dash().add_section("iommu", 1);
    dashboard::dash().add_section("cpu", 4);
    dashboard::dash().add_section("sys", 1);
    dashboard::dash().add_section("cap", 1);
    dashboard::dash().add_section("sched", 1);
    #[cfg(not(lancer_test))]
    dashboard::dash().add_section("proc", 2);
    #[cfg(not(lancer_test))]
    dashboard::dash().add_section("fb", 1);
    #[cfg(not(lancer_test))]
    dashboard::dash().add_section("boot", 6);

    #[cfg(not(lancer_test))]
    let boot_start = tsc::read_tsc();
    let mut serial_wcet: WcetTracker<16> = WcetTracker::new("serial_init");
    let mut memmap_wcet: WcetTracker<16> = WcetTracker::new("memmap_parse");
    let mut gdt_wcet: WcetTracker<16> = WcetTracker::new("gdt_init");
    let mut idt_wcet: WcetTracker<16> = WcetTracker::new("idt_init");
    let mut phys_wcet: WcetTracker<16> = WcetTracker::new("phys_alloc_init");
    let mut virt_wcet: WcetTracker<16> = WcetTracker::new("virt_mem_init");
    let mut apic_wcet: WcetTracker<16> = WcetTracker::new("apic_init");

    serial_wcet.record(0);

    phase_begin!(kern);
    arch::boot::validate_base_revision();
    show!(kern, "limine base revision validated");

    let hhdm_offset = arch::boot::hhdm_offset();
    show!(kern, "hhdm offset {:#x}", hhdm_offset);

    let (phys_base, virt_base) = arch::boot::executable_address();
    show!(kern, "kernel phys {:#x} virt {:#x}", phys_base, virt_base);
    phase_end!(kern);

    phase_begin!(mem);
    {
        let _guard = WcetGuard::new(&mut memmap_wcet);
        let memory_map = arch::boot::memory_map();
        show!(mem, "memory map {} entries", memory_map.len());
        mem::init(memory_map, hhdm_offset);
    }
    proc::address_space::capture_boot_pml4();
    console::init(16);
    {
        let stats = mem::phys::frame_stats();
        show!(mem, "{} total, {} free", stats.0, stats.1);
    }
    phase_end!(mem);

    phase_begin!(gdt);
    {
        let _guard = WcetGuard::new(&mut gdt_wcet);
        arch::gdt::init();
    }
    show!(gdt, "loaded");
    phase_end!(gdt);

    phase_begin!(idt);
    {
        let _guard = WcetGuard::new(&mut idt_wcet);
        arch::idt::init();
    }
    show!(idt, "32 exception handlers + apic vectors");
    phase_end!(idt);

    phase_begin!(phys);
    {
        let _guard = WcetGuard::new(&mut phys_wcet);
    }
    {
        let stats = mem::phys::frame_stats();
        show!(phys, "{} frames available", stats.1);
    }
    phase_end!(phys);

    phase_begin!(virt);
    let mut mapper = unsafe { arch::paging::init(hhdm_offset) };
    let mut allocator = mem::phys::BitmapFrameAllocator;
    {
        let _guard = WcetGuard::new(&mut virt_wcet);
        mem::virt::test_mapping(&mut mapper, &mut allocator);
    }
    show!(virt, "page tables initialized");
    phase_end!(virt);

    {
        let mut ptable = proc::PROCESSES.lock();
        ptable.init(&mut allocator);
    }

    phase_begin!(apic);
    {
        let _guard = WcetGuard::new(&mut apic_wcet);
        let apic_base = unsafe { x2apic::lapic::xapic_base() };
        mem::virt::map_mmio(&mut mapper, &mut allocator, apic_base, hhdm_offset).expect("map_mmio");
        arch::apic::init(hhdm_offset);
    }
    show!(apic, "lapic enabled, timer configured");
    phase_end!(apic);

    phase_begin!(acpi);
    let acpi_info = match arch::boot::rsdp_address() {
        Some(rsdp) => {
            show!(acpi, "rsdp at {:#x}", rsdp);
            acpi::parse(rsdp, &mut mapper, &mut allocator, hhdm_offset)
        }
        None => {
            show!(acpi, "no rsdp from bootloader");
            None
        }
    };
    phase_end!(acpi);

    phase_begin!(ioapic);
    let (com1_gsi, kbd_gsi) = match &acpi_info {
        Some(info) => {
            mem::virt::map_mmio(&mut mapper, &mut allocator, info.ioapic_addr, hhdm_offset)
                .expect("map_mmio");
            arch::ioapic::init(info.ioapic_addr);
            let com1 = info.gsi_for_irq(4);
            arch::ioapic::route_irq(com1, arch::idt::COM1_VECTOR, 0);
            let kbd = info.gsi_for_irq(1);
            arch::ioapic::route_irq(kbd, arch::idt::KBD_VECTOR, 0);
            show!(
                ioapic,
                "com1 gsi {} -> vec {}",
                com1,
                arch::idt::COM1_VECTOR
            );
            show!(ioapic, "kbd gsi {} -> vec {}", kbd, arch::idt::KBD_VECTOR);
            (Some(com1), Some(kbd))
        }
        None => {
            show!(ioapic, "skipped, no acpi");
            (None, None)
        }
    };
    phase_end!(ioapic);

    phase_begin!(pci);
    match &acpi_info {
        Some(info) if !info.mcfg_entries.is_empty() => {
            acpi::mcfg::cache_mcfg_entries(&info.mcfg_entries);
            pci::enumerate(&info.mcfg_entries, &mut mapper, &mut allocator, hhdm_offset);
            let count = pci::DEVICE_TABLE.lock().len();
            show!(pci, "{} devices enumerated", count);
        }
        _ => {
            show!(pci, "skipped, no mcfg");
        }
    }
    phase_end!(pci);

    phase_begin!(iommu);
    match &acpi_info {
        Some(info) if info.dmar_info.is_some() => {
            let dmar = info.dmar_info.as_ref().unwrap();
            match iommu::init(dmar, &mut mapper, &mut allocator, hhdm_offset) {
                Ok(()) => match iommu::enable(&mut allocator, hhdm_offset) {
                    Ok(()) => show!(iommu, "enabled"),
                    Err(e) => {
                        show!(iommu, error, "enable failed {:?}", e);
                        phase_fail!(iommu);
                    }
                },
                Err(e) => {
                    show!(iommu, error, "init failed {:?}", e);
                    phase_fail!(iommu);
                }
            }
        }
        _ => {
            show!(iommu, "no dmar table");
        }
    }
    match dashboard::domain_to_section("iommu") {
        Some(idx)
            if !matches!(
                dashboard::dash().section_status(idx),
                lancer_log::phase::PhaseStatus::Failed
            ) =>
        {
            phase_end!(iommu)
        }
        _ => {}
    }

    let virtio_net_gsi = match &acpi_info {
        Some(info) => {
            let dev_table = pci::DEVICE_TABLE.lock();
            dev_table
                .iter()
                .find(|d| d.vendor_id == 0x1AF4 && d.class_code == 0x02 && d.subclass == 0x00)
                .map(|d| {
                    let irq = d.interrupt_line;
                    let gsi = info.gsi_for_irq(irq);
                    match d.msix_cap.is_some() {
                        true => {}
                        false => {
                            arch::ioapic::route_irq(gsi, arch::idt::VIRTIO_NET_VECTOR, 0);
                        }
                    }
                    gsi
                })
        }
        None => None,
    };

    syscall::platform::store_platform_info(com1_gsi, kbd_gsi, virtio_net_gsi);

    phase_begin!(cpu);
    {
        use x86_64::registers::control::{Cr0, Cr0Flags, Cr4, Cr4Flags};
        use x86_64::registers::model_specific::{Efer, EferFlags};

        unsafe {
            Efer::update(|flags| *flags |= EferFlags::NO_EXECUTE_ENABLE);
        }
        show!(cpu, "nxe enabled");

        unsafe {
            Cr0::update(|f| {
                f.remove(Cr0Flags::EMULATE_COPROCESSOR);
                f.remove(Cr0Flags::TASK_SWITCHED);
                f.insert(Cr0Flags::MONITOR_COPROCESSOR);
                f.insert(Cr0Flags::NUMERIC_ERROR);
            });
            Cr4::update(|f| {
                f.insert(Cr4Flags::OSFXSR);
                f.insert(Cr4Flags::OSXMMEXCPT_ENABLE);
            });
        }
        show!(cpu, "fpu/sse enabled");

        arch::xsave::init();
        arch::xsave::propagate_to_percpu();

        let leaf7 = core::arch::x86_64::__cpuid_count(7, 0);
        let has_smep = leaf7.ebx & (1 << 7) != 0;
        let has_smap = leaf7.ebx & (1 << 20) != 0;

        if has_smep {
            unsafe { Cr4::update(|f| *f |= Cr4Flags::SUPERVISOR_MODE_EXECUTION_PROTECTION) };
            show!(cpu, "smep enabled");
        }

        if has_smap {
            unsafe { Cr4::update(|f| *f |= Cr4Flags::SUPERVISOR_MODE_ACCESS_PREVENTION) };
            arch::syscall::set_smap_active(true);
            show!(cpu, "smap enabled");
        }
    }
    phase_end!(cpu);

    phase_begin!(sys);
    arch::syscall::init();
    show!(sys, "syscall/sysret configured");
    phase_end!(sys);

    phase_begin!(cap);
    cap::init(&mut allocator);
    show!(cap, "initialized");
    phase_end!(cap);

    phase_begin!(sched);
    sched::init();
    wcet::init();
    sched::seed_timer_wheel();
    x86_64::instructions::interrupts::enable();
    show!(sched, "scheduler + interrupts enabled");
    phase_end!(sched);

    #[cfg(lancer_test)]
    {
        dashboard::dash().transition_to_streaming();
        let _ = (com1_gsi, kbd_gsi, virtio_net_gsi);
        {
            let mut ptable = proc::PROCESSES.lock();
            let created = ptable
                .allocate(&mut allocator)
                .expect("Failed to create process 0");
            assert!(created.pid().raw() == 0, "First process must be pid 0");
            ptable.start(created).expect("Created -> Ready");
        }
        {
            let memory_map = arch::boot::memory_map();
            let reserved = cap::boot_untyped::build_reserved_regions();
            let boot_untypeds = {
                let ptable = proc::PROCESSES.lock();
                let mut pool = cap::pool::POOL.lock_after(&ptable);
                cap::boot_untyped::create_boot_untypeds(memory_map, &reserved, &mut pool)
            };

            cap::boot_untyped::store_metadata(boot_untypeds);
        }
        crate::test_harness::run_all();
    }

    #[cfg(not(lancer_test))]
    {
        phase_begin!(proc);

        fn extract_filename(path: &[u8]) -> &[u8] {
            let last_slash = (0..path.len()).rev().find(|&i| path[i] == b'/');
            match last_slash {
                Some(pos) => &path[pos + 1..],
                None => path,
            }
        }

        let modules = arch::boot::modules().unwrap_or(&[]);
        show!(proc, "{} modules available", modules.len());

        let init_file = modules
            .iter()
            .find(|file| {
                let filename = extract_filename(file.path().to_bytes());
                filename.len() == 4 && filename.iter().zip(b"init").all(|(a, b)| a == b)
            })
            .expect("init module not found");

        let init_pid = {
            let mut ptable = proc::PROCESSES.lock();
            let created = ptable
                .allocate(&mut allocator)
                .expect("failed to create init process");
            let pid = created.pid();
            ptable.start(created).expect("Created -> Ready");
            drop(ptable);

            let data =
                unsafe { core::slice::from_raw_parts(init_file.addr(), init_file.size() as usize) };
            proc::loader::spawn_module(pid, data, &mut allocator).expect("init module load failed");
            {
                let mut ptable = proc::PROCESSES.lock();
                let exec = ptable.exec_mut(pid).expect("init pid missing");
                exec.set_name(b"init");
                exec.saved_context.rdi = u64::MAX;
                exec.seal_context();
            }
            pid
        };

        show!(proc, "init loaded as pid {}", init_pid.raw());
        phase_end!(proc);

        phase_begin!(fb);
        match arch::boot::framebuffer() {
            Some(fb) => show!(
                fb,
                "{}x{} {}bpp pitch {}",
                fb.width(),
                fb.height(),
                fb.bpp(),
                fb.pitch()
            ),
            None => show!(fb, "no framebuffer available"),
        }
        phase_end!(fb);

        phase_begin!(boot);
        proc::bootstrap::bootstrap_self_cap(init_pid, 10).expect("init self-cap failed");

        {
            const KERNEL_FRAME_RESERVE: usize = 5120;
            let reserve_base = allocator
                .allocate_contiguous(KERNEL_FRAME_RESERVE)
                .expect("kernel frame reserve allocation failed");

            let memory_map = arch::boot::memory_map();
            let reserved = cap::boot_untyped::build_reserved_regions();
            let boot_untypeds = {
                let ptable = proc::PROCESSES.lock();
                let mut pool = cap::pool::POOL.lock_after(&ptable);
                cap::boot_untyped::create_boot_untypeds(memory_map, &reserved, &mut pool)
            };

            cap::boot_untyped::seal_bitmap(&boot_untypeds);

            (0..KERNEL_FRAME_RESERVE).for_each(|i| {
                let phys = x86_64::PhysAddr::new(reserve_base.as_u64() + (i as u64) * 4096);
                allocator.deallocate_frame(
                    x86_64::structures::paging::PhysFrame::containing_address(phys),
                );
            });

            let total_bytes: u64 = boot_untypeds
                .as_slice()
                .iter()
                .filter(|u| !u.is_device)
                .map(|u| 1u64 << u.size_bits)
                .fold(0u64, u64::saturating_add);
            show!(
                boot,
                "{} boot untypeds, {}MB usable exported, {}KB kernel reserve",
                boot_untypeds.len(),
                total_bytes / (1024 * 1024),
                (KERNEL_FRAME_RESERVE * 4096) / 1024,
            );

            cap::boot_untyped::install_boot_untypeds(init_pid, &boot_untypeds)
                .expect("install boot untypeds");
            cap::boot_untyped::store_metadata(boot_untypeds);
        }

        proc::address_space::sync_kernel_entries();
        phase_end!(boot);

        let boot_end = tsc::read_tsc();
        let boot_cycles = boot_end - boot_start;

        dashboard::dash().transition_to_streaming();
        show!(
            boot,
            "complete {} cycles {}ns",
            boot_cycles,
            tsc::cycles_to_ns(boot_cycles)
        );

        sched::rescue();
    }
}