//! The bootable kernel binary. All the real code is in the library
//! (src/lib.rs and friends); this file just runs the demos in order.
//!
//! Boot flow: BIOS -> `bootloader` crate (switches to 64-bit long mode, sets up
//! paging, maps all physical memory) -> `kernel_main` below, with a
//! `BootInfo` describing the memory map.

#![no_std] // don't link the Rust standard library (it needs an OS)
#![no_main] // no normal `main`: the bootloader calls our entry point

extern crate alloc; // Box, Vec, Rc... once the heap is set up

use alloc::{boxed::Box, rc::Rc, vec, vec::Vec};
use bootloader::{BootInfo, entry_point};
use core::panic::PanicInfo;
use rust_kernel::task::{Task, executor::Executor, keyboard};
use rust_kernel::{allocator, fs, kprintln, memory, println, process, serial_println, thread};
use rust_kernel::vga_buffer;
use x86_64::{VirtAddr, structures::paging::{Page, Translate}};

// `entry_point!` defines the real `_start` for us and checks at compile time
// that `kernel_main` has the right signature (takes `&'static BootInfo`).
entry_point!(kernel_main);

fn kernel_main(boot_info: &'static BootInfo) -> ! {
    kprintln!("Hello from Tyler's Rust kernel!");

    rust_kernel::init(); // GDT, IDT, PICs, enable interrupts

    // Fire a breakpoint exception on purpose: the handler prints a message,
    // returns, and we carry on as if nothing happened.
    x86_64::instructions::interrupts::int3();
    kprintln!("Survived the breakpoint, still running.");

    #[cfg(feature = "stack-overflow")]
    {
        serial_println!("Overflowing the stack on purpose...");
        stack_overflow();
    }

    paging_demo(boot_info); // also sets up the heap, needed from here on

    // --- Preemptive threads (thread.rs) -----------------------------------
    // From now on the timer interrupt switches between these and the code
    // below (the "main" thread) every tick (~55 ms).
    thread::spawn("A", thread_a);
    thread::spawn("B", thread_b);
    thread::spawn("C", thread_c);
    kprintln!("[thread] spawned threads A, B, C (status on the top line)");

    // --- VFS, PCI, disks (fs/, pci.rs, block/) ------------------------------
    fs::init();
    let (pci_devices, storage_log) = rust_kernel::init_storage();
    for d in &pci_devices {
        serial_println!(
            "[pci] {:02x}:{:02x}.{} {:04x}:{:04x} {}",
            d.bus, d.device, d.function, d.vendor_id, d.device_id, d.class_name()
        );
    }
    kprintln!("[pci] {} devices on the PCI bus", pci_devices.len());
    for line in &storage_log {
        kprintln!("[disk] {}", line);
    }
    for (path, fs_type) in fs::mounts() {
        serial_println!("[vfs] mounted {} at {}", fs_type, path);
    }

    // --- User mode (process.rs, syscall.rs) ---------------------------------
    // The programs (built from user/*.rs) are ELF files on the FAT32 disk.
    // Each gets its own address space and runs in ring 3.
    const BIN: &str = "/mnt/disk0/bin";
    if fs::lookup(BIN).is_err() {
        kprintln!("[process] no {} -- was QEMU started with the disk image?", BIN);
    }
    let programs = [
        "hello", // syscalls + ramfs demo, exits 0
        "evil",  // touches kernel memory -> killed
        "spin",  // busy loop, gets preempted
        "demo",  // stage 6: FAT file, cwd, getdents, spawn + wait, pipes
        "echo",  // reads keyboard lines via read(0)
    ];
    for name in programs {
        let path = alloc::format!("{}/{}", BIN, name);
        if let Err(e) = process::spawn_path(&path, process::SpawnArgs::kernel(name)) {
            kprintln!("[process] could not start {}: {:?}", path, e);
        }
    }

    // --- Cooperative async tasks (task/) ------------------------------------
    // The main thread runs the executor; the keyboard is just another task.
    let mut executor = Executor::new();
    executor.spawn(Task::new(example_task()));
    executor.spawn(Task::new(keyboard::print_keypresses()));
    kprintln!("Type a line + Enter: the echo program (ring 3) reads it");
    executor.run(); // never returns: runs tasks, `hlt`s when idle
}

// --- async demo -------------------------------------------------------------

async fn async_number() -> u32 {
    42
}

async fn example_task() {
    let number = async_number().await;
    kprintln!("[async] async number: {}", number);
    for i in 1..=3 {
        rust_kernel::task::yield_now().await; // let other tasks run
        serial_println!("[async] example task step {} (after yield)", i);
    }
}

// --- thread demo ------------------------------------------------------------

/// Threads A and B: sleep a bit, count, show the count on the top line.
/// They print their first 5 counts to serial too (more would drown out the
/// keyboard echo in the terminal).
fn counter_thread(name: &str, period_ticks: u64, col: usize) {
    for n in 1u64.. {
        if n <= 5 {
            serial_println!("[thread {}] count {}", name, n);
        }
        vga_buffer::print_at(0, col, format_args!(" {}:{:>5} ", name, n));
        thread::sleep_ticks(period_ticks); // hlt -> other threads run meanwhile
    }
}

fn thread_a() {
    counter_thread("A", 5, 0);
}

fn thread_b() {
    counter_thread("B", 7, 10);
}

