use types::address::{PhysicalAddress, VirtualAddress};
use types::id::ProcessId;

use crate::error::KernelError;

const USER_HEAP_START: usize = 0x40000;
const PAGE_SIZE: usize = 4096;
const MAX_ALLOCATIONS: usize = 16;

pub struct ProcessMemory {
    process: ProcessId,
    page_count: usize,
    limit: usize,
    next_heap_va: usize,
    allocations: [(PhysicalAddress, usize); MAX_ALLOCATIONS],
    alloc_count: usize,
}

impl ProcessMemory {
    pub const DEFAULT_PAGE_LIMIT: usize = 1024;

    pub fn new(process: ProcessId, limit: usize) -> Self {
        Self {
            process,
            page_count: 0,
            limit,
            next_heap_va: USER_HEAP_START,
            allocations: [(PhysicalAddress::from_raw(0), 0); MAX_ALLOCATIONS],
            alloc_count: 0,
        }
    }

    pub fn alloc_user_va(&mut self, pages: usize) -> VirtualAddress {
        let va = self.next_heap_va;
        self.next_heap_va += pages * PAGE_SIZE;
        VirtualAddress::from_raw(va)
    }

    pub fn record_allocation(&mut self, pa: PhysicalAddress, count: usize) {
        if self.alloc_count < MAX_ALLOCATIONS {
            self.allocations[self.alloc_count] = (pa, count);
            self.alloc_count += 1;
        }
    }

    pub fn allocations(&self) -> &[(PhysicalAddress, usize)] {
        &self.allocations[..self.alloc_count]
    }

    pub fn process(&self) -> ProcessId {
        self.process
    }

    pub fn page_count(&self) -> usize {
        self.page_count
    }

    pub fn limit(&self) -> usize {
        self.limit
    }

    pub fn grant(&mut self, count: usize) -> Result<(), KernelError> {
        if self.page_count + count > self.limit {
            return Err(KernelError::OutOfMemory);
        }
        self.page_count += count;
        Ok(())
    }

    pub fn release(&mut self, count: usize) -> Result<(), KernelError> {
        if count > self.page_count {
            return Err(KernelError::OutOfMemory);
        }
        self.page_count -= count;
        Ok(())
    }
}