// Copyright (c) 2025 shibedrill
// SPDX-License-Identifier: MIT

use crate::boot::modules::MODULE_RESPONSE;
use crate::boot::params::EXECUTABLE_FILE_RESPONSE;
use crate::{LOGGER, LogLevel, format};
use alloc::string::String;

use free_list::{AllocError, FreeList, PAGE_SIZE, PageRange};
use lazy_static::lazy_static;
use limine::response::MemoryMapResponse;
use limine::{
    memory_map::EntryType,
    request::{ExecutableAddressRequest, HhdmRequest, MemoryMapRequest, PagingModeRequest},
    response::HhdmResponse,
};
use spin::Mutex;
use talc::*;

pub extern crate alloc;

pub static FREELIST: Mutex<FreeList<32>> = Mutex::new(FreeList::<32>::new());
pub static HW_FREELIST: Mutex<FreeList<32>> = Mutex::new(FreeList::<32>::new());

lazy_static! {
    pub static ref MEMMAP_RESPONSE: &'static MemoryMapResponse = MEMMAP_REQUEST
        .get_response()
        .expect("Bootloader did not supply memory map");
}

pub fn round_up_to(multiple: usize, input: usize) -> usize {
    if input.is_multiple_of(multiple) {
        input
    } else {
        input - (input % multiple) + multiple
    }
}

fn round_down_to(multiple: usize, input: usize) -> usize {
    if input.is_multiple_of(multiple) {
        input
    } else {
        input - (input % multiple)
    }
}

pub fn deallocate_usable() -> Result<(), AllocError> {
    init_statics();
    let mut any_error: Option<AllocError> = None;
    let mut bytes_deallocated: u64 = 0;
    for entry in MEMMAP_RESPONSE.entries() {
        if (entry.entry_type == EntryType::USABLE)
            | (entry.entry_type == EntryType::BOOTLOADER_RECLAIMABLE)
        {
            if let Err(error) = deallocate(**entry) {
                any_error = Some(error);
            }
            bytes_deallocated += entry.length;
        }
    }
    log_trace!("Deallocated 0x{:x} bytes", bytes_deallocated);
    if let Some(error) = any_error {
        Err(error)
    } else {
        Ok(())
    }
}

pub fn deallocate_hardware() -> Result<(), AllocError> {
    let mut any_error: Option<AllocError> = None;
    let mut bytes_deallocated: u64 = 0;
    for entry in MEMMAP_RESPONSE.entries() {
        if (entry.entry_type == EntryType::ACPI_NVS)
            | (entry.entry_type == EntryType::RESERVED)
            | (entry.entry_type == EntryType::FRAMEBUFFER)
        {
            if let Err(error) = {
                let base: usize = round_down_to(PAGE_SIZE, entry.base as usize);
                let mut length = entry.length as usize + (entry.base as usize - base);
                length = round_up_to(PAGE_SIZE, length);
                let range = PageRange::from_start_len(base, length);
                match range {
                    Ok(range_inner) => unsafe { HW_FREELIST.lock().deallocate(range_inner) },
                    Err(err) => {
                        log_error!(
                            "Failed to convert range 0x{:x} @ 0x{:x} (original: 0x{:x} @ 0x{:x}): {}",
                            length,
                            base,
                            entry.length,
                            entry.base,
                            err
                        );
                        Ok(())
                    }
                }
            } {
                any_error = Some(error);
            }
            bytes_deallocated += entry.length;
        }
    }
    log_trace!(
        "Deallocated 0x{:x} bytes of hardware reserved memory",
        bytes_deallocated
    );
    if let Some(error) = any_error {
        Err(error)
    } else {
        Ok(())
    }
}

pub fn deallocate(entry: limine::memory_map::Entry) -> Result<(), AllocError> {
    let range = PageRange::from_start_len(entry.base as usize, entry.length as usize)
        .expect("Could not convert map entry to range!");
    unsafe { FREELIST.lock().deallocate(range) }
}