/// Thread C never sleeps or yields: a pure busy loop. It still only gets
/// its share of the CPU, because the timer preempts it. It finishes after a
/// while to show that threads can exit.
fn thread_c() {
    let mut count: u64 = 0;
    for round in 1..=5 {
        for _ in 0..2_000_000 {
            count = core::hint::black_box(count + 1); // don't let the compiler skip the work
        }
        serial_println!("[thread C] busy loop round {} (count = {})", round, count);
        vga_buffer::print_at(0, 20, format_args!(" C:{:>5} ", round));
    }
    serial_println!("[thread {}] done, exiting", thread::current_name());
    vga_buffer::print_at(0, 20, format_args!(" C: done "));
}

/// Chapters "Introduction to Paging", "Paging Implementation", "Heap Allocation".
fn paging_demo(boot_info: &'static BootInfo) {
    use x86_64::registers::control::Cr3;

    // --- Introduction to Paging -------------------------------------------
    // CR3 = physical address of the level-4 page table currently in use.
    let (level_4_frame, _) = Cr3::read();
    kprintln!("[paging] CR3 (level 4 table) at phys {:#x}", level_4_frame.start_address().as_u64());

    #[cfg(feature = "page-fault")]
    {
        kprintln!("[paging] writing to unmapped address 0xdeadbeaf...");
        unsafe { *(0xdeadbeaf as *mut u8) = 42 }; // -> page fault handler
    }

    // --- Paging Implementation --------------------------------------------
    let phys_mem_offset = VirtAddr::new(boot_info.physical_memory_offset);
    serial_println!("[paging] physical memory mapped at virt {:#x}", phys_mem_offset.as_u64());
    serial_println!(
        "[memory] {} memory map regions, {} KiB usable RAM",
        boot_info.memory_map.iter().count(),
        memory::usable_memory_bytes(&boot_info.memory_map) / 1024
    );

    // Safety: the bootloader mapped all physical memory at that offset, and
    // this is the only place we call `init`.
    let mut mapper = unsafe { memory::init(phys_mem_offset) };
    let mut frame_allocator = unsafe { memory::BootInfoFrameAllocator::init(&boot_info.memory_map) };

    // Translate some virtual addresses to physical ones by walking the tables.
    let addresses = [
        0xb8000,                         // VGA buffer: identity-mapped
        0x201008,                        // somewhere in our kernel's code
        0x0100_0020_1a10,                // somewhere in our stack
        boot_info.physical_memory_offset, // -> physical address 0
    ];
    for &address in &addresses {
        let virt = VirtAddr::new(address);
        let phys = mapper.translate_addr(virt);
        serial_println!("[paging] translate {:?} -> {:?}", virt, phys);
    }

    // Map a brand-new page (0xdeadbeaf000, previously unmapped) to the VGA
    // buffer frame, then write "New!" (white on red) through it.
    let page = Page::containing_address(VirtAddr::new(0xdeadbeaf000));
    memory::create_example_mapping(page, &mut mapper, &mut frame_allocator);
    let page_ptr: *mut u64 = page.start_address().as_mut_ptr();
    // Offset 400 u64s = 3200 bytes = row 20 of the screen (160 bytes/row).
    unsafe { page_ptr.offset(400).write_volatile(0x_f021_f077_f065_f04e) };
    kprintln!("[paging] mapped page 0xdeadbeaf000 -> VGA frame, wrote \"New!\" through it");

    // --- Heap Allocation --------------------------------------------------
    allocator::init_heap(&mut mapper, &mut frame_allocator).expect("heap initialization failed");
    kprintln!(
        "[heap] {} KiB heap at {:#x} (fixed-size-block allocator)",
        allocator::HEAP_SIZE / 1024,
        allocator::HEAP_START
    );

    let heap_value = Box::new(41);
    kprintln!("[heap] Box::new(41) = {} at {:p}", heap_value, heap_value);

    let mut vec = Vec::new();
    for i in 0..500 {
        vec.push(i);
    }
    let sum: u64 = vec.iter().sum();
    kprintln!("[heap] Vec of 500 numbers at {:p}, sum = {}", vec.as_ptr(), sum);

    let reference_counted = Rc::new(vec![1, 2, 3]);
    let cloned_reference = reference_counted.clone();
    kprintln!("[heap] Rc count after clone = {}", Rc::strong_count(&cloned_reference));
    core::mem::drop(reference_counted);
    kprintln!("[heap] Rc count after drop  = {}", Rc::strong_count(&cloned_reference));

    // From now on the kernel owns the frame allocator globally (processes
    // need frames for their page tables and pages).
    memory::init_frame_allocator(frame_allocator);
}

/// Recurse forever until we run off the end of the stack into the guard page.
#[cfg(feature = "stack-overflow")]
#[allow(unconditional_recursion)]
fn stack_overflow() {
    stack_overflow();
    // Stop the compiler turning the recursion into a loop (tail call).
    unsafe { core::ptr::read_volatile(&0u8) };
}

/// Called on `panic!`. Print the message to both outputs, then stop.
#[panic_handler]
fn panic(info: &PanicInfo) -> ! {
    println!("{}", info);
    serial_println!("KERNEL PANIC: {}", info);
    #[cfg(feature = "headless")]
    rust_kernel::exit_qemu(rust_kernel::QemuExitCode::Failed);
    rust_kernel::hlt_loop();
}