use core::arch::x86_64::{_rdrand64_step, _rdseed64_step};

use {crate::bootmodules::BootModule, core::arch::asm, log::warn};
pub mod memory;

mod cpuid;
mod device_info_collector;
mod gdt;
pub mod graphics;
pub(crate) mod interrupts;
pub mod logging;
pub mod pci;
// pub mod virtio;

pub use {logging::log, memory::PAGE_SIZE};

use {
    crate::allocator,
    limine::{HhdmRequest, KernelFileRequest, MemmapRequest, ModuleRequest},
    x86_64::VirtAddr,
};

extern "C" {
    fn _initial_kernel_heap_start();
    fn _initial_kernel_heap_size();
}

const INITIAL_KERNEL_HEAP_START: *mut u8 = _initial_kernel_heap_start as _;
const INITIAL_KERNEL_HEAP_SIZE: *const () = _initial_kernel_heap_size as _;

#[no_mangle]
#[naked]
#[cfg(not(target_feature = "avx2"))]
unsafe extern "C" fn _kernel_start() -> ! {
    // Initialise SSE, then jump to kernel entrypoint
    core::arch::asm!(
        // Initialise SSE
        "mov rax, cr0",
        "and ax, 0xfffb",
        "or ax, 0x2",
        "mov cr0, rax",
        "mov rax, cr4",
        "or ax, 3 << 9",
        "mov cr4, rax",

        // Jump to the kernel entry point
        "jmp {}",
        sym start,
        options(noreturn),
    )
}

#[no_mangle]
#[naked]
#[cfg(target_feature = "avx2")]
unsafe extern "C" fn _kernel_start() -> ! {
    core::arch::asm!(
        // Enable protected mode and configure control registers
        "mov rax, cr0",
        "and ax, 0xFFFB",  // Clear CR0.EM (bit 2) for coprocessor emulation
        "or ax, 0x2",      // Set CR0.MP (bit 1) for coprocessor monitoring
        "mov cr0, rax",

        "mov rax, cr4",
        "or ax, (1 << 9) | (1 << 10)", // Set CR4.OSFXSR (bit 9) and CR4.OSXMMEXCPT (bit 10)
        "mov cr4, rax",

        // Enable OSXSAVE (required for AVX, AVX2, and XSAVE)
        "mov rax, cr4",
        "or eax, 1 << 18", // Set CR4.OSXSAVE (bit 18)
        "mov cr4, rax",

        // Enable AVX and AVX2 state saving
        "xor rcx, rcx",
        "xgetbv",
        "or eax, 7",       // Enable SSE, AVX, and AVX2 state saving
        "xsetbv",

        // Check for AVX and XSAVE support
        "mov eax, 1",
        "cpuid",
        "and ecx, 0x18000000",
        "cmp ecx, 0x18000000",
        "jne {1}",          // Jump if AVX/OSXSAVE is not supported

        // Check for BMI2 and AVX2 support
        "mov eax, 7",
        "xor ecx, ecx",
        "cpuid",
        "and ebx, (1 << 8) | (1 << 5)", // Check BMI2 (bit 8) and AVX2 (bit 5)
        "cmp ebx, (1 << 8) | (1 << 5)", // Compare to ensure both are supported

        // Check for LZCNT and POPCNT support
        "mov eax, 1",
        "cpuid",
        "and ecx, (1 << 5) | (1 << 23)", // Check LZCNT (bit 5) and POPCNT (bit 23)
        "cmp ecx, (1 << 5) | (1 << 23)", // Compare to ensure both are supported

        // Jump to the kernel entry point
        "jmp {0}",
        sym start,
        sym oops,
        options(noreturn),
    )
}

unsafe extern "C" fn oops() -> ! {
    panic!("your cpu is ancient >:(")
}

