#![no_std]
#![no_main]

extern crate alloc;

use core::arch::asm;
use core::fmt::{self, Write};
use alloc::vec::Vec;

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

use crate::mem::paging::{PageTable, PageTableFlags};
use crate::mem::address::{PhysAddr, VirtAddr};
use crate::cap::{Relation, CapRights, Capability, CapObject};

pub mod cap;
mod mem;
#[macro_use]
pub mod debug;

pub mod allocator {
    pub use crate::mem::allocator::ALLOCATOR;
}

use embedded_graphics::{
    mono_font::{ascii::FONT_6X10, MonoTextStyle},
    pixelcolor::Rgb888,
    prelude::*,
    text::Text,
    Drawable,
};

/// Low-level pixel rendering engine
struct FramebufferDisplay<'a> {
    framebuffer: &'a limine::framebuffer::Framebuffer<'a>,
}

impl DrawTarget for FramebufferDisplay<'_> {
    type Color = Rgb888;
    type Error = core::convert::Infallible;

    fn draw_iter<I>(&mut self, pixels: I) -> Result<(), Self::Error>
    where I: IntoIterator<Item = Pixel<Self::Color>> {
        for Pixel(coord, color) in pixels {
            if coord.x >= 0 && coord.x < self.framebuffer.width() as i32 &&
               coord.y >= 0 && coord.y < self.framebuffer.height() as i32 {

                let offset = (coord.y as usize * self.framebuffer.pitch() as usize) + (coord.x as usize * 4);

                unsafe {
                    self.framebuffer.addr().add(offset).cast::<u32>().write(color.into_storage());
                }
            }
        }
        Ok(())
    }
}

impl OriginDimensions for FramebufferDisplay<'_> {
    fn size(&self) -> Size {
        Size::new(self.framebuffer.width() as u32, self.framebuffer.height() as u32)
    }
}

/// High-level console for terminal-like output
struct Console<'a> {
    display: FramebufferDisplay<'a>,
    x: i32,
    y: i32,
    current_color: Rgb888,
}

impl<'a> Console<'a> {
    pub fn new(display: FramebufferDisplay<'a>) -> Self {
        Self { display, x: 10, y: 20, current_color: Rgb888::WHITE }
    }

    pub fn set_color(&mut self, color: Rgb888) {
        self.current_color = color;
    }

    pub fn clear(&mut self) {
        let fb = self.display.framebuffer;
        unsafe {
            core::ptr::write_bytes(fb.addr(), 0, fb.pitch() as usize * fb.height() as usize);
        }
        self.x = 10;
        self.y = 20;
    }

    /// Shifts pixels upward to make room for new lines
    fn scroll(&mut self) {
        let fb = self.display.framebuffer;
        let pitch = fb.pitch() as usize;
        let height = fb.height() as usize;
        let font_height = 10;
        let shift = font_height * pitch;

        unsafe {
            let addr = fb.addr();
            core::ptr::copy(addr.add(shift), addr, pitch * (height - font_height));
            core::ptr::write_bytes(addr.add(pitch * (height - font_height)), 0, shift);
        }
        self.y -= font_height as i32;
    }
}

impl fmt::Write for Console<'_> {
    fn write_str(&mut self, s: &str) -> fmt::Result {
        let style = MonoTextStyle::new(&FONT_6X10, self.current_color);

        for c in s.chars() {
            if c == '\n' {
                self.x = 10;
                self.y += 10;
            } else {
                if self.x + 6 > self.display.framebuffer.width() as i32 {
                    self.x = 10;
                    self.y += 10;
                }

                let mut buf = [0u8; 4];
                let s_char = c.encode_utf8(&mut buf);
                let _ = Text::new(s_char, Point::new(self.x, self.y), style).draw(&mut self.display);

                self.x += 6;
            }

            if self.y > self.display.framebuffer.height() as i32 - 10 {
                self.scroll();
            }
        }
        Ok(())
    }
}

// --- Limine Requests ---

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

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

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

#[used] #[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_start_marker")]
static _START_MARKER: RequestsStartMarker = RequestsStartMarker::new();

#[used] #[unsafe(link_section = ".requests_end_marker")]
static _END_MARKER: RequestsEndMarker = RequestsEndMarker::new();

