//!
//! Distributed physical-memory actor that owns a fixed physical capital range
//! and processes allocation requests asynchronously via a lock-free MPSC queue.
//!
//! ## Key changes vs. previous version
//!
//! 1. **Real allocator**: `PMActor` now contains a `BuddyAllocator` instead of a
//!    bare page counter. Allocations produce real relative frame indices that
//!    translate to absolute `PhysAddr` values.
//!
//! 2. **Closed response loop**: `process_messages()` returns `Vec<PMResponse>`.
//!    The caller (scheduler / boot code) routes each `PMResponse` to the
//!    requesting entity identified by `channel_id`.  `channel_id = 0` means
//!    "discard" and is safe to ignore.
//!
//! 3. **Revised message packing** — fits all fields into a single `u64` without
//!    truncating critical identifiers:
//!    ```text
//!    bits 63:56  OPCODE     (8 bits)   1=Alloc, 2=Free, 3=Carve
//!    bits 55:40  CHANNEL_ID (16 bits)  response routing; 0 = discard
//!    bits 39:20  ARG1       (20 bits)  size_pages / local_frame_idx / offset_pages
//!    bits 19: 0  ARG2       (20 bits)  token_sig (low 20 bits) / size_pages
//!    ```
//!    Maximum ARG size: 2^20 = 1 048 576 pages = 4 GiB per allocation.
//!    Maximum CHANNEL_ID: 65 535 simultaneous pending requesters.
//!
//! ## Ownership model
//!
//! `PMActor` holds a `Strong` capability over its entire managed range.
//! Sub-capabilities carved out via `Allocate` or `Carve` are `Strong` children;
//! freeing them (via `Free`) coalesces the buddy tree.  The actor **never**
//! touches the global `BitmapPMM` — that belongs to the kernel boot path.
//!
//! ## Ballooning hook
//!
//! When `BuddyAllocator::alloc_pages` returns `None`, the actor emits a
//! `PMResponse::OutOfMemory { channel_id }`.  The future ballooning subsystem
//! will intercept this before it reaches the requester and negotiate memory
//! transfer from a neighbouring actor.

use alloc::vec::Vec;
use core::sync::atomic::{AtomicU64, AtomicUsize, Ordering};

use crate::cap::{CapObject, CapRights, Capability, Relation};
use crate::mem::address::PhysAddr;
use crate::mem::buddy::BuddyAllocator;
use crate::mem::pmm::PAGE_SIZE;

extern crate alloc;

// Message packing constants

const QUEUE_SIZE: usize = 1024; // power-of-two
const QUEUE_MASK: usize = QUEUE_SIZE - 1;

/// Bit layout of a packed PMRequest u64.
mod packing {
    pub const OPCODE_SHIFT:  u32 = 56;
    pub const CHAN_SHIFT:    u32 = 40;
    pub const ARG1_SHIFT:   u32 = 20;
    // ARG2 sits in bits 19:0 (no shift).

    pub const CHAN_MASK:  u64 = 0x0000_FFFF;   // 16 bits
    pub const ARG_MASK:  u64 = 0x000F_FFFF;   // 20 bits
}

// Request type
/// Asynchronous request submitted to a `PMActor` inbox.
///
/// `channel_id` identifies where the response should be routed.
/// Use `channel_id = 0` if you do not need a reply.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PMRequest {
    /// Allocate `size_pages` pages; tag the resulting capability with `token_sig`.
    Allocate {
        size_pages: usize,
        token_sig:  u32,
        channel_id: u16,
    },
    /// Free a previously allocated block.
    ///
    /// `local_frame_idx` is the **relative** page index returned inside the
    /// `Allocated` capability (i.e. `cap.object.phys - managed_range.start`).
    /// `order` is the buddy order used at allocation time (stored in the cap).
    Free {
        local_frame_idx: usize,
        order:           usize,
    },
    /// Carve a fixed sub-region without consulting the buddy allocator.
    /// Used for static layout decisions made at boot (e.g. framebuffer alias).
    Carve {
        offset_pages: usize,
        size_pages:   usize,
        channel_id:   u16,
    },
    None,
}