unsafe extern "C" fn start() -> ! {
    logging::init();
    crate::logger::init().expect("failed to set logger");
    log::info!("Initialising AKern {}", crate::VERSION);

    static HDHM_REQ: HhdmRequest = HhdmRequest::new(0);
    memory::init_pt(VirtAddr::new(
        HDHM_REQ
            .get_response()
            .get()
            .expect("tried to get physical memory mapping offset from Limine")
            .offset,
    ));
    allocator::init(INITIAL_KERNEL_HEAP_START, INITIAL_KERNEL_HEAP_SIZE as _);

    static MMAP_REQ: MemmapRequest = MemmapRequest::new(0);
    memory::initialize(
        MMAP_REQ
            .get_response()
            .get()
            .expect("tried to get memory map from Limine")
            .memmap(),
    );

    gdt::init();
    interrupts::init();

    static KFILE_REQ: KernelFileRequest = KernelFileRequest::new(0);
    static MOD_REQ: ModuleRequest = ModuleRequest::new(0);

    device_info_collector::collect_device_info();

    // Graphics test
    // {
    //     graphics::init();
    //     let mut dis = DISPLAY.lock();
    //     use embedded_graphics::prelude::RgbColor;

    //     let _ = dis.set_color(Rgb888::YELLOW);
    //     let thick = 6;
    //     let p1 = (400, 30);
    //     let p2 = (200, 150);
    //     let p3 = (600, 150);
    //     let p4 = (200, 350);
    //     let p5 = (600, 350);
    //     let p6 = (400, 470);

    //     {
    //         //HEXAGON

    //         let _ = dis.line(p1.0, p1.1, p2.0, p2.1, thick);
    //         let _ = dis.line(p1.0, p1.1, p3.0, p3.1, thick);
    //         let _ = dis.line(p2.0, p2.1, p4.0, p4.1, thick);
    //         let _ = dis.line(p3.0, p3.1, p5.0, p5.1, thick);
    //         let _ = dis.line(p6.0, p6.1, p4.0, p4.1, thick);
    //         let _ = dis.line(p6.0, p6.1, p5.0, p5.1, thick);
    //     }
    //     {
    //         let _ = dis.line(600, 150, 200, 350, thick);
    //         let _ = dis.line(600, 350, 400, 250, thick);
    //     }

    //     {
    //         let _ = dis.set_color(Rgb888::WHITE);
    //         let hp1 = (350, 150);
    //         let hp2 = (350, 350);
    //         let hp3 = (450, 250);
    //         let hp4 = (350, 250);
    //         let hp5 = (450, 150);
    //         let hp6 = (450, 350);

    //         let _ = dis.line(hp1.0, hp1.1, hp2.0, hp2.1, thick);
    //         let _ = dis.line(hp3.0, hp3.1, hp4.0, hp4.1, thick);
    //         let _ = dis.line(hp5.0, hp5.1, hp6.0, hp6.1, thick);
    //     }

    //     dis.swap_buffers();
    // };

    // TODO: Add in rdseed and rdrand as sources for randomness
    let _rand = xml::XMLElement::new("Random");

    log::trace!("Getting boot modules");
    let bm = MOD_REQ.get_response().get();

    let mut bootmodules = alloc::vec::Vec::new();

    if bm.is_some() {
        let bm = bm.unwrap();
        for x in 0..bm.module_count {
            let file = bm.modules().get(x as usize);
            if file.is_some() {
                let file = file.unwrap();
                let raw_bytes = core::slice::from_raw_parts(
                    file.base.as_ptr().expect("invalid initrd"),
                    file.length as usize,
                );

                let file_path = file.path.to_str().unwrap().to_str();
                if file_path.is_err() {
                    panic!("invalid file path: {:?}", file_path);
                }
                let file_cmd = file.cmdline.to_str().unwrap().to_str();
                if file_cmd.is_err() {
                    panic!("invalid module cmd: {:?}", file_cmd);
                }

                log::trace!("module path: {:?}", file_path);
                log::trace!("module cmd: {:?}", file_cmd);

                bootmodules.push(BootModule::new(
                    file_path.unwrap(),
                    raw_bytes,
                    file_cmd.unwrap(),
                ));
            } else {
                log::error!("You should not be here");
                break;
            }
        }
        log::info!("Boot module count: {:?}", bootmodules.len());
        assert_eq!(bm.module_count, bootmodules.len() as u64);
    }

    crate::kmain::kmain(
        KFILE_REQ
            .get_response()
            .get()
            .and_then(|r| r.kernel_file.get())
            .expect("failed to get kernel file from Limine")
            .cmdline
            .to_str()
            .map(core::ffi::CStr::to_str)
            .transpose()
            .expect("expected valid cmdline string")
            .unwrap_or_default(),
        bootmodules,
    )
}

/// Spin loop
pub fn spin_loop() -> ! {
    loop {
        core::hint::spin_loop();
        x86_64::instructions::hlt()
    }
}

pub fn hardware_random_u64() -> u64 {
    let mut out: u64 = 0;
    match unsafe { _rdrand64_step(&mut out) } {
        1 => out,
        _ => {
            warn!("RDRand not supported.");
            // Try rdseed
            match unsafe { _rdseed64_step(&mut out) } {
                1 => out,
                _ => panic!("Neither RDRand or RDSeed are supported"),
            }
        }
    }
}

pub fn get_edid() {}

#[allow(unused)]
pub fn register_dump() {
    let rax: u64;
    let rbx: u64 = 0;
    let rcx: u64;
    let rdx: u64;
    let si: u64;
    let di: u64;

    let r8: u64; // TODO: r8-r15
    let r9: u64;
    let r10: u64;
    let r11: u64;
    let r12: u64;
    let r13: u64;
    let r14: u64;
    let r15: u64;

    unsafe {
        asm!("",
        out("rax") rax,
        out("rcx") rcx,
        out("rdx") rdx,
        out("si") si,
        out("di") di,
        out("r8") r8,
        out("r9") r9,
        out("r10") r10,
        out("r11") r11,
        out("r12") r12,
        out("r13") r13,
        out("r14") r14,
        out("r15") r15,
        )
    };

    log::error!(
        "Kernel Panic!\r
Register Dump\r
    rax: {:#x}\r
    rbx: {:#x}\r
    rcx: {:#x}\r
    rdx: {:#x}\r
    si : {:#x}\r
    di : {:#x}\r
    r8 : {:#x}\r
    r9 : {:#x}\r
    r11: {:#x}\r
    r12: {:#x}\r
    r13: {:#x}\r
    r14: {:#x}\r
    r15: {:#x}\r
",
        rax,
        rbx,
        rcx,
        rdx,
        si,
        di,
        r8,
        r9,
        r11,
        r12,
        r13,
        r14,
        r15,
    );
}