pub fn log_memory() {
    if !MEMMAP_RESPONSE.entries().is_empty() {
        let mut log_msg: String = String::from("Memory map:");
        let mut usable: u64 = 0;
        let mut reclaimable: u64 = 0;
        let mut hardware: u64 = 0;
        let mut unusable: u64 = 0;
        for entry in MEMMAP_RESPONSE.entries() {
            log_msg.push_str(&format!(
                "\n\t0x{:x} bytes ({} pages) @ 0x{:x}: {}",
                entry.length,
                entry.length as usize / PAGE_SIZE,
                entry.base,
                match entry.entry_type {
                    EntryType::ACPI_NVS => {
                        hardware += entry.length;
                        "ACPI (reserved)"
                    }
                    EntryType::ACPI_RECLAIMABLE => {
                        reclaimable += entry.length;
                        "ACPI (reclaimable)"
                    }
                    EntryType::BAD_MEMORY => {
                        unusable += entry.length;
                        "damaged/unusable"
                    }
                    EntryType::BOOTLOADER_RECLAIMABLE => {
                        reclaimable += entry.length;
                        "bootloader (reclaimable)"
                    }
                    #[allow(unreachable_patterns, deprecated)]
                    EntryType::EXECUTABLE_AND_MODULES | EntryType::KERNEL_AND_MODULES => {
                        unusable += entry.length;
                        "executable & modules"
                    }
                    EntryType::FRAMEBUFFER => {
                        hardware += entry.length;
                        "framebuffer"
                    }
                    EntryType::RESERVED => {
                        unusable += entry.length;
                        "reserved"
                    }
                    EntryType::USABLE => {
                        usable += entry.length;
                        "usable"
                    }
                    _ => "unidentified",
                }
            ));
        }
        log_trace!("{log_msg}");
        let total = usable + reclaimable + hardware + unusable;
        log_info!(
            "Memory report: \n\tFree: 0x{usable:x}\n\tAvailable: 0x{reclaimable:x}\n\tUsable: 0x{:x}\n\tHardware: 0x{hardware:x}\n\tUnusable: 0x{unusable:x}\n\tTotal: 0x{total:x}",
            usable + reclaimable,
        );
        log_trace!("Deallocating available memory...");
        let _ = deallocate_usable();
        log_trace!("Deallocating hardware reserved memory...");
        let _ = deallocate_hardware();
        let free_bytes = { FREELIST.lock().free_space() };
        log_info!(
            "Freelist: 0x{:x} bytes ({} pages)",
            free_bytes,
            free_bytes / PAGE_SIZE
        );
    } else {
        panic!("Memory map contains no entries");
    }
}

pub fn log_address() {
    if let Some(resp) = ADDRESS_REQUEST.get_response() {
        log_info!(
            "Kernel physical start address: 0x{:x}",
            resp.physical_base()
        );
        log_info!("Kernel virtual start address: 0x{:x}", resp.virtual_base());
    } else {
        log_warning!("No kernel address response provided.");
    }
    log_info!(
        "Higher Half direct map offset: 0x{:x}",
        HHDM_RESPONSE.offset()
    );
}

// Have to initialize all of these before reclaiming bootloader memory
fn init_statics() {
    let _ = EXECUTABLE_FILE_RESPONSE;
    let _ = HHDM_RESPONSE;
    let _ = MODULE_RESPONSE;
    let _ = MEMMAP_RESPONSE;
}

#[used]
#[unsafe(link_section = ".requests")]
pub static PAGING_REQUEST: PagingModeRequest = limine::request::PagingModeRequest::new();

#[used]
#[unsafe(link_section = ".requests")]
pub static ADDRESS_REQUEST: ExecutableAddressRequest =
    limine::request::ExecutableAddressRequest::new();

#[used]
#[unsafe(link_section = ".requests")]
pub static MEMMAP_REQUEST: MemoryMapRequest = limine::request::MemoryMapRequest::new();

#[used]
#[unsafe(link_section = ".requests")]
pub static HHDM_REQUEST: HhdmRequest = limine::request::HhdmRequest::new();

lazy_static! {
    pub static ref HHDM_RESPONSE: &'static HhdmResponse = HHDM_REQUEST
        .get_response()
        .expect("Did not get HHDM response from bootloader");
}

// TODO: 1mb kernel heap. Need to figure out how to make this less stupid...
static mut ARENA: [u8; 100000] = [0; 100000];

#[global_allocator]
static ALLOCATOR: Talck<spin::Mutex<()>, ClaimOnOom> =
    Talc::new(unsafe { ClaimOnOom::new(Span::from_array(core::ptr::addr_of!(ARENA).cast_mut())) })
        .lock();