impl PMRequest {
    /// Encode the request into one `u64` atom for the lock-free buffer.
    ///
    /// Opcode 0 (None) is reserved as the "slot empty" sentinel; `None` returns
    /// 0 intentionally and the sender rejects it before pushing.
    #[inline]
    pub fn pack(self) -> u64 {
        use packing::*;
        match self {
            Self::Allocate { size_pages, token_sig, channel_id } => {
                (1u64 << OPCODE_SHIFT)
                    | ((channel_id as u64 & CHAN_MASK) << CHAN_SHIFT)
                    | ((size_pages  as u64 & ARG_MASK) << ARG1_SHIFT)
                    | (token_sig as u64 & ARG_MASK)
            }
            Self::Free { local_frame_idx, order } => {
                (2u64 << OPCODE_SHIFT)
                    // channel_id not needed — Free has no reply
                    | ((local_frame_idx as u64 & ARG_MASK) << ARG1_SHIFT)
                    | (order as u64 & ARG_MASK)
            }
            Self::Carve { offset_pages, size_pages, channel_id } => {
                (3u64 << OPCODE_SHIFT)
                    | ((channel_id  as u64 & CHAN_MASK) << CHAN_SHIFT)
                    | ((offset_pages as u64 & ARG_MASK) << ARG1_SHIFT)
                    | (size_pages   as u64 & ARG_MASK)
            }
            Self::None => 0,
        }
    }

    /// Decode a packed `u64` back into a `PMRequest`.
    #[inline]
    pub fn unpack(val: u64) -> Self {
        use packing::*;
        let opcode     = val >> OPCODE_SHIFT;
        let channel_id = ((val >> CHAN_SHIFT) & CHAN_MASK) as u16;
        let arg1       = ((val >> ARG1_SHIFT) & ARG_MASK) as usize;
        let arg2       = (val & ARG_MASK) as usize;

        match opcode {
            1 => Self::Allocate {
                size_pages: arg1,
                token_sig:  arg2 as u32,
                channel_id,
            },
            2 => Self::Free {
                local_frame_idx: arg1,
                order:           arg2,
            },
            3 => Self::Carve {
                offset_pages: arg1,
                size_pages:   arg2,
                channel_id,
            },
            _ => Self::None,
        }
    }
}

// Response type

#[derive(Debug, Clone, Copy)]
pub struct PMResponse {
    pub channel_id: u16,
    pub result: PMResult,
}

/// Outcome of a single PM operation.
#[derive(Debug, Clone, Copy)]
pub enum PMResult {
    Allocated { cap: Capability, order: usize },
    OutOfMemory { size_pages: usize },
    Carved { cap: Capability },
    Freed { pages_returned: usize },
}

// Lock-free MPSC inbox
/// MPSC ring buffer for PMRequest values.
///
/// Producers (any core, any context) call `send`; the owning PMActor calls
/// `pop` — exactly one consumer, no locking on the read side.
pub struct PMActorQueue {
    buffer: [AtomicU64; QUEUE_SIZE],
    _pad0: [u8; 64],
    head: AtomicUsize,
    _pad1: [u8; 64],
    tail: AtomicUsize,
}

impl PMActorQueue {
    pub const fn new() -> Self {
        Self {
            #[allow(clippy::declare_interior_mutable_const)]
            buffer: [const { AtomicU64::new(0) }; QUEUE_SIZE],
            _pad0: [0u8; 64],
            head:  AtomicUsize::new(0),
            _pad1: [0u8; 64],
            tail:  AtomicUsize::new(0),
        }
    }

    /// Push a request into the inbox (Multi-Producer path).
    ///
    /// Returns `Err` only if the ring buffer is full — which should trigger
    /// backpressure or a ballooning request in production.
    pub fn send(&self, req: PMRequest) -> Result<(), &'static str> {
        let packed = req.pack();
        if packed == 0 {
            return Err("PMActorQueue::send: PMRequest::None is not sendable");
        }

