//! The bootable kernel binary. All the real code is in the library
//! (src/lib.rs and friends); this file brings the machine up and starts the
//! shell. Build with `--features demos` to also see the step-by-step demos
//! of the earlier stages (paging, heap, threads, async, user programs).
//!
//! 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 bootloader::{BootInfo, entry_point};
use core::panic::PanicInfo;
use rustbelt::task::{Task, executor::Executor, keyboard};
use rustbelt::{fs, kprintln, println, process, serial_println, thread};
// Only the demos need these.
#[cfg(feature = "demos")]
use {
    alloc::{boxed::Box, rc::Rc, vec, vec::Vec},
    rustbelt::{allocator, memory, vga_buffer},
    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!("Rustbelt {} -- a hobby x86_64 kernel in Rust", env!("CARGO_PKG_VERSION"));

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

    #[cfg(feature = "demos")]
    {
        // 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();
    }

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

    // Paging, heap and the frame allocator. The `demos` feature does the
    // same steps one at a time and prints what happens (chapters 8-11).
    #[cfg(feature = "demos")]
    paging_demo(boot_info);
    #[cfg(not(feature = "demos"))]
    rustbelt::init_memory(boot_info);

    // --- More CPUs (smp.rs, apic.rs) -----------------------------------------
    // Start the other CPU cores; from now on threads run on all of them.
    let cpus = rustbelt::init_smp();
    kprintln!("[smp] {} CPU{} online", cpus, if cpus == 1 { "" } else { "s" });
    kprintln!("[fpu] {}", rustbelt::fpu::describe()); // stage 15.1: per-thread FPU/SSE/AVX

    // --- Preemptive threads demo (thread.rs) --------------------------------
    // From now on the timer interrupt switches between threads every tick.
    #[cfg(feature = "demos")]
    {
        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) = rustbelt::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);
    }

    // --- Virtual devices (virtio/): rng, second console, GPU framebuffer ---
    // Before the network: smoltcp gets its random seed from virtio-rng.
    for line in rustbelt::init_drivers() {
        kprintln!("[virtio] {}", line);
    }

    // --- Network (net/): loopback + virtio-net + smoltcp; DHCP runs in the
    // background ----------------------------------------------------------
    let net = rustbelt::init_net();
    for line in &net {
        kprintln!("[net] {}", line);
    }
    if !rustbelt::net::present() {
        kprintln!("[net] no virtio-net card (start QEMU with -device virtio-net-pci)");
    }

    // --- The shell (user/sh.rs), run by its own kernel thread ---------------
    thread::spawn("shell", shell_supervisor);

    // --- Cooperative async tasks (task/) ------------------------------------
    // The main thread runs the executor; the keyboard is just another task.
    let mut executor = Executor::new();
    #[cfg(feature = "demos")]
    executor.spawn(Task::new(example_task()));
    executor.spawn(Task::new(keyboard::print_keypresses()));
    executor.run(); // never returns: runs tasks, `hlt`s when idle
}

/// The shell and the programs live on the FAT32 disk.
const SHELL: &str = "/mnt/disk0/bin/sh";
/// Commands the shell runs in headless mode instead of reading the keyboard.
#[cfg(feature = "headless")]
const SELFTEST: &str = "/mnt/disk0/scripts/selftest.sh";

/// Kernel thread: start the shell, wait for it, start a new one when it
/// exits (like `init` respawning a login shell). In headless builds the
/// shell runs the self-test script and its exit code ends the QEMU run.
fn shell_supervisor() {
    #[cfg(feature = "demos")]
    run_demo_programs();

    if fs::lookup(SHELL).is_err() {
        kprintln!("[kernel] no {} -- was QEMU started with the disk image?", SHELL);
        #[cfg(feature = "headless")]
        rustbelt::exit_qemu(rustbelt::QemuExitCode::Failed);
        return;
    }
    // Start hvc0's tty (and its input thread `hvcd`) now rather than on
    // the first open, so its kernel stack isn't counted as a leak below.
    if rustbelt::virtio::console::present() {
        rustbelt::tty::hvc();
    }
    loop {
        #[cfg(feature = "headless")]
        let frames_before = rustbelt::memory::frames_in_use();
        #[allow(unused_mut)]
        let mut args = process::SpawnArgs::kernel("sh");
        #[cfg(feature = "headless")]
        args.args.push(alloc::string::String::from(SELFTEST));
        let status = match process::spawn_path(SHELL, args) {
            Ok(shell) => process::wait(&shell),
            Err(e) => {
                kprintln!("[kernel] could not start {}: {:?}", SHELL, e);
                #[cfg(feature = "headless")]
                rustbelt::exit_qemu(rustbelt::QemuExitCode::Failed);
                return;
            }
        };
        #[cfg(feature = "headless")]
        {
            let ok = status == process::ExitStatus::Exited(0);
            serial_println!("[{}] shell self-test finished: {:?}", if ok { "ok" } else { "failed" }, status);
            // Every process the self-test ran is gone: all their memory
            // (page tables, user pages, kernel stacks) must be back.
            serial_println!(
                "[memory] {} frames in use before the shell, {} after it and all its children exited",
                frames_before,
                rustbelt::memory::frames_in_use()
            );
            rustbelt::exit_qemu(if ok { rustbelt::QemuExitCode::Success } else { rustbelt::QemuExitCode::Failed });
        }
        kprintln!("[kernel] shell exited ({:?}); starting a new one", status);
    }
}

/// `demos` feature: the stage 5/6 user programs, one after the other.
#[cfg(feature = "demos")]
fn run_demo_programs() {
    for name in ["hello", "evil", "spin", "demo"] {
        let path = alloc::format!("/mnt/disk0/bin/{}", name);
        match process::spawn_path(&path, process::SpawnArgs::kernel(name)) {
            Ok(p) => {
                process::wait(&p);
            }
            Err(e) => kprintln!("[process] could not start {}: {:?}", path, e),
        }
    }
}

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

#[cfg(feature = "demos")]
async fn async_number() -> u32 {
    42
}

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

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

#[cfg(feature = "demos")]
/// 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
    }
}

#[cfg(feature = "demos")]
fn thread_a() {
    counter_thread("A", 5, 0);
}

#[cfg(feature = "demos")]
fn thread_b() {
    counter_thread("B", 7, 10);
}

#[cfg(feature = "demos")]
/// 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 "));
}

#[cfg(feature = "demos")]
/// 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());

    // --- 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")]
    rustbelt::exit_qemu(rustbelt::QemuExitCode::Failed);
    rustbelt::hlt_loop();
}