#[unsafe(no_mangle)]
unsafe extern "C" fn kmain() -> ! {
    assert!(BASE_REVISION.is_supported());

    let fb_res = FRAMEBUFFER_REQUEST.get_response().expect("Limine: No Framebuffer");
    let mmap_res = MEMORY_MAP_REQUEST.get_response().expect("Limine: No Memory Map");
    let hhdm_res = HHDM_REQUEST.get_response().expect("Limine: No HHDM");
    let kaddr_res = KERNEL_ADDR_REQUEST.get_response().expect("Limine: No Kernel Address");
    let hhdm_offset = hhdm_res.offset();

    let fb = fb_res.framebuffers().next().expect("Limine: No active framebuffer found");
    let mut console = Console::new(FramebufferDisplay { framebuffer: &fb });
    console.clear();

    info!(console, "BOOT", "LISA Kernel Starting...");

    unsafe { mem::pmm::BitmapPMM::init(&mmap_res, hhdm_offset); }
    info!(console, "MEM", "Physical Memory Manager initialized.");

    let p4_phys = mem::pmm::alloc_frame().expect("OOM: Failed to allocate P4 table");
    let p4 = unsafe { &mut *p4_phys.to_virt(hhdm_offset).as_mut_ptr::<PageTable>() };
    unsafe { core::ptr::write_bytes(p4 as *mut _ as *mut u8, 0, 4096); }

    let flags = PageTableFlags::PRESENT | PageTableFlags::WRITABLE;

    for entry in mmap_res.entries() {
        let phys = PhysAddr(entry.base);
        let virt_hhdm = phys.to_virt(hhdm_offset);
        p4.map_region(virt_hhdm, phys, entry.length, flags, hhdm_offset);
        if entry.entry_type != limine::memory_map::EntryType::RESERVED {
            p4.map_region(VirtAddr(entry.base), phys, entry.length, flags, hhdm_offset);
        }
    }

    p4.map_region(VirtAddr(kaddr_res.virtual_base()), PhysAddr(kaddr_res.physical_base()), 0x1000 * 1024, flags, hhdm_offset);

    info!(console, "MMU", "Switching to Kernel Page Tables...");
    unsafe { p4.activate(p4_phys); }

    let heap_start = 0xFFFF_9000_0000_0000;
    let heap_size = 8 * 1024 * 1024;
    for i in (0..heap_size).step_by(4096) {
        let frame = mem::pmm::alloc_frame().expect("OOM: Heap allocation failed");
        p4.map_page(VirtAddr(heap_start + i as u64), frame, flags, hhdm_offset);
    }

    {
        let mut allocator = allocator::ALLOCATOR.lock();
        allocator.init(heap_start as usize, heap_size);
    }
    info!(console, "HEAP", "Global Allocator is online.");


    let root_cnode = cap::CNode::new(256);

    if let Some(frame) = mem::pmm::alloc_frame() {
        let mem_cap = Capability {
            object: CapObject::Memory { phys: frame, size_pages: 1 },
            rights: CapRights::all(),
            relation: Relation::Strong,
            token_sig: 0x1,
        };

        root_cnode.insert(0, mem_cap).unwrap();
        info!(console, "CAP", "Root capability created at slot 0");
    }

    root_cnode.mint(0, 10, Relation::Borrow, CapRights::READ | CapRights::WRITE).unwrap();

    if let Some(c) = root_cnode.get_cap(10) {
        info!(console, "CAP", "Slot 10 (Borrowed): {:?}", c.relation);
    }

    root_cnode.revoke(0);
    if let Some(c) = root_cnode.get_cap(10) {
        if !c.is_valid() {
            info!(console, "CAP", "Slot 10 successfully revoked.");
        }
    }
    let logo = r#"
            ###########
        ##################
      ######          ######
    #####                #####
   ####         #####     ####
  ####    ###      ######    ####
 ####    ####        ####    ####
 ####   ####          ###    ####
 ####   ###  ##        ###    ####
 ####   ########     ####    ####
 ####    ####        ####    ####
  ####    ######     ###    ####
   ####      #######        ####
    #####                #####
      ######          ######
        ##################
           ############
"#;
    let _ = writeln!(console, "{}", logo);

    hcf();
}

fn hcf() -> ! {
    loop {
        unsafe { asm!("hlt", options(nomem, nostack, preserves_flags)); }
    }
}

#[panic_handler]
fn rust_panic(info: &core::panic::PanicInfo) -> ! {
    let mut sp = unsafe { debug::serial::SerialPort::init() };
    let _ = writeln!(sp, "KERNEL PANIC: {:?}", info);
    hcf();
}