use crate::cap_table::{CapRef, CapSlot};
use crate::error::KernelError;
use crate::object_tag::ObjectTag;
use crate::types::{Generation, ObjectId};

pub const MIN_CNODE_BITS: u8 = 1;
pub const MAX_CNODE_BITS: u8 = 12;
pub const MAX_RESOLVE_DEPTH: u8 = 16;
pub const CAPSLOT_SIZE: usize = core::mem::size_of::<CapSlot>();
pub const PAGE_SIZE: usize = 4096;

const _: () = {
    let max_slots = 1usize << MAX_CNODE_BITS;
    let max_bytes = max_slots * CAPSLOT_SIZE;
    let max_frames = max_bytes.div_ceil(PAGE_SIZE);
    assert!(max_frames <= u8::MAX as usize);
};

pub fn frames_for_cnode(size_bits: u8) -> u8 {
    let slots = 1usize << size_bits;
    let bytes = slots * CAPSLOT_SIZE;
    bytes.div_ceil(PAGE_SIZE) as u8
}

pub fn extract_index(address: u64, remaining_depth: u8, size_bits: u8) -> u64 {
    let shift = remaining_depth - size_bits;
    let mask = (1u64 << size_bits) - 1;
    (address >> shift) & mask
}

pub fn validate_size_bits(size_bits: u8) -> Result<(), KernelError> {
    match (MIN_CNODE_BITS..=MAX_CNODE_BITS).contains(&size_bits) {
        true => Ok(()),
        false => Err(KernelError::InvalidParameter),
    }
}

pub trait CNodeStore {
    fn read_slot(
        &self,
        cnode_id: ObjectId,
        cnode_gen: Generation,
        index: usize,
    ) -> Result<CapSlot, KernelError>;

    fn size_bits(
        &self,
        cnode_id: ObjectId,
        cnode_gen: Generation,
    ) -> Result<u8, KernelError>;
}

pub fn resolve<S: CNodeStore>(
    store: &S,
    cnode_id: ObjectId,
    cnode_gen: Generation,
    address: u64,
    depth: u8,
) -> Result<CapSlot, KernelError> {
    resolve_inner(store, cnode_id, cnode_gen, address, depth, 0)
}

fn resolve_inner<S: CNodeStore>(
    store: &S,
    cnode_id: ObjectId,
    cnode_gen: Generation,
    address: u64,
    depth: u8,
    recursion: u8,
) -> Result<CapSlot, KernelError> {
    if recursion >= MAX_RESOLVE_DEPTH {
        return Err(KernelError::InvalidSlot);
    }

    let size_bits = store.size_bits(cnode_id, cnode_gen)?;

    if depth < size_bits {
        return Err(KernelError::InvalidSlot);
    }

    let index = extract_index(address, depth, size_bits) as usize;
    let remaining = depth - size_bits;
    let slot = store.read_slot(cnode_id, cnode_gen, index)?;

    match remaining {
        0 => Ok(slot),
        _ => match slot {
            CapSlot::Active(cap) if cap.tag() == ObjectTag::CNode => resolve_inner(
                store,
                cap.object_id(),
                cap.generation(),
                address,
                remaining,
                recursion + 1,
            ),
            _ => Err(KernelError::InvalidSlot),
        },
    }
}

pub fn resolve_and_validate<S: CNodeStore>(
    store: &S,
    cnode_id: ObjectId,
    cnode_gen: Generation,
    address: u64,
    depth: u8,
    expected_tag: ObjectTag,
    required_rights: crate::cap_table::Rights,
) -> Result<CapRef, KernelError> {
    match resolve(store, cnode_id, cnode_gen, address, depth)? {
        CapSlot::Empty => Err(KernelError::SlotEmpty),
        CapSlot::Active(cap) => {
            if cap.tag() != expected_tag {
                return Err(KernelError::InvalidType);
            }
            if !cap.rights().contains(required_rights) {
                return Err(KernelError::InsufficientRights);
            }
            Ok(cap)
        }
    }
}

pub fn resolve_and_read<S: CNodeStore>(
    store: &S,
    cnode_id: ObjectId,
    cnode_gen: Generation,
    address: u64,
    depth: u8,
) -> Result<CapRef, KernelError> {
    match resolve(store, cnode_id, cnode_gen, address, depth)? {
        CapSlot::Empty => Err(KernelError::SlotEmpty),
        CapSlot::Active(cap) => Ok(cap),
    }
}

#[cfg(kani)]
mod kani_proofs {
    use super::*;
    use crate::cap_table::Rights;

    struct SingleCNode {
        slots: [CapSlot; 16],
        size_bits: u8,
        expected_id: ObjectId,
        expected_gen: Generation,
    }

    impl SingleCNode {
        fn new(id: ObjectId, generation: Generation) -> Self {
            Self {
                slots: [CapSlot::Empty; 16],
                size_bits: 4,
                expected_id: id,
                expected_gen: generation,
            }
        }
    }

    impl CNodeStore for SingleCNode {
        fn read_slot(
            &self,
            cnode_id: ObjectId,
            cnode_gen: Generation,
            index: usize,
        ) -> Result<CapSlot, KernelError> {
            if cnode_id != self.expected_id || cnode_gen != self.expected_gen {
                return Err(KernelError::StaleGeneration);
            }
            match index < (1usize << self.size_bits) {
                true => Ok(self.slots[index]),
                false => Err(KernelError::InvalidSlot),
            }
        }

