#![no_std]
#![no_main]

mod board;
mod boot_fs;
mod hw;
mod kernel_cell;
mod mmu;
mod platform;
mod trap;

use core::arch::global_asm;
use hw::uart::Uart;
use kernel_cell::KernelCell;
use types::id::{ProcessId, ThreadId};

static KERNEL: KernelCell<kernel::Kernel<platform::RiscVPlatform>> = KernelCell::uninit();

pub fn kernel_ref() -> &'static mut kernel::Kernel<platform::RiscVPlatform> {
    KERNEL.get_mut()
}

global_asm!(
    r#"
.section .text.entry
.globl _start
_start:
    la sp, _stack_top
    call kmain
spin:
    wfi
    j spin
"#
);

unsafe extern "C" {
    static _bss_start: u8;
    static _bss_end: u8;
    static _heap_start: u8;
}

fn zero_bss() {
    unsafe {
        let start = &_bss_start as *const u8 as *mut u8;
        let end = &_bss_end as *const u8 as *mut u8;
        let len = end as usize - start as usize;
        core::ptr::write_bytes(start, 0, len);
    }
}

const USER_VA_BASE: usize = 0x10000;
const MAX_FILE_SIZE: usize = 64 * 1024;
const SHARED_BUF_VA: usize = 0x30000;
const SHARED_BUF_PAGES: usize = 16;
const STORAGE_VIRTIO_VA: usize = 0x20000;
const STORAGE_VQUEUE_VA: usize = 0x21000;
const STORAGE_REQHDR_VA: usize = 0x23000;
const STORAGE_PHYSINFO_VA: usize = 0x24000;

const DISPLAY_GPU_VA: usize = 0x20000;
const DISPLAY_VQUEUE_VA: usize = 0x21000;
const DISPLAY_CMDBUF_VA: usize = 0x23000;
const DISPLAY_PHYSINFO_VA: usize = 0x24000;
const DISPLAY_FB_VA: usize = 0x100000;
const DISPLAY_FB_WIDTH: u32 = 640;
const DISPLAY_FB_HEIGHT: u32 = 480;
const DISPLAY_FB_BPP: usize = 4;
const DISPLAY_FB_PAGES: usize = (DISPLAY_FB_WIDTH as usize * DISPLAY_FB_HEIGHT as usize * DISPLAY_FB_BPP + 4095) / 4096;

