#![no_std]

extern crate alloc;

/// Minimal command shell — reads from PS/2 keyboard, writes to GOP + serial.
use rsnix::ps2::{self, KeyEvent};
use rsnix::{gop_print, gop_println, serial_print, serial_println};

/// Maximum input line length.
const LINE_BUF_SIZE: usize = 256;

/// Run the shell forever.
pub fn run() -> ! {
    gop_println!(
        "rs-dos v{} — type 'help' for commands",
        env!("CARGO_PKG_VERSION")
    );
    serial_println!(
        "========== rs-dos v{} ==========",
        env!("CARGO_PKG_VERSION")
    );

    loop {
        // Print prompt
        gop_print!("> ");
        serial_print!("> ");

        // Read one line from keyboard
        let line = read_line();

        // Parse and execute
        let trimmed = line.trim();
        if trimmed.is_empty() {
            continue;
        }

        let (cmd, args) = split_first_word(trimmed);

        match cmd {
            "help" => cmd_help(),
            "echo" => cmd_echo(args),
            "clear" => cmd_clear(),
            "halt" => cmd_halt(),
            "poweroff" | "exit" => cmd_poweroff(),
            "hello" => {
                gop_println!("World!");
                serial_println!("World!");
            }
            "info" => cmd_info(),
            "mem" => cmd_mem(),
            "mmap" => cmd_mmap(),
            "stress" => cmd_stress(args),
            _ => {
                gop_println!("unknown command: {}", cmd);
                serial_println!("unknown command: {}", cmd);
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Line input (PS/2 keyboard, no heap)
// ---------------------------------------------------------------------------
fn read_line() -> heapless_str::String<LINE_BUF_SIZE> {
    let mut buf = heapless_str::String::<LINE_BUF_SIZE>::new();
    loop {
        match ps2::poll_key() {
            KeyEvent::Enter => {
                rsnix::gop::put_char('\n');
                rsnix::serial::write_str("\r\n");
                return buf;
            }
            KeyEvent::Char(c) if c.is_ascii_graphic() || c == ' ' => {
                if buf.push(c).is_ok() {
                    rsnix::gop::put_char(c);
                    // Echo to serial too
                    let mut s = [0u8; 4];
                    let s = c.encode_utf8(&mut s);
                    rsnix::serial::write_str(s);
                }
            }
            KeyEvent::Backspace => {
                if buf.pop().is_some() {
                    // GOP backspace: move cursor back and erase
                    rsnix::gop::put_char('\x08');
                    // Serial backspace: send BS SP BS (standard terminal protocol)
                    rsnix::serial::write_str("\x08 \x08");
                }
            }
            KeyEvent::Tab => {
                for _ in 0..4 {
                    if buf.push(' ').is_ok() {
                        rsnix::gop::put_char(' ');
                        rsnix::serial::write_str(" ");
                    }
                }
            }
            _ => {}
        }
    }
}

/// Split a command line into the first word and the rest.
fn split_first_word(s: &str) -> (&str, &str) {
    let s = s.trim_start();
    if let Some(i) = s.find(char::is_whitespace) {
        (&s[..i], s[i..].trim_start())
    } else {
        (s, "")
    }
}

// ---------------------------------------------------------------------------
// Built-in commands
// ---------------------------------------------------------------------------

fn cmd_help() {
    gop_println!("available commands:");
    gop_println!("  help     — show this message");
    gop_println!("  echo     — echo arguments");
    gop_println!("  clear    — clear screen");
    gop_println!("  hello    — print 'World!'");
    gop_println!("  info     — show system info");
    gop_println!("  mem      — show memory info");
    gop_println!("  mmap     — dump the UEFI memory map");
    gop_println!("  stress   — stress-test heap/page allocators");
    gop_println!("  halt     — halt the system (hlt loop)");
    gop_println!("  poweroff — shut down via ACPI (ResetSystem)");
    gop_println!("  exit     — alias for poweroff");
    serial_println!("available commands:");
    serial_println!("  help, echo, clear, hello, info, mem, mmap, stress, halt, poweroff, exit");
}

fn cmd_echo(args: &str) {
    gop_println!("{}", args);
}

fn cmd_clear() {
    rsnix::gop::clear_screen();
}

fn cmd_info() {
    // Test that the global allocator works: build a String on the heap
    let mut s = alloc::string::String::from("rs-dos v");
    s.push_str(env!("CARGO_PKG_VERSION"));
    s.push_str(" — rsnix kernel");
    gop_println!("{}", s);
    gop_println!("arch: x86_64, freestanding, no userspace");
    serial_println!("{}", s);
    serial_println!("arch: x86_64, freestanding, no userspace");
    // String is freed (dropped) when this scope ends — tests dealloc too
}

fn cmd_mem() {
    let (map, key) = rsnix::memory::get_map();

    let mut total_kb = 0u64;
    let mut conventional_kb = 0u64;
    let mut loader_kb = 0u64;
    let mut runtime_kb = 0u64;
    let mut boot_code_kb = 0u64;
    let mut boot_data_kb = 0u64;
    let mut reserved_kb = 0u64;
    let mut acpi_kb = 0u64;
    let mut mmio_kb = 0u64;

    for desc in map {
        let pages = desc.number_of_pages;
        let bytes_4k = pages * 4;
        total_kb += bytes_4k;
        match desc.kind {
            1 | 2 => loader_kb += bytes_4k, // LoaderCode + LoaderData
            3 => boot_code_kb += bytes_4k,
            4 => boot_data_kb += bytes_4k,
            5 | 6 => runtime_kb += bytes_4k, // RuntimeCode + RuntimeData
            7 => conventional_kb += bytes_4k,
            9 | 10 => acpi_kb += bytes_4k, // ACPIReclaim + ACPINVS
            11 | 12 => mmio_kb += bytes_4k,
            _ => reserved_kb += bytes_4k, // Reserved, Unusable, Pal, ...
        }
    }

    // Build a big heap string with format! to test the allocator
    let mut out = alloc::string::String::new();
    use alloc::fmt::Write;
    let _ = writeln!(out, "memory map (key=0x{:x}, {} entries)", key, map.len());
    let _ = writeln!(
        out,
        "  usable    {:4} MB   Conventional",
        conventional_kb / 1024
    );
    let _ = writeln!(
        out,
        "  boot code {:4} MB   BootServicesCode",
        boot_code_kb / 1024
    );
    let _ = writeln!(
        out,
        "  boot data {:4} MB   BootServicesData",
        boot_data_kb / 1024
    );
    let _ = writeln!(
        out,
        "  loader    {:4} MB   LoaderCode + LoaderData",
        loader_kb / 1024
    );
    let _ = writeln!(
        out,
        "  runtime   {:4} MB   RuntimeServices",
        runtime_kb / 1024
    );
    let _ = writeln!(out, "  ACPI      {:4} MB", acpi_kb / 1024);
    let _ = writeln!(out, "  MMIO      {:4} MB", mmio_kb / 1024);
    let _ = writeln!(out, "  reserved  {:4} MB", reserved_kb / 1024);
    let _ = writeln!(
        out,
        "  raw total {:4} MB   (incl. PCIe MMIO)",
        total_kb / 1024
    );

    // Heap (free-list allocator with coalescing)
    let h = rsnix::memory::heap_stats();
    let _ = writeln!(out, "heap:");
    let _ = writeln!(
        out,
        "  range    {:#x}..{:#x}  {} KB total",
        h.heap_start,
        h.heap_end,
        h.heap_size / 1024
    );
    let _ = writeln!(
        out,
        "  free     {} KB in {} block(s), allocated {} KB",
        h.free_bytes / 1024,
        h.free_blocks,
        h.heap_size.saturating_sub(h.free_bytes) / 1024
    );

    // Physical pages (bitmap allocator)
    let ps = rsnix::memory::page_stats();
    let _ = writeln!(out, "pages:");
    let _ = writeln!(
        out,
        "  span     base {:#x}, {} MB covered",
        ps.base_page * rsnix::memory::PAGE_SIZE,
        ps.span_bytes / 1024 / 1024
    );
    let _ = writeln!(
        out,
        "  free     {} pages ({} MB), {} used",
        ps.free_pages,
        ps.free_pages * rsnix::memory::PAGE_SIZE / 1024 / 1024,
        ps.used_pages
    );

    gop_print!("{}", out);
    serial_print!("{}", out);
}

// ---------------------------------------------------------------------------
// mmap — dump every entry of the UEFI memory map with type names
// ---------------------------------------------------------------------------
fn cmd_mmap() {
    let (map, key) = rsnix::memory::get_map();
    let mut out = alloc::string::String::new();
    use alloc::fmt::Write;
    let _ = writeln!(out, "memory map (key=0x{:x}, {} entries):", key, map.len());
    let _ = writeln!(
        out,
        "  #  type              start           pages      size"
    );
    for (i, d) in map.iter().enumerate() {
        let _ = writeln!(
            out,
            "  {:>2} {:<16} {:#014x} {:>9} {:>8} KB",
            i,
            rsnix::uefi::memory_type_name(d.kind),
            d.physical_start,
            d.number_of_pages,
            d.number_of_pages * 4
        );
    }
    gop_print!("{}", out);
    serial_print!("{}", out);
}

// ---------------------------------------------------------------------------
// stress — allocate <chunks> blocks of <kb> KB, hold them, then free all.
// Pushes the heap past its base region so growth into the page allocator
// is exercised; also verifies a clean return to the free-list after.
// ---------------------------------------------------------------------------
fn cmd_stress(args: &str) {
    let mut it = args.split_whitespace();
    let chunks: usize = match it.next().and_then(|s| s.parse().ok()) {
        Some(n) if (1..=64).contains(&n) => n,
        _ => {
            gop_println!("usage: stress <chunks(1-64)> <kb>");
            serial_println!("usage: stress <chunks(1-64)> <kb>");
            return;
        }
    };
    let kb: usize = match it.next().and_then(|s| s.parse().ok()) {
        Some(n) if n > 0 && n <= 1024 * 1024 => n,
        _ => {
            gop_println!("usage: stress <chunks> <kb>");
            serial_println!("usage: stress <chunks> <kb>");
            return;
        }
    };

    let mut out = alloc::string::String::new();
    use alloc::fmt::Write;

    let mut allocated = 0usize;
    {
        let mut blocks: [Option<alloc::vec::Vec<u8>>; 64] = core::array::from_fn(|_| None);
        for i in 0..chunks {
            let mut b = alloc::vec::Vec::<u8>::new();
            match b.try_reserve_exact(kb * 1024) {
                Ok(()) => {
                    b.resize(kb * 1024, 0x42);
                    blocks[i] = Some(b);
                    allocated += 1;
                }
                Err(_) => {
                    let _ = writeln!(out, "  block {}: out of memory, stopping", i);
                    break;
                }
            }
        }
        let h = rsnix::memory::heap_stats();
        let ps = rsnix::memory::page_stats();
        let _ = writeln!(
            out,
            "  held {} x {} KB = {} KB",
            allocated,
            kb,
            allocated * kb
        );
        let _ = writeln!(
            out,
            "  heap free {} KB ({} blocks), pages free {} MB",
            h.free_bytes / 1024,
            h.free_blocks,
            ps.free_pages * 4 / 1024
        );
        // blocks dropped here → everything freed
    }
    let h = rsnix::memory::heap_stats();
    let _ = writeln!(
        out,
        "  freed: heap free {} KB in {} blocks",
        h.free_bytes / 1024,
        h.free_blocks
    );

    gop_print!("{}", out);
    serial_print!("{}", out);
}

fn cmd_halt() -> ! {
    gop_println!("halting...");
    serial_println!("halting...");
    rsnix::power::halt()
}

fn cmd_poweroff() -> ! {
    gop_println!("shutting down...");
    serial_println!("shutting down...");
    rsnix::power::poweroff()
}

// ---------------------------------------------------------------------------
// Minimal fixed-capacity string for no-stack use (avoids alloc dependency)
// ---------------------------------------------------------------------------
mod heapless_str {
    use core::fmt;

    /// A fixed-capacity string allocated entirely on the stack.
    pub struct String<const N: usize> {
        buf: [u8; N],
        len: usize,
    }

    impl<const N: usize> String<N> {
        pub const fn new() -> Self {
            String {
                buf: [0u8; N],
                len: 0,
            }
        }

        pub fn as_str(&self) -> &str {
            // SAFETY: we only write valid ASCII bytes.
            unsafe { core::str::from_utf8_unchecked(&self.buf[..self.len]) }
        }

        pub fn push(&mut self, c: char) -> Result<(), ()> {
            if self.len >= N {
                return Err(());
            }
            self.buf[self.len] = c as u8;
            self.len += 1;
            Ok(())
        }

        pub fn pop(&mut self) -> Option<char> {
            if self.len == 0 {
                return None;
            }
            self.len -= 1;
            Some(self.buf[self.len] as char)
        }

        pub fn is_empty(&self) -> bool {
            self.len == 0
        }

        pub fn clear(&mut self) {
            self.len = 0;
        }

        pub fn trim(&self) -> &str {
            self.as_str().trim()
        }
    }

    impl<const N: usize> fmt::Display for String<N> {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            write!(f, "{}", self.as_str())
        }
    }

    impl<const N: usize> fmt::Debug for String<N> {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            write!(f, "{:?}", self.as_str())
        }
    }
}