#![no_std]
#![no_main]
#![cfg_attr(target_arch = "x86_64", feature(abi_x86_interrupt))]
#![feature(naked_functions_rustic_abi)]

extern crate alloc;

use core::fmt::{self, Write};
use core::sync::atomic::{AtomicU64, Ordering};

use limine::{
    BaseRevision,
    request::{
        DateAtBootRequest, ExecutableAddressRequest, FramebufferRequest, HhdmRequest,
        MemoryMapRequest, RequestsEndMarker, RequestsStartMarker,
    },
};

#[cfg(target_arch = "aarch64")]
use crate::arch::aarch64::paging::AArchPageTable;
#[cfg(target_arch = "x86_64")]
use crate::arch::arch::KERNEL_MAPPER;
use crate::driver::io::input::init_input;

#[cfg(target_arch = "aarch64")]
use aarch64_cpu::registers::{DAIF, Writeable};

#[cfg(target_arch = "x86_64")]
use crate::driver::io::ps2::init_keyboard;

use crate::arch::arch::{HHDM_OFFSET, infinite_idle, init, kernel_crash, serial_print};
use crate::driver::{
    elf::loader::run_elf,
    framebuffer::{init_framebuffer, with_framebuffer},
    io::fs::assets::INIT_ELF,
    // serial::{ConsoleWriter, init_serial_console, with_serial_console},
    timer::TIMER,
};

pub mod arch;
pub mod config;
pub mod driver;
pub mod mm;
pub mod task;
pub mod util;

/// Sets the base revision to the latest revision supported by the crate.
/// See specification for further info.
/// Be sure to mark all limine requests with #[used], otherwise they may be removed by the compiler.
#[used]
// The .requests section allows limine to find the requests faster and more safely.
#[unsafe(link_section = ".requests")]
static BASE_REVISION: BaseRevision = BaseRevision::new();

#[used]
#[unsafe(link_section = ".requests")]
static FRAMEBUFFER_REQUEST: FramebufferRequest = FramebufferRequest::new();

#[used]
#[unsafe(link_section = ".requests")]
static HHDM_REQUEST: HhdmRequest = HhdmRequest::new();

#[used]
#[unsafe(link_section = ".requests")]
static MEMORY_MAP_REQUEST: MemoryMapRequest = MemoryMapRequest::new();

#[used]
#[unsafe(link_section = ".requests")]
static DATE_AT_BOOT_REQUEST: DateAtBootRequest = DateAtBootRequest::new();

#[used]
#[unsafe(link_section = ".requests")]
static EXECUTABLE_ADDRESS_REQUEST: ExecutableAddressRequest = ExecutableAddressRequest::new();

/// Define the stand and end markers for Limine requests.
#[used]
#[unsafe(link_section = ".requests_start_marker")]
static _START_MARKER: RequestsStartMarker = RequestsStartMarker::new();
#[used]
#[unsafe(link_section = ".requests_end_marker")]
static _END_MARKER: RequestsEndMarker = RequestsEndMarker::new();

pub static KERNEL_PHYS_BASE: AtomicU64 = AtomicU64::new(0);
pub static KERNEL_VIRT_BASE: AtomicU64 = AtomicU64::new(0);

#[macro_export]
macro_rules! print {
    ($($arg:tt)*) => {
        $crate::_print(core::format_args!($($arg)*))
    };
}

#[macro_export]
macro_rules! println {
    () => {
        $crate::_print(core::format_args!("\n"))
    };
    ($($arg:tt)*) => {
        $crate::_print(core::format_args!("{}\n", core::format_args!($($arg)*)))
    };
}

pub struct SerialWriter {}

impl Write for SerialWriter {
    fn write_str(&mut self, s: &str) -> fmt::Result {
        serial_print(s);
        Ok(())
    }
}

#[doc(hidden)]
pub fn _print(args: core::fmt::Arguments) {
    let mut writer = SerialWriter {};
    let _ = writer.write_fmt(args);
}

#[unsafe(no_mangle)]
unsafe extern "C" fn kmain() -> ! {
    // mask everything early to prevent stack corruption and other undefined behavior
    #[cfg(target_arch = "aarch64")]
    {
        DAIF.write(DAIF::D::Masked + DAIF::A::Masked + DAIF::I::Masked + DAIF::F::Masked);
    }

    if let Some(hhdm_response) = HHDM_REQUEST.get_response() {
        HHDM_OFFSET.store(
            hhdm_response.offset(),
            core::sync::atomic::Ordering::Relaxed,
        );
        #[allow(unused_variables)]
        if let Some(memory_map_response) = MEMORY_MAP_REQUEST.get_response() {
            if let Some(executable_address_response) = EXECUTABLE_ADDRESS_REQUEST.get_response() {
                KERNEL_PHYS_BASE.store(
                    executable_address_response.physical_base(),
                    Ordering::Relaxed,
                );

                KERNEL_VIRT_BASE.store(
                    executable_address_response.virtual_base(),
                    Ordering::Relaxed,
                );

                #[cfg(target_arch = "aarch64")]
                use crate::arch::aarch64::init::preinit_aarch64;
                #[cfg(target_arch = "aarch64")]
                preinit_aarch64(
                    hhdm_response,
                    memory_map_response,
                    executable_address_response,
                );

                #[cfg(target_arch = "x86_64")]
                unsafe {
                    kernel_main_x86_64()
                }

                #[cfg(target_arch = "aarch64")]
                loop {}
            } else {
                kernel_crash()
            }
        } else {
            kernel_crash(); // Could not get required info from Limine's memory map.
        }
    } else {
        kernel_crash(); // Could not get required info from Limine's higher-half direct mapping.
    }
}