        let mut tail = self.tail.load(Ordering::Relaxed);
        loop {
            let head = self.head.load(Ordering::Acquire);
            if tail.wrapping_sub(head) >= QUEUE_SIZE {
                return Err("PMActorQueue overflow — apply backpressure or balloon");
            }

            match self.tail.compare_exchange_weak(
                tail,
                tail.wrapping_add(1),
                Ordering::AcqRel,
                Ordering::Relaxed,
            ) {
                Ok(_) => {
                    self.buffer[tail & QUEUE_MASK].store(packed, Ordering::Release);
                    return Ok(());
                }
                Err(actual) => tail = actual,
            }
        }
    }

    /// Pop one request (Single-Consumer path — actor's own context only).
    pub fn pop(&self) -> Option<PMRequest> {
        let head = self.head.load(Ordering::Relaxed);
        let tail = self.tail.load(Ordering::Acquire);
        if head == tail {
            return None;
        }
        let packed = self.buffer[head & QUEUE_MASK].load(Ordering::Acquire);
        self.head.store(head.wrapping_add(1), Ordering::Release);
        Some(PMRequest::unpack(packed))
    }
}

// PM Actor
/// An autonomous physical-memory actor.
///
/// Owns a `BuddyAllocator` over its capital range and a lock-free MPSC inbox.
/// The actor must be driven externally: when the scheduler grants it CPU time,
/// it calls `process_messages()` to drain the inbox and produce `PMResponse`s.
///
/// # Single-consumer guarantee
/// Only `process_messages()` ever mutates `buddy` and `queue.head`.
/// This is enforced by taking `&mut self` on `process_messages`.
pub struct PMActor {
    pub actor_id: u64,
    pub root_untyped: Capability,
    pub managed_range: (PhysAddr, PhysAddr),
    queue: PMActorQueue,
    buddy: BuddyAllocator,
}

impl PMActor {
    /// Construct a new actor managing `size_bytes` bytes starting at `start`.
    ///
    /// `size_bytes` must be a multiple of `PAGE_SIZE`; non-multiples are
    /// rounded down silently.
    pub fn new(
        actor_id:  u64,
        root_cap:  Capability,
        start:     PhysAddr,
        size_bytes: u64,
    ) -> Self {
        let total_pages = (size_bytes / PAGE_SIZE) as usize;
        Self {
            actor_id,
            root_untyped: root_cap,
            managed_range: (start, PhysAddr(start.0 + total_pages as u64 * PAGE_SIZE)),
            queue: PMActorQueue::new(),
            buddy: BuddyAllocator::new(total_pages, start.0),
        }
    }

    //Producer API (callable from any context)

    /// Submit a request to this actor's inbox.
    ///
    /// Thread-safe (MPSC producer side).  Returns `Err` if the inbox is full.
    #[inline]
    pub fn submit_request(&self, req: PMRequest) -> Result<(), &'static str> {
        self.queue.send(req)
    }

    //Consumer API (actor's own scheduled context)

    /// Drain the inbox and execute all pending requests.
    ///
    /// Returns a `Vec` of responses that the caller (scheduler) must route to
    /// the appropriate process or actor mailbox based on `channel_id`.
    ///
    /// Responses with `channel_id == 0` can be discarded.
    ///
    /// # No locks held
    /// `buddy` and `queue.head` are mutated — no other thread touches them.
    /// The only shared state is the queue's `tail`, which is written by producers
    /// via `AtomicUsize::compare_exchange_weak`, never by this path.
    /// Maximum requests to process in a single `process_messages` call.
    /// Prevents kernel starvation when the inbox is deep.
    const MAX_MESSAGES_PER_CALL: usize = 64;

    pub fn process_messages(&mut self) -> Vec<PMResponse> {
        let mut responses = Vec::new();
        let mut remaining = Self::MAX_MESSAGES_PER_CALL;

        while let Some(req) = self.queue.pop() {
            remaining -= 1;
            let resp = match req {
                PMRequest::Allocate { size_pages, token_sig, channel_id } => {
                    self.handle_allocate(size_pages, token_sig, channel_id)
                }
                PMRequest::Free { local_frame_idx, order } => {
                    self.handle_free(local_frame_idx, order)
                }
                PMRequest::Carve { offset_pages, size_pages, channel_id } => {
                    self.handle_carve(offset_pages, size_pages, channel_id)
                }
                PMRequest::None => continue,
            };

            if let Some(r) = resp {
                responses.push(r);
            }

            if remaining == 0 {
                break;
            }
        }

        responses
    }

    //Introspection

    /// Free pages remaining in this actor's buddy pool.
    #[inline]
    pub fn free_pages(&self) -> usize {
        self.buddy.free_pages()
    }

    /// Total pages this actor was initialised with.
    #[inline]
    pub fn total_pages(&self) -> usize {
        self.buddy.total_pages()
    }
}