#[unsafe(no_mangle)]
extern "C" fn kmain() -> ! {
    zero_bss();

    let uart = Uart::new();
    uart.write_str("k-os: booted on RISC-V\n");

    let heap_start = unsafe { &_heap_start as *const u8 as usize };
    mmu::init_allocator(heap_start);
    let kernel_pt = mmu::init_kernel_page_table();
    mmu::activate(kernel_pt);
    uart.write_str("k-os: MMU enabled (Sv39)\n");

    trap::init();
    hw::timer::init();
    hw::plic::Plic::init();
    uart.write_str("k-os: traps and timer initialized\n");

    let riscv_platform = platform::RiscVPlatform::new(kernel_pt as *const mmu::PageTable as usize);
    let pool = domain::memory::page_pool::PagePool::new(
        types::address::PhysicalAddress::from_raw(board::USER_POOL_BASE),
        board::USER_POOL_PAGES,
    );
    KERNEL.init(kernel::Kernel::new(riscv_platform, pool));

    // --- Read init + storage from disk ---
    let mut blk = hw::virtio_blk::VirtioBlk::probe();
    if blk.is_none() {
        halt(&uart, "no block device");
    }
    uart.write_str("k-os: VirtIO block device found\n");

    // Boot init (pid 0) — kern.boot() grants ProcessCreate cap at slot 0
    let (init_pid, init_tid) = kernel_ref()
        .boot(types::bounded::Priority::DEFAULT, types::time::TimeSlice::DEFAULT)
        .unwrap_or_else(|_| halt(&uart, "boot failed"));

    let init_entry = load_elf_into_process(&uart, blk.as_mut(), kernel_pt, init_pid, b"init.elf");
    if init_entry == 0 { halt(&uart, "init.elf not found"); }
    setup_thread(init_tid, init_entry);
    uart.write_str("k-os: init loaded\n");

    // Load storage server
    let storage_pid = kernel_ref().create_process_internal()
        .unwrap_or_else(|| halt(&uart, "storage process failed"));
    let storage_entry = load_elf_into_process(&uart, blk.as_mut(), kernel_pt, storage_pid, b"storage.elf");
    if storage_entry == 0 { halt(&uart, "storage.elf not found"); }
    let storage_tid = kernel_ref().spawn_thread_internal(storage_pid)
        .unwrap_or_else(|| halt(&uart, "storage thread failed"));
    setup_thread(storage_tid, storage_entry);
    uart.write_str("k-os: storage loaded\n");

    // Load display server
    let display_pid = kernel_ref().create_process_internal()
        .unwrap_or_else(|| halt(&uart, "display process failed"));
    let display_entry = load_elf_into_process(&uart, blk.as_mut(), kernel_pt, display_pid, b"display.elf");
    if display_entry == 0 { halt(&uart, "display.elf not found"); }
    let display_tid = kernel_ref().spawn_thread_internal(display_pid)
        .unwrap_or_else(|| halt(&uart, "display thread failed"));
    setup_thread(display_tid, display_entry);
    uart.write_str("k-os: display loaded\n");

    // Done reading from disk — drop the adapter's VirtIO driver
    let virtio_base = blk.as_ref().map(|b| b.base()).unwrap_or(0);
    drop(blk);

    // --- Set up VirtIO for storage server ---
    let storage_pt = kernel_ref().platform_mut().page_table_for(storage_pid);
    let storage_pt_ref = unsafe { &mut *(storage_pt as *mut mmu::PageTable) };

    // Map VirtIO MMIO into storage's address space
    mmu::map(storage_pt_ref, mmu::Mapping {
        va: types::address::VirtualAddress::from_raw(STORAGE_VIRTIO_VA),
        pa: types::address::PhysicalAddress::from_raw(virtio_base),
        pages: 1,
        flags: mmu::Sv39Flags::user_rw(),
    });

    // Allocate virtqueue pages (2 pages) and map into storage
    let vqueue_phys = mmu::allocator::alloc_zeroed_page();
    mmu::allocator::alloc_zeroed_page(); // second page (contiguous)
    for i in 0..2usize {
        mmu::map(storage_pt_ref, mmu::Mapping {
            va: types::address::VirtualAddress::from_raw(STORAGE_VQUEUE_VA + i * 4096),
            pa: types::address::PhysicalAddress::from_raw(vqueue_phys + i * 4096),
            pages: 1,
            flags: mmu::Sv39Flags::user_rw(),
        });
    }

    // Allocate request header page and map into storage
    let reqhdr_phys = mmu::allocator::alloc_zeroed_page();
    mmu::map(storage_pt_ref, mmu::Mapping {
        va: types::address::VirtualAddress::from_raw(STORAGE_REQHDR_VA),
        pa: types::address::PhysicalAddress::from_raw(reqhdr_phys),
        pages: 1,
        flags: mmu::Sv39Flags::user_rw(),
    });

    // Allocate phys info page — storage reads DMA addresses from here
    let physinfo_phys = mmu::allocator::alloc_zeroed_page();
    mmu::map(storage_pt_ref, mmu::Mapping {
        va: types::address::VirtualAddress::from_raw(STORAGE_PHYSINFO_VA),
        pa: types::address::PhysicalAddress::from_raw(physinfo_phys),
        pages: 1,
        flags: mmu::Sv39Flags::user_rw(),
    });
    // Write physical addresses: [req_hdr_pa, status_pa, sector_buf_pa]
    unsafe {
        let info = physinfo_phys as *mut u64;
        core::ptr::write_volatile(info, reqhdr_phys as u64);          // req_hdr_pa
        core::ptr::write_volatile(info.add(1), (reqhdr_phys + 16) as u64); // status_pa
        core::ptr::write_volatile(info.add(2), (reqhdr_phys + 32) as u64); // sector_buf_pa
    }

    // Re-initialize VirtIO device with storage's virtqueue
    reinit_virtio_for_storage(virtio_base, vqueue_phys);
    uart.write_str("k-os: VirtIO configured for storage\n");

    // --- Set up VirtIO GPU for display server ---
    let display_pt = kernel_ref().platform_mut().page_table_for(display_pid);
    let display_pt_ref = unsafe { &mut *(display_pt as *mut mmu::PageTable) };

    if let Some(gpu) = hw::virtio_gpu::VirtioGpu::probe() {
        let gpu_base = gpu.base();

        // Map GPU MMIO into display's address space
        mmu::map(display_pt_ref, mmu::Mapping {
            va: types::address::VirtualAddress::from_raw(DISPLAY_GPU_VA),
            pa: types::address::PhysicalAddress::from_raw(gpu_base),
            pages: 1,
            flags: mmu::Sv39Flags::user_rw(),
        });

        // Allocate GPU virtqueue (2 pages)
        let gpu_vq_phys = mmu::allocator::alloc_zeroed_page();
        mmu::allocator::alloc_zeroed_page();
        for i in 0..2usize {
            mmu::map(display_pt_ref, mmu::Mapping {
                va: types::address::VirtualAddress::from_raw(DISPLAY_VQUEUE_VA + i * 4096),
                pa: types::address::PhysicalAddress::from_raw(gpu_vq_phys + i * 4096),
                pages: 1,
                flags: mmu::Sv39Flags::user_rw(),
            });
        }

        // Allocate command buffer (1 page)
        let gpu_cmd_phys = mmu::allocator::alloc_zeroed_page();
        mmu::map(display_pt_ref, mmu::Mapping {
            va: types::address::VirtualAddress::from_raw(DISPLAY_CMDBUF_VA),
            pa: types::address::PhysicalAddress::from_raw(gpu_cmd_phys),
            pages: 1,
            flags: mmu::Sv39Flags::user_rw(),
        });

        // Allocate framebuffer (300 pages, contiguous)
        let fb_phys = mmu::allocator::alloc_zeroed_page();
        for _ in 1..DISPLAY_FB_PAGES {
            mmu::allocator::alloc_zeroed_page();
        }
        for i in 0..DISPLAY_FB_PAGES {
            mmu::map(display_pt_ref, mmu::Mapping {
                va: types::address::VirtualAddress::from_raw(DISPLAY_FB_VA + i * 4096),
                pa: types::address::PhysicalAddress::from_raw(fb_phys + i * 4096),
                pages: 1,
                flags: mmu::Sv39Flags::user_rw(),
            });
        }

        // Write phys-info for display server
        let gpu_physinfo = mmu::allocator::alloc_zeroed_page();
        mmu::map(display_pt_ref, mmu::Mapping {
            va: types::address::VirtualAddress::from_raw(DISPLAY_PHYSINFO_VA),
            pa: types::address::PhysicalAddress::from_raw(gpu_physinfo),
            pages: 1,
            flags: mmu::Sv39Flags::user_rw(),
        });
        unsafe {
            let info = gpu_physinfo as *mut u64;
            core::ptr::write_volatile(info, gpu_cmd_phys as u64);
            core::ptr::write_volatile(info.add(1), (gpu_cmd_phys + 512) as u64); // response area
            core::ptr::write_volatile(info.add(2), fb_phys as u64);
            core::ptr::write_volatile(info.add(3), DISPLAY_FB_WIDTH as u64);
            core::ptr::write_volatile(info.add(4), DISPLAY_FB_HEIGHT as u64);
        }

        // Initialize GPU device and set up framebuffer
        gpu.init_device(gpu_vq_phys);
        gpu.setup_framebuffer(gpu_cmd_phys, gpu_vq_phys, fb_phys, DISPLAY_FB_WIDTH, DISPLAY_FB_HEIGHT);
        uart.write_str("k-os: VirtIO GPU configured\n");
    } else {
        uart.write_str("k-os: no GPU found (display will be headless)\n");
    }

    // --- Allocate shared buffer and map into both processes ---
    let init_pt = kernel_ref().platform_mut().page_table_for(init_pid);
    let init_pt_ref = unsafe { &mut *(init_pt as *mut mmu::PageTable) };

    for i in 0..SHARED_BUF_PAGES {
        let page = mmu::allocator::alloc_zeroed_page();
        let va = SHARED_BUF_VA + i * 4096;
        mmu::map(init_pt_ref, mmu::Mapping {
            va: types::address::VirtualAddress::from_raw(va),
            pa: types::address::PhysicalAddress::from_raw(page),
            pages: 1,
            flags: mmu::Sv39Flags::user_rw(),
        });
        mmu::map(storage_pt_ref, mmu::Mapping {
            va: types::address::VirtualAddress::from_raw(va),
            pa: types::address::PhysicalAddress::from_raw(page),
            pages: 1,
            flags: mmu::Sv39Flags::user_rw(),
        });
    }
    uart.write_str("k-os: shared buffer mapped\n");

    // --- Grant capabilities ---
    let all_perms = types::permission::PermissionSet::empty()
        .with(types::permission::Permission::Read)
        .with(types::permission::Permission::Write)
        .with(types::permission::Permission::Execute)
        .with(types::permission::Permission::Delegate);

    // Create storage endpoint
    let storage_ep = kernel_ref().create_endpoint()
        .unwrap_or_else(|_| halt(&uart, "endpoint create failed"));

    // Init caps: slot 0 = ProcessCreate (from boot), slot 1 = Process(storage), slot 2 = Endpoint
    kernel_ref().grant_capability(
        init_pid,
        domain::resource::Resource::Process(storage_pid),
        all_perms,
    ).unwrap_or_else(|_| halt(&uart, "grant Process cap failed"));
    kernel_ref().grant_capability(
        init_pid,
        domain::resource::Resource::Endpoint(storage_ep),
        all_perms,
    ).unwrap_or_else(|_| halt(&uart, "grant Endpoint to init failed"));

    // Storage caps: slot 0 = Endpoint (Read+Execute for receive/send)
    let storage_ep_perms = types::permission::PermissionSet::empty()
        .with(types::permission::Permission::Read)
        .with(types::permission::Permission::Execute);
    kernel_ref().grant_capability(
        storage_pid,
        domain::resource::Resource::Endpoint(storage_ep),
        storage_ep_perms,
    ).unwrap_or_else(|_| halt(&uart, "grant Endpoint to storage failed"));

    // Create display endpoint
    let display_ep = kernel_ref().create_endpoint()
        .unwrap_or_else(|_| halt(&uart, "display endpoint failed"));

    // Display caps: slot 0 = Endpoint (Read+Execute)
    kernel_ref().grant_capability(
        display_pid,
        domain::resource::Resource::Endpoint(display_ep),
        storage_ep_perms, // same Read+Execute
    ).unwrap_or_else(|_| halt(&uart, "grant Endpoint to display failed"));

    // Init also gets display endpoint (slot 3) and display process cap (slot 4) for IPC
    kernel_ref().grant_capability(
        init_pid,
        domain::resource::Resource::Endpoint(display_ep),
        all_perms,
    ).unwrap_or_else(|_| halt(&uart, "grant display ep to init failed"));
    kernel_ref().grant_capability(
        init_pid,
        domain::resource::Resource::Process(display_pid),
        all_perms,
    ).unwrap_or_else(|_| halt(&uart, "grant display proc to init failed"));

    uart.write_str("k-os: capabilities granted\n");

    // --- Enter userspace with init, passing shared_buf VA in a0 ---
    uart.write_str("k-os: entering userspace\n");
    enter_userspace(init_pid, init_entry, SHARED_BUF_VA);
}