#[cfg(target_arch = "aarch64")]
#[unsafe(no_mangle)]
pub unsafe extern "C" fn kernel_main_aarch64(mapper: &mut AArchPageTable) -> ! {
    // All limine requests must also be referenced in a called function, otherwise they may be
    // removed by the linker.
    assert!(BASE_REVISION.is_supported());
    init(mapper);

    if let Some(framebuffer_response) = FRAMEBUFFER_REQUEST.get_response() {
        if let Some(limine_framebuffer) = framebuffer_response.framebuffers().next() {
            init_framebuffer(&limine_framebuffer);
            with_framebuffer(|fb| fb.setup_aarch64());
        } else {
            panic!("no framebuffers found");
        }
    } else {
        panic!("no framebuffers found");
    }

    init_input();

    if let Some(date_at_boot_response) = DATE_AT_BOOT_REQUEST.get_response() {
        TIMER.set_date_at_boot(date_at_boot_response.timestamp().as_secs());
    } else {
        println!("Could not get date at boot. Will default to 0.")
    }
    DAIF.write(DAIF::D::Masked + DAIF::A::Masked + DAIF::I::Unmasked + DAIF::F::Masked);

    run_elf(INIT_ELF, false, true);

    loop {}
}

#[cfg(target_arch = "x86_64")]
pub unsafe fn kernel_main_x86_64() -> ! {
    // All limine requests must also be referenced in a called function, otherwise they may be
    // removed by the linker.
    assert!(BASE_REVISION.is_supported());

    let (hhdm_response, memory_map_response) = {
        if let Some(a) = HHDM_REQUEST.get_response() {
            if let Some(b) = MEMORY_MAP_REQUEST.get_response() {
                (a, b)
            } else {
                kernel_crash(); // Could not get required info from Limine's memory map.
            }
        } else {
            kernel_crash(); // Could not get required info from Limine's higher-half direct mapping.
        }
    };

    let mut mapper = init(hhdm_response, memory_map_response);

    if let Some(framebuffer_response) = FRAMEBUFFER_REQUEST.get_response() {
        if let Some(limine_framebuffer) = framebuffer_response.framebuffers().next() {
            init_framebuffer(&limine_framebuffer);
            with_framebuffer(|fb| fb.setup_x86_64(&mut mapper));
        } else {
            serial_print("no framebuffers found");
        }
    }

    init_keyboard();
    init_input();

    #[allow(static_mut_refs)]
    unsafe {
        *KERNEL_MAPPER.get_mut() = Some(mapper)
    };

    if let Some(date_at_boot_response) = DATE_AT_BOOT_REQUEST.get_response() {
        TIMER.set_date_at_boot(date_at_boot_response.timestamp().as_secs());
    } else {
        println!("Could not get date at boot. Will default to 0.")
    }

    run_elf(INIT_ELF, false, true);

    loop {}
}

struct BufWriter<'a> {
    buf: &'a mut [u8],
    pos: usize,
}

impl<'a> BufWriter<'a> {
    fn new(buf: &'a mut [u8]) -> Self {
        Self { buf, pos: 0 }
    }
}

impl core::fmt::Write for BufWriter<'_> {
    fn write_str(&mut self, s: &str) -> core::fmt::Result {
        let bytes = s.as_bytes();
        let remaining = &mut self.buf[self.pos..];
        let len = bytes.len().min(remaining.len());
        remaining[..len].copy_from_slice(&bytes[..len]);
        self.pos += len;
        Ok(())
    }
}

#[panic_handler]
fn rust_panic(_info: &core::panic::PanicInfo) -> ! {
    serial_print("\nKERNEL PANIC:\n");
    serial_print(
        _info
            .message()
            .as_str()
            .unwrap_or("Could not get message str"),
    );
    let mut buf = [0u8; 512];
    let msg = {
        let mut w = BufWriter::new(&mut buf);
        let _ = core::fmt::write(&mut w, core::format_args!("{}", _info));
        let len = w.pos;
        core::str::from_utf8(&buf[..len]).unwrap_or("(utf8 error)")
    };
    serial_print(msg);
    serial_print("\n");
    // with_framebuffer(|mut fb| {
    //     fb.clear(rgb(180, 0, 0));

    //     with_serial_console(|console| {
    //         console.clear();

    //         let mut writer = ConsoleWriter {
    //             fb: &mut fb,
    //             console,
    //             should_center: true,
    //         };

    //         let _ = writer.write_str("KERNEL PANIC\n\n");
    //         let _ = writer.write_fmt(core::format_args!("{}", _info));
    //     });
    // });

    infinite_idle();
}