        fn size_bits(
            &self,
            cnode_id: ObjectId,
            cnode_gen: Generation,
        ) -> Result<u8, KernelError> {
            if cnode_id != self.expected_id || cnode_gen != self.expected_gen {
                return Err(KernelError::StaleGeneration);
            }
            Ok(self.size_bits)
        }
    }

    #[kani::proof]
    fn extract_index_bounded() {
        let address: u64 = kani::any();
        let size_bits: u8 = kani::any();
        let remaining_depth: u8 = kani::any();

        kani::assume(size_bits >= MIN_CNODE_BITS && size_bits <= MAX_CNODE_BITS);
        kani::assume(remaining_depth >= size_bits && remaining_depth <= 64);

        let idx = extract_index(address, remaining_depth, size_bits);
        assert!(idx < (1u64 << size_bits));
    }

    #[kani::proof]
    fn extract_index_shift_correctness() {
        let size_bits: u8 = kani::any();
        let remaining_depth: u8 = kani::any();
        let address: u64 = kani::any();

        kani::assume(size_bits >= 1 && size_bits <= 8);
        kani::assume(remaining_depth >= size_bits && remaining_depth <= 16);

        let idx = extract_index(address, remaining_depth, size_bits);
        let shift = remaining_depth - size_bits;
        let reconstructed = idx << shift;

        let mask = ((1u64 << size_bits) - 1) << shift;
        assert!((reconstructed ^ (address & mask)) == 0);
    }

    #[kani::proof]
    fn two_level_address_partitioning() {
        let outer_bits: u8 = kani::any();
        let inner_bits: u8 = kani::any();

        kani::assume(outer_bits >= 1 && outer_bits <= 6);
        kani::assume(inner_bits >= 1 && inner_bits <= 6);

        let total_depth = outer_bits + inner_bits;
        kani::assume(total_depth <= 12);

        let outer_idx: u64 = kani::any();
        let inner_idx: u64 = kani::any();
        kani::assume(outer_idx < (1u64 << outer_bits));
        kani::assume(inner_idx < (1u64 << inner_bits));

        let address = (outer_idx << inner_bits) | inner_idx;

        let got_outer = extract_index(address, total_depth, outer_bits);
        let remaining = total_depth - outer_bits;
        let got_inner = extract_index(address, remaining, inner_bits);

        assert!(got_outer == outer_idx);
        assert!(got_inner == inner_idx);
    }

    #[kani::proof]
    fn frames_for_cnode_monotonic() {
        let a: u8 = kani::any();
        let b: u8 = kani::any();
        kani::assume(a >= MIN_CNODE_BITS && a <= MAX_CNODE_BITS);
        kani::assume(b >= MIN_CNODE_BITS && b <= MAX_CNODE_BITS);
        kani::assume(a <= b);
        assert!(frames_for_cnode(a) <= frames_for_cnode(b));
    }

    #[kani::proof]
    fn resolve_empty_cnode_returns_empty() {
        let id = ObjectId::new(0);
        let generation = Generation::new(0);
        let store = SingleCNode::new(id, generation);

        let address: u8 = kani::any();
        kani::assume(address < 16);

        let result = resolve(&store, id, generation, address as u64, 4);
        assert!(matches!(result, Ok(CapSlot::Empty)));
    }

    #[kani::proof]
    fn resolve_finds_inserted_cap() {
        let id = ObjectId::new(0);
        let generation = Generation::new(0);
        let mut store = SingleCNode::new(id, generation);

        let slot_idx: u8 = kani::any();
        kani::assume(slot_idx < 16);

        let cap = CapRef::new(ObjectTag::Endpoint, ObjectId::new(1), Rights::ALL, Generation::new(0));
        store.slots[slot_idx as usize] = CapSlot::Active(cap);

        let result = resolve(&store, id, generation, slot_idx as u64, 4);
        match result {
            Ok(CapSlot::Active(found)) => {
                assert!(found.tag() == ObjectTag::Endpoint);
                assert!(found.object_id() == ObjectId::new(1));
            }
            _ => panic!("expected Active slot"),
        }
    }

    #[kani::proof]
    fn resolve_stale_generation_rejected() {
        let id = ObjectId::new(0);
        let generation = Generation::new(0);
        let store = SingleCNode::new(id, generation);

        let stale = Generation::new(1);
        let result = resolve(&store, id, stale, 0, 4);
        assert!(matches!(result, Err(KernelError::StaleGeneration)));
    }

    #[kani::proof]
    fn resolve_depth_less_than_size_bits_rejected() {
        let id = ObjectId::new(0);
        let generation = Generation::new(0);
        let store = SingleCNode::new(id, generation);

        let depth: u8 = kani::any();
        kani::assume(depth < 4);
        let result = resolve(&store, id, generation, 0, depth);
        assert!(matches!(result, Err(KernelError::InvalidSlot)));
    }

    #[kani::proof]
    fn validate_size_bits_range() {
        let s: u8 = kani::any();
        let result = validate_size_bits(s);
        match (MIN_CNODE_BITS..=MAX_CNODE_BITS).contains(&s) {
            true => assert!(result.is_ok()),
            false => assert!(matches!(result, Err(KernelError::InvalidParameter))),
        }
    }
}