//Private handler implementations

impl PMActor {
    fn handle_allocate(
        &mut self,
        size_pages: usize,
        token_sig:  u32,
        channel_id: u16,
    ) -> Option<PMResponse> {
        match self.buddy.alloc_pages(size_pages) {
            Some((rel_idx, order)) => {
                // Translate relative index → absolute physical address.
                let phys = PhysAddr(self.managed_range.0.0 + rel_idx as u64 * PAGE_SIZE);

                // Build a strong child capability.
                // token_sig is widened and XOR'd with actor_id for uniqueness.
                let cap = Capability {
                    object: CapObject::Memory {
                        phys,
                        // Report actual (rounded-up) page count so the caller
                        // passes the correct `order` back to Free.
                        size_pages: 1 << order,
                    },
                    rights:     CapRights::READ | CapRights::WRITE | CapRights::GRANT,
                    relation:   Relation::Strong,
                    token_sig:  (token_sig as u64) ^ self.actor_id ^ (rel_idx as u64),
                };

                Some(PMResponse {
                    channel_id,
                    result: PMResult::Allocated { cap, order },
                })
            }
            None => {
                // TODO: emit a ballooning request to a neighbour actor before
                // propagating OutOfMemory upward.
                Some(PMResponse {
                    channel_id,
                    result: PMResult::OutOfMemory { size_pages },
                })
            }
        }
    }

    fn handle_free(
        &mut self,
        local_frame_idx: usize,
        order: usize,
    ) -> Option<PMResponse> {
        // Validate bounds before touching the buddy tree.
        let block_end = local_frame_idx.saturating_add(1usize << order);
        if block_end > self.buddy.total_pages() {
            // Corrupted or spoofed index — ignore silently in release, panic in debug.
            debug_assert!(
                false,
                "PMActor {}: Free out of range (idx={}, order={}, total={})",
                self.actor_id,
                local_frame_idx,
                order,
                self.buddy.total_pages()
            );
            return None;
        }

        self.buddy.free(local_frame_idx, order);

        None
    }

    fn handle_carve(
        &mut self,
        offset_pages: usize,
        size_pages:   usize,
        channel_id:   u16,
    ) -> Option<PMResponse> {
        let range_pages = self.buddy.total_pages();
        if offset_pages >= range_pages
            || size_pages == 0
            || offset_pages.saturating_add(size_pages) > range_pages
        {
            return Some(PMResponse {
                channel_id,
                result: PMResult::OutOfMemory { size_pages },
            });
        }

        let phys = PhysAddr(self.managed_range.0.0 + offset_pages as u64 * PAGE_SIZE);
        let cap  = Capability {
            object: CapObject::Memory { phys, size_pages },
            rights:    CapRights::READ | CapRights::WRITE | CapRights::GRANT,
            relation:  Relation::Strong,
            token_sig: self.actor_id ^ offset_pages as u64,
        };

        Some(PMResponse {
            channel_id,
            result: PMResult::Carved { cap },
        })
    }
}