fn reinit_virtio_for_storage(base: usize, vqueue_phys: usize) {
    use hw::virtio_blk::{mmio_read, mmio_write};

    // Reset
    mmio_write(base, 0x070, 0);
    while mmio_read(base, 0x070) != 0 {}

    // Acknowledge + Driver
    mmio_write(base, 0x070, 1); // ACK
    mmio_write(base, 0x070, mmio_read(base, 0x070) | 2); // DRIVER

    // Legacy: set guest page size
    mmio_write(base, 0x028, 4096);

    // Features: accept none
    mmio_write(base, 0x014, 0);
    mmio_write(base, 0x024, 0);
    mmio_write(base, 0x020, 0);

    // Queue 0 setup
    mmio_write(base, 0x030, 0); // select queue 0
    let queue_max = mmio_read(base, 0x034);
    let queue_num = 8u32.min(queue_max);
    mmio_write(base, 0x038, queue_num);

    // Point device to storage's virtqueue (physical page number)
    mmio_write(base, 0x040, (vqueue_phys / 4096) as u32);

    // DRIVER_OK
    mmio_write(base, 0x070, mmio_read(base, 0x070) | 4);
}

fn load_elf_into_process(
    _uart: &Uart,
    blk: Option<&mut hw::virtio_blk::VirtioBlk>,
    kernel_pt: &mmu::PageTable,
    pid: ProcessId,
    filename: &[u8],
) -> usize {
    let buf_base = mmu::allocator::alloc_zeroed_page();
    let file_pages = MAX_FILE_SIZE / 4096;
    for _ in 1..file_pages {
        mmu::allocator::alloc_zeroed_page();
    }
    let buf = unsafe { core::slice::from_raw_parts_mut(buf_base as *mut u8, MAX_FILE_SIZE) };

    let file_size = match blk {
        Some(b) => boot_fs::read_file_from_disk(b, filename, buf).unwrap_or(0),
        None => 0,
    };
    if file_size == 0 { return 0; }

    let elf_info = match elf::parse(&buf[..file_size]) {
        Ok(info) => info,
        Err(_) => return 0,
    };

    let user_pt = mmu::create_user_page_table(
        kernel_pt,
        types::address::PhysicalAddress::from_raw(0),
        0,
        types::address::VirtualAddress::from_raw(USER_VA_BASE),
    );

    for segment in elf_info.segments() {
        let page_start = segment.vaddr & !0xFFF;
        let offset_in_page = segment.vaddr & 0xFFF;
        let total_size = offset_in_page + segment.mem_size;
        let seg_pages = (total_size + 4095) / 4096;

        for p in 0..seg_pages {
            let page = mmu::allocator::alloc_zeroed_page();
            let page_va = page_start + p * 4096;
            let seg_file_start = if p == 0 { 0 } else { p * 4096 - offset_in_page };
            let dest_off = if p == 0 { offset_in_page } else { 0 };
            let copy_len = if seg_file_start < segment.file_size {
                (segment.file_size - seg_file_start).min(4096 - dest_off)
            } else { 0 };

            if copy_len > 0 {
                let src = segment.file_offset + seg_file_start;
                if src + copy_len <= file_size {
                    unsafe {
                        core::ptr::copy_nonoverlapping(
                            buf.as_ptr().add(src),
                            (page as *mut u8).add(dest_off),
                            copy_len,
                        );
                    }
                }
            }

            let flags = if segment.flags.execute {
                mmu::Sv39Flags::user_rx()
            } else if segment.flags.write {
                mmu::Sv39Flags::user_rw()
            } else {
                mmu::Sv39Flags::user_rx()
            };
            mmu::map(user_pt, mmu::Mapping {
                va: types::address::VirtualAddress::from_raw(page_va),
                pa: types::address::PhysicalAddress::from_raw(page),
                pages: 1,
                flags,
            });
        }
    }

    // Stack: 4 pages below USER_VA_BASE
    for i in 0..4usize {
        let sp = mmu::allocator::alloc_zeroed_page();
        mmu::map(user_pt, mmu::Mapping {
            va: types::address::VirtualAddress::from_raw(USER_VA_BASE - (4 - i) * 4096),
            pa: types::address::PhysicalAddress::from_raw(sp),
            pages: 1,
            flags: mmu::Sv39Flags::user_rw(),
        });
    }

    kernel_ref().platform_mut().set_page_table(pid, user_pt as *const mmu::PageTable as usize);
    elf_info.entry_point
}

