//! Two-slot, task-owned kernel handle table. A number is never an address;
//! the caller may use it only while the Scheduler names that exact TaskId.
use crate::task::TaskId;

pub const CAPACITY: usize = 2;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Resource {
    Serial,
    Clock,
    Channel(u64),
}

impl Resource {
    pub fn number(self) -> u64 {
        match self {
            Self::Serial => 1,
            Self::Clock => 2,
            Self::Channel(_) => 3,
        }
    }

    pub fn from_number(number: u64) -> Option<Self> {
        match number {
            1 => Some(Self::Serial),
            2 => Some(Self::Clock),
            _ => None,
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
    Full,
    IdExhausted,
    InvalidHandle,
    SameOwner,
}

#[derive(Clone, Copy)]
struct Entry {
    id: u64,
    owner: TaskId,
    resource: Resource,
}

pub struct Table {
    entries: [Option<Entry>; CAPACITY],
    next_id: u64,
}

impl Table {
    pub const fn new() -> Self {
        Self {
            entries: [None; CAPACITY],
            next_id: 1,
        }
    }

    pub fn len(&self) -> usize {
        self.entries.iter().filter(|entry| entry.is_some()).count()
    }

    pub fn open(&mut self, owner: TaskId, resource: Resource) -> Result<u64, Error> {
        let vacant = self
            .entries
            .iter_mut()
            .find(|entry| entry.is_none())
            .ok_or(Error::Full)?;
        let following = self.next_id.checked_add(1).ok_or(Error::IdExhausted)?;
        let id = self.next_id;
        *vacant = Some(Entry {
            id,
            owner,
            resource,
        });
        self.next_id = following;
        Ok(id)
    }

    pub fn query(&self, owner: TaskId, id: u64) -> Result<Resource, Error> {
        self.entries
            .iter()
            .flatten()
            .find(|entry| entry.id == id && entry.owner == owner)
            .map(|entry| entry.resource)
            .ok_or(Error::InvalidHandle)
    }

    pub fn close(&mut self, owner: TaskId, id: u64) -> Result<(), Error> {
        let entry = self
            .entries
            .iter_mut()
            .find(|entry| entry.is_some_and(|entry| entry.id == id && entry.owner == owner))
            .ok_or(Error::InvalidHandle)?;
        *entry = None;
        Ok(())
    }

    /// Move exactly one existing capability to a distinct task. Reissue the
    /// opaque number so a stale sender cannot regain access after a later
    /// transfer back. The caller must bind both task IDs to live scheduler
    /// entries; this table cannot decide whether a recipient is running.
    pub fn transfer(&mut self, owner: TaskId, recipient: TaskId, id: u64) -> Result<u64, Error> {
        let entry = self
            .entries
            .iter_mut()
            .flatten()
            .find(|entry| entry.id == id && entry.owner == owner)
            .ok_or(Error::InvalidHandle)?;
        if owner == recipient {
            return Err(Error::SameOwner);
        }
        let following = self.next_id.checked_add(1).ok_or(Error::IdExhausted)?;
        let issued = self.next_id;
        entry.id = issued;
        entry.owner = recipient;
        self.next_id = following;
        Ok(issued)
    }

    /// Revoke every handle naming an object removed by the task-exit hook,
    /// including a capability that was previously transferred to a peer.
    pub fn revoke_channel(&mut self, channel_id: u64) -> usize {
        let mut revoked = 0;
        for entry in &mut self.entries {
            if entry.is_some_and(|item| item.resource == Resource::Channel(channel_id)) {
                *entry = None;
                revoked += 1;
            }
        }
        revoked
    }

    /// Close resources still owned by this exact task generation. Opaque IDs
    /// remain monotonic and never become valid again after cleanup.
    pub fn close_all(&mut self, owner: TaskId) -> usize {
        let mut closed = 0;
        for entry in &mut self.entries {
            if entry.is_some_and(|item| item.owner == owner) {
                *entry = None;
                closed += 1;
            }
        }
        closed
    }
}

impl Default for Table {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    extern crate std;

    use super::*;
    const A: TaskId = TaskId {
        slot: 0,
        generation: 1,
    };
    const B: TaskId = TaskId {
        slot: 1,
        generation: 2,
    };

    #[test]
    fn foreign_owner_cannot_query_or_close_live_handle() {
        let mut table = Table::new();
        let a = table.open(A, Resource::Serial).unwrap();
        assert_eq!(a, 1);
        assert_eq!(table.query(B, a), Err(Error::InvalidHandle));
        assert_eq!(table.close(B, a), Err(Error::InvalidHandle));
        assert_eq!(table.len(), 1);
        assert_eq!(table.query(A, a), Ok(Resource::Serial));
        assert_eq!(table.close(A, a), Ok(()));
        assert_eq!(table.query(A, a), Err(Error::InvalidHandle));
    }

    #[test]
    fn capacity_and_stale_number_are_fail_closed() {
        let mut table = Table::new();
        let first = table.open(A, Resource::Serial).unwrap();
        let second = table.open(B, Resource::Clock).unwrap();
        assert_eq!(table.open(A, Resource::Clock), Err(Error::Full));
        assert_eq!(table.len(), CAPACITY);
        assert_eq!(table.close(A, first), Ok(()));
        let third = table.open(A, Resource::Clock).unwrap();
        assert_eq!((first, second, third), (1, 2, 3));
        assert_eq!(table.query(A, first), Err(Error::InvalidHandle));
        assert_eq!(table.query(A, third), Ok(Resource::Clock));
    }

    #[test]
    fn old_generation_and_invalid_resource_are_rejected() {
        let mut table = Table::new();
        let first = table.open(A, Resource::Serial).unwrap();
        let recycled = TaskId {
            slot: A.slot,
            generation: A.generation + 2,
        };
        assert_eq!(table.query(recycled, first), Err(Error::InvalidHandle));
        assert_eq!(table.close(recycled, first), Err(Error::InvalidHandle));
        assert_eq!(Resource::from_number(0), None);
        assert_eq!(Resource::from_number(3), None);
        assert_eq!(Resource::from_number(1), Some(Resource::Serial));
        assert_eq!(Resource::from_number(2), Some(Resource::Clock));
        assert_eq!(table.query(A, 0), Err(Error::InvalidHandle));
    }

    #[test]
    fn id_overflow_leaves_vacant_slot_untouched() {
        let mut table = Table::new();
        table.next_id = u64::MAX;
        assert_eq!(table.open(A, Resource::Serial), Err(Error::IdExhausted));
        assert_eq!(table.len(), 0);
        assert_eq!(table.next_id, u64::MAX);
    }

    #[test]
    fn k314_transfer_moves_one_capability_and_invalidates_old_number() {
        let mut table = Table::new();
        let first = table.open(A, Resource::Channel(7)).unwrap();
        assert_eq!(table.transfer(A, B, first), Ok(2));
        assert_eq!(table.len(), 1);
        assert_eq!(table.query(A, first), Err(Error::InvalidHandle));
        assert_eq!(table.query(B, first), Err(Error::InvalidHandle));
        assert_eq!(table.query(A, 2), Err(Error::InvalidHandle));
        assert_eq!(table.query(B, 2), Ok(Resource::Channel(7)));
        assert_eq!(table.transfer(B, A, 2), Ok(3));
        assert_eq!(table.query(B, 2), Err(Error::InvalidHandle));
        assert_eq!(table.query(A, 3), Ok(Resource::Channel(7)));
        assert_eq!(table.len(), 1);
        assert_eq!(table.close(A, 3), Ok(()));
        assert_eq!(table.len(), 0);
    }

    #[test]
    fn k314_foreign_self_stale_generation_and_overflow_never_move_capability() {
        let mut table = Table::new();
        let first = table.open(A, Resource::Clock).unwrap();
        let stale_a = TaskId {
            slot: A.slot,
            generation: A.generation + 1,
        };
        assert_eq!(table.transfer(B, A, first), Err(Error::InvalidHandle));
        assert_eq!(table.transfer(stale_a, B, first), Err(Error::InvalidHandle));
        assert_eq!(table.transfer(A, A, first), Err(Error::SameOwner));
        assert_eq!(table.len(), 1);
        assert_eq!(table.next_id, 2);
        assert_eq!(table.query(A, first), Ok(Resource::Clock));
        table.next_id = u64::MAX;
        assert_eq!(table.transfer(A, B, first), Err(Error::IdExhausted));
        assert_eq!(table.query(A, first), Ok(Resource::Clock));
        assert_eq!(table.query(B, first), Err(Error::InvalidHandle));
        assert_eq!(table.next_id, u64::MAX);
        assert_eq!(table.len(), 1);
    }

    #[test]
    fn k451_transfer_matrix_never_duplicates_or_resurrects_capabilities() {
        let mut accepted = 0usize;
        let mut invalid = 0usize;
        let mut same_owner = 0usize;
        let mut exhausted = 0usize;
        for epoch in 0..64u64 {
            let a = TaskId {
                slot: 0,
                generation: epoch + 1,
            };
            let b = TaskId {
                slot: 1,
                generation: epoch + 2,
            };
            let stale_a = TaskId {
                slot: 0,
                generation: epoch,
            };
            let stale_b = TaskId {
                slot: 1,
                generation: epoch + 1,
            };
            let c = TaskId {
                slot: 2,
                generation: epoch + 3,
            };
            let identities = [a, b, stale_a, stale_b, c];
            let numbers = [0, 1, 2, 3, epoch + 100, u64::MAX];
            for occupied in 0..=2 {
                for &owner in &identities {
                    for &recipient in &identities {
                        for &number in &numbers {
                            for overflow in [false, true] {
                                let mut table = Table::new();
                                if occupied >= 1 {
                                    assert_eq!(table.open(a, Resource::Channel(epoch)), Ok(1));
                                }
                                if occupied == 2 {
                                    assert_eq!(table.open(b, Resource::Clock), Ok(2));
                                }
                                if overflow {
                                    table.next_id = u64::MAX;
                                }
                                let before: [Option<(u64, TaskId, Resource)>; CAPACITY] =
                                    table.entries.map(|entry| {
                                        entry.map(|item| (item.id, item.owner, item.resource))
                                    });
                                let next_before = table.next_id;
                                let found = before.iter().enumerate().find_map(|(index, entry)| {
                                    entry
                                        .filter(|item| item.0 == number && item.1 == owner)
                                        .map(|item| (index, item.2))
                                });
                                let expected = match found {
                                    None => Err(Error::InvalidHandle),
                                    Some(_) if owner == recipient => Err(Error::SameOwner),
                                    Some(_) if next_before == u64::MAX => Err(Error::IdExhausted),
                                    Some(_) => Ok(next_before),
                                };
                                let actual = table.transfer(owner, recipient, number);
                                assert_eq!(
                                    actual, expected,
                                    "epoch={epoch} occupied={occupied} id={number} overflow={overflow}"
                                );
                                let after: [Option<(u64, TaskId, Resource)>; CAPACITY] =
                                    table.entries.map(|entry| {
                                        entry.map(|item| (item.id, item.owner, item.resource))
                                    });
                                match actual {
                                    Err(error) => {
                                        match error {
                                            Error::InvalidHandle => invalid += 1,
                                            Error::SameOwner => same_owner += 1,
                                            Error::IdExhausted => exhausted += 1,
                                            Error::Full => {
                                                panic!("transfer cannot allocate a new slot")
                                            }
                                        }
                                        assert_eq!(
                                            after, before,
                                            "failed transfer changed a capability"
                                        );
                                        assert_eq!(table.next_id, next_before);
                                    }
                                    Ok(issued) => {
                                        accepted += 1;
                                        let (slot, resource) = found.unwrap();
                                        assert_eq!(
                                            after[slot],
                                            Some((issued, recipient, resource))
                                        );
                                        assert_eq!(table.next_id, issued + 1);
                                        assert_eq!(table.len(), occupied);
                                        for other in 0..CAPACITY {
                                            if other != slot {
                                                assert_eq!(
                                                    after[other], before[other],
                                                    "transfer duplicated or changed the other slot"
                                                );
                                            }
                                        }
                                        for &identity in &identities {
                                            assert_eq!(
                                                table.query(identity, number),
                                                Err(Error::InvalidHandle),
                                                "old handle resurrected"
                                            );
                                            assert_eq!(
                                                table.query(identity, issued),
                                                if identity == recipient {
                                                    Ok(resource)
                                                } else {
                                                    Err(Error::InvalidHandle)
                                                }
                                            );
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
        assert_eq!(accepted + invalid + same_owner + exhausted, 57_600);
        let mut table = Table::new();
        let mut owner = A;
        let mut number = table.open(A, Resource::Channel(7)).unwrap();
        for turn in 0..512u64 {
            let recipient = if owner == A { B } else { A };
            let issued = table.transfer(owner, recipient, number).unwrap();
            assert_eq!(issued, turn + 2);
            assert_eq!(table.len(), 1);
            assert_eq!(table.query(owner, number), Err(Error::InvalidHandle));
            assert_eq!(table.query(recipient, number), Err(Error::InvalidHandle));
            assert_eq!(table.query(recipient, issued), Ok(Resource::Channel(7)));
            for stale in 1..issued {
                assert_eq!(table.query(A, stale), Err(Error::InvalidHandle));
                assert_eq!(table.query(B, stale), Err(Error::InvalidHandle));
            }
            owner = recipient;
            number = issued;
        }
        std::println!(
            "K451 TRANSFER-MATRIX: CASES=57600 OK={accepted} INVALID={invalid} SAME={same_owner} EXHAUSTED={exhausted} STALE-EPOCHS=512"
        );
    }

    #[test]
    fn k315_exit_releases_exact_owner_and_revokes_transferred_channel() {
        let mut table = Table::new();
        let channel = table.open(A, Resource::Channel(7)).unwrap();
        assert_eq!(table.transfer(A, B, channel), Ok(2));
        let serial = table.open(A, Resource::Serial).unwrap();
        assert_eq!(serial, 3);
        assert_eq!(table.revoke_channel(8), 0);
        assert_eq!(table.revoke_channel(7), 1);
        assert_eq!(table.query(B, 2), Err(Error::InvalidHandle));
        let stale_a = TaskId {
            slot: A.slot,
            generation: A.generation + 1,
        };
        assert_eq!(table.close_all(stale_a), 0);
        assert_eq!(table.query(A, serial), Ok(Resource::Serial));
        assert_eq!(table.close_all(A), 1);
        assert_eq!(table.len(), 0);
        assert_eq!(table.open(B, Resource::Clock), Ok(4));
        assert_eq!(table.query(A, channel), Err(Error::InvalidHandle));
    }
}