#![no_std]
#![no_main]

mod console;
mod drivers;
mod efi;
mod elf;
mod font;
mod framebuffer;
mod fs;
mod gdt;
mod interrupts;
mod io;
mod ipc;
mod log;
mod memory;
mod paging;
mod services;
mod syscall;
mod task;
mod tty;
mod user;

use core::panic::PanicInfo;

use efi::{EFI_SUCCESS, EfiHandle, EfiStatus, EfiSystemTable};
use framebuffer::Framebuffer;
use memory::PhysicalMemoryManager;
use paging::{AddressSpace, PageFlags};
use tty::Tty;

static mut MEMORY_MAP_BUFFER: [u8; 65536] = [0; 65536];
static mut PHYSICAL_MEMORY: PhysicalMemoryManager = PhysicalMemoryManager::new();

#[unsafe(no_mangle)]
extern "efiapi" fn efi_main(
    image_handle: EfiHandle,
    system_table: *mut EfiSystemTable,
) -> EfiStatus {
    unsafe {
        io::disable_interrupts();

        let Some(boot_services) = efi::boot_services(system_table) else {
            halt_forever();
        };

        let Some(gop) = efi::locate_gop(boot_services) else {
            halt_forever();
        };

        let Some(framebuffer) = Framebuffer::from_gop(gop) else {
            halt_forever();
        };

        let mut tty = Tty::new(framebuffer);
        tty.clear();
        console::init(&raw mut tty);
        let mut logger = log::KernelLogger::new(&mut tty);
        logger.info(b"ZeroOS kernel initializing");

        drivers::platform::pic::remap_and_mask_all();
        logger.ok(b"pic");

        if let Some(unix_seconds) = efi::read_unix_time(system_table) {
            drivers::platform::timer::set_wall_clock_boot_time(unix_seconds);
            logger.ok(b"UEFI time");
        } else {
            logger.info(b"UEFI time unavailable");
        }

        let Some(memory_map) = efi::get_memory_map(
            boot_services,
            (&raw mut MEMORY_MAP_BUFFER).cast::<u8>(),
            65536,
        ) else {
            halt_forever();
        };
        logger.ok(b"UEFI memory map");

        let memory = &mut *(&raw mut PHYSICAL_MEMORY);
        task::init();
        memory.init_from_uefi_memory_map(memory_map);
        let stats = memory.stats();
        let _ = (stats.total_pages, stats.free_pages, stats.used_pages);
        logger.ok(b"physical memory manager");

        fs::vfs::init();
        logger.ok(b"vfs");
        ipc::init();
        logger.ok(b"ipc");

        let ahci_disks = drivers::storage::ahci::init(memory);
        let mut ext4_user_image = None;
        let mut ext4_user_path = None;
        if ahci_disks > 0 {
            logger.ok(b"ahci");
            if let Some(mut disk) = drivers::storage::ahci::first_disk() {
                match drivers::storage::partition::find_linux_partition(&mut disk) {
                    Ok(partition) => {
                        let mut root =
                            drivers::storage::partition::PartitionBlockDevice::new(disk, partition);
                        match fs::ext4::mount(&mut root) {
                            Ok(ext4) => {
                                logger.ok(b"ext4");
                                let _ = fs::ext4::mount_vfs_tree(&mut root, &ext4);
                                fs::ext4::register_root(root, ext4);
                                match fs::ext4::load_path(
                                    &mut root,
                                    &ext4,
                                    memory,
                                    b"/usr/bin/busybox",
                                ) {
                                    Ok(image) => {
                                        ext4_user_image = Some(image);
                                        ext4_user_path = Some(b"/usr/bin/busybox" as &'static [u8]);
                                        logger.ok(b"ext4 /usr/bin/busybox");
                                    }
                                    Err(_) => match fs::ext4::load_path(
                                        &mut root,
                                        &ext4,
                                        memory,
                                        b"/usr/bin/sh",
                                    ) {
                                        Ok(image) => {
                                            ext4_user_image = Some(image);
                                            ext4_user_path = Some(b"/usr/bin/sh" as &'static [u8]);
                                            logger.ok(b"ext4 /usr/bin/sh");
                                        }
                                        Err(_) => logger.info(b"ext4 userspace not loaded"),
                                    },
                                }
                            }
                            Err(_) => logger.info(b"ext4 not mounted"),
                        }
                    }
                    Err(_) => logger.info(b"no linux gpt partition"),
                }
            }
        } else {
            logger.info(b"no ahci disk");
        }

        let (loaded_user_image, loaded_user_path) = if let Some(image) = ext4_user_image {
            (Some(image), ext4_user_path)
        } else if let Some(image) =
            fs::uefi::load_file(image_handle, boot_services, memory, fs::uefi::BASH_PATH)
        {
            (Some(image), Some(b"/bin/bash" as &'static [u8]))
        } else if let Some(image) =
            fs::uefi::load_file(image_handle, boot_services, memory, fs::uefi::SH_PATH)
        {
            (Some(image), Some(b"/bin/sh" as &'static [u8]))
        } else {
            (None, None)
        };
        let userspace_image_found = loaded_user_image.is_some();
        if loaded_user_image.is_some() {
            if let (Some(image), Some(path)) = (&loaded_user_image, loaded_user_path) {
                let _ = fs::vfs::mount_static_file(path, image.ptr, image.len);
                let _ = fs::vfs::mount_static_file(b"/usr/bin/busybox", image.ptr, image.len);
                let _ = fs::vfs::mount_static_file(b"/bin/busybox", image.ptr, image.len);
                let _ = fs::vfs::mount_static_file(b"/usr/bin/sh", image.ptr, image.len);
                let _ = fs::vfs::mount_static_file(b"/usr/bin/bash", image.ptr, image.len);
                let _ = fs::vfs::mount_static_file(b"/bin/sh", image.ptr, image.len);
                let _ = fs::vfs::mount_static_file(b"/bin/bash", image.ptr, image.len);
            }
            logger.ok(b"userspace image loaded from esp");
        } else {
            logger.info(b"no /bin/bash or /bin/sh on esp, using built-in user program");
        }

        gdt::init();
        logger.ok(b"gdt");

        if interrupts::init(memory).is_none() {
            halt_forever();
        }
        logger.ok(b"idt");

        drivers::platform::timer::init();
        logger.ok(b"pit timer");

        drivers::input::ps2::init();
        logger.ok(b"ps/2 keyboard");

        let Some(mut address_space) = AddressSpace::new_kernel(memory) else {
            halt_forever();
        };
        logger.ok(b"address space");

        let mut userspace_program = None;
        let user_entry = if let Some(image) = loaded_user_image {
            match elf::load_user_elf(&image, memory) {
                Some(program) => {
                    logger.ok(b"elf userspace program");
                    userspace_program = Some(user::UserAux {
                        entry: program.entry,
                        phdr: program.phdr,
                        phent: program.phent,
                        phnum: program.phnum,
                    });
                    Some(program.entry)
                }
                None => {
                    logger.info(b"elf load failed, using built-in user program");
                    None
                }
            }
        } else {
            None
        };

        let Some(test_page) = memory.alloc_page() else {
            halt_forever();
        };
        if address_space
            .map_page(
                memory,
                0xffff_8000_0000_0000,
                test_page as u64,
                PageFlags::WRITABLE,
            )
            .is_none()
        {
            halt_forever();
        }
        let _ = address_space.translate(0xffff_8000_0000_0000);
        let _ = address_space.unmap_page(0xffff_8000_0000_0000);
        memory.free_page(test_page);
        let _ = (address_space.root_table(), address_space.identity_tables());
        logger.ok(b"paging smoke test");

        let Some(final_memory_map) = efi::get_memory_map(
            boot_services,
            (&raw mut MEMORY_MAP_BUFFER).cast::<u8>(),
            65536,
        ) else {
            logger.info(b"final UEFI memory map failed");
            halt_forever();
        };
        logger.ok(b"final UEFI memory map");

        let exit_status = ((*boot_services).exit_boot_services)(image_handle, final_memory_map.key);
        if exit_status != EFI_SUCCESS {
            logger.info(b"exit boot services failed");
            halt_forever();
        }
        logger.ok(b"exit boot services");

        address_space.activate();
        logger.ok(b"cr3");
        tty.clear();
        if user_entry.is_some() {
            if loaded_user_path == Some(b"/usr/bin/busybox" as &'static [u8]) {
                console::write(b"enter userspace: /usr/bin/busybox sh\r\n");
            } else if loaded_user_path == Some(b"/usr/bin/sh" as &'static [u8]) {
                console::write(b"enter userspace: /usr/bin/sh\r\n");
            } else if loaded_user_path == Some(b"/usr/bin/bash" as &'static [u8]) {
                console::write(b"enter userspace: /usr/bin/bash\r\n");
            } else {
                console::write(b"enter userspace: /bin/bash\r\n");
            }
        } else if userspace_image_found {
            console::write(b"/bin/bash found but ELF load failed; enter built-in userspace\r\n");
        } else {
            console::write(b"/bin/bash not found on ESP; enter built-in userspace\r\n");
        }
        syscall::init(memory, &raw mut address_space);
        io::enable_interrupts();
        if let Some(entry) = user_entry {
            if loaded_user_path == Some(b"/usr/bin/busybox" as &'static [u8]) {
                user::enter_elf_with_args(
                    entry,
                    &[b"/usr/bin/busybox", b"sh"],
                    user_program_or_default(user_entry, userspace_program),
                );
            } else if loaded_user_path == Some(b"/usr/bin/sh" as &'static [u8]) {
                user::enter_elf_with_args(
                    entry,
                    &[b"/usr/bin/sh"],
                    user_program_or_default(user_entry, userspace_program),
                );
            } else if loaded_user_path == Some(b"/usr/bin/bash" as &'static [u8]) {
                user::enter_elf_with_args(
                    entry,
                    &[b"/usr/bin/bash"],
                    user_program_or_default(user_entry, userspace_program),
                );
            } else if loaded_user_path == Some(b"/bin/sh" as &'static [u8]) {
                user::enter_elf_with_args(
                    entry,
                    &[b"/bin/sh"],
                    user_program_or_default(user_entry, userspace_program),
                );
            } else {
                user::enter_elf_with_args(
                    entry,
                    &[b"/bin/bash"],
                    user_program_or_default(user_entry, userspace_program),
                );
            }
        } else {
            user::enter();
        }
    }
}

#[allow(dead_code)]
fn kernel_loop(tty: &mut Tty) -> ! {
    loop {
        while let Some(ch) = drivers::input::ps2::read_char() {
            match ch {
                0x08 => tty.backspace(),
                b'\r' => tty.write_bytes(b"\r\n"),
                ch => tty.put_char(ch),
            }
        }
        halt();
    }
}

fn halt() {
    unsafe {
        core::arch::asm!("hlt", options(nomem, nostack, preserves_flags));
    }
}

fn user_program_or_default(entry: Option<u64>, aux: Option<user::UserAux>) -> user::UserAux {
    aux.unwrap_or(user::UserAux {
        entry: entry.unwrap_or(0),
        phdr: 0,
        phent: 0,
        phnum: 0,
    })
}

fn halt_forever() -> ! {
    loop {
        halt();
    }
}

#[panic_handler]
fn panic(_info: &PanicInfo) -> ! {
    halt_forever();
}