fn setup_thread(tid: ThreadId, entry: usize) {
    if let Some(ctx) = kernel_ref().thread_context_mut(tid) {
        ctx.set_pc(entry);
        ctx.registers_mut()[2] = USER_VA_BASE;
        ctx.set_status(1 << 5);
    }
}

fn enter_userspace(pid: ProcessId, entry: usize, a0: usize) -> ! {
    let kern = kernel_ref();
    let pt_addr = kern.platform_mut().page_table_for(pid);
    if pt_addr != 0 {
        mmu::activate(unsafe { &*(pt_addr as *const mmu::PageTable) });
    }

    let kernel_sp: usize;
    unsafe { core::arch::asm!("mv {0}, sp", out(reg) kernel_sp) };
    trap::set_kernel_stack(kernel_sp);

    unsafe {
        core::arch::asm!(
            "csrc sstatus, {spp}",
            "csrs sstatus, {spie}",
            "csrw sepc, {entry}",
            "mv sp, {stack}",
            "mv a0, {arg}",
            "sret",
            spp = in(reg) 1usize << 8,
            spie = in(reg) 1usize << 5,
            entry = in(reg) entry,
            stack = in(reg) USER_VA_BASE,
            arg = in(reg) a0,
            options(noreturn),
        );
    }
}

fn halt(uart: &Uart, msg: &str) -> ! {
    uart.write_str("k-os: FATAL — ");
    uart.write_str(msg);
    uart.put_char(b'\n');
    loop { unsafe { core::arch::asm!("wfi") }; }
}

#[panic_handler]
fn panic(info: &core::panic::PanicInfo) -> ! {
    let uart = Uart::new();
    uart.write_str("PANIC: ");
    if let Some(location) = info.location() {
        uart.write_str(location.file());
        uart.put_char(b':');
        uart.write_dec(location.line() as usize);
    }
    uart.put_char(b'\n');
    loop { unsafe { core::arch::asm!("wfi") }; }
}