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

pub const MIN_CNODE_BITS: u8 = 1;
pub const MAX_CNODE_BITS: u8 = 15;
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 <= u16::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 extract_guard(address: u64, depth: u8, guard_bits: u8) -> u64 {
    match guard_bits {
        0 => 0,
        64 => address,
        _ => {
            let shift = depth - guard_bits;
            let mask = (1u64 << guard_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_phys: ObjPhys,
        cnode_gen: Generation,
        index: usize,
    ) -> Result<CapSlot, KernelError>;

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

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

#[allow(clippy::too_many_arguments)]
fn resolve_inner<S: CNodeStore>(
    store: &S,
    cnode_phys: ObjPhys,
    cnode_gen: Generation,
    address: u64,
    depth: u8,
    guard_value: u64,
    guard_bits: u8,
    recursion: u8,
) -> Result<CapSlot, KernelError> {
    if recursion >= MAX_RESOLVE_DEPTH {
        return Err(KernelError::InvalidSlot);
    }

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

    let extracted = extract_guard(address, depth, guard_bits);
    if extracted != guard_value {
        return Err(KernelError::GuardMismatch);
    }

    let post_guard_depth = depth - guard_bits;

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

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

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

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

#[allow(clippy::too_many_arguments)]
pub fn resolve_and_validate<S: CNodeStore>(
    store: &S,
    cnode_phys: ObjPhys,
    cnode_gen: Generation,
    address: u64,
    depth: u8,
    guard_value: u64,
    guard_bits: u8,
    expected_tag: ObjectTag,
    required_rights: crate::cap_table::Rights,
) -> Result<CapRef, KernelError> {
    match resolve(
        store,
        cnode_phys,
        cnode_gen,
        address,
        depth,
        guard_value,
        guard_bits,
    )? {
        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_phys: ObjPhys,
    cnode_gen: Generation,
    address: u64,
    depth: u8,
    guard_value: u64,
    guard_bits: u8,
) -> Result<CapRef, KernelError> {
    match resolve(
        store,
        cnode_phys,
        cnode_gen,
        address,
        depth,
        guard_value,
        guard_bits,
    )? {
        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_phys: ObjPhys,
        expected_gen: Generation,
    }

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

    impl CNodeStore for SingleCNode {
        fn read_slot(
            &self,
            cnode_phys: ObjPhys,
            cnode_gen: Generation,
            index: usize,
        ) -> Result<CapSlot, KernelError> {
            if cnode_phys != self.expected_phys || 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_phys: ObjPhys, cnode_gen: Generation) -> Result<u8, KernelError> {
            if cnode_phys != self.expected_phys || 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 phys = ObjPhys::new(0x1000);
        let generation = Generation::new(0);
        let store = SingleCNode::new(phys, generation);

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

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

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

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

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

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

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

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

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

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

    #[kani::proof]
    fn extract_guard_bounded() {
        let address: u64 = kani::any();
        let depth: u8 = kani::any();
        let guard_bits: u8 = kani::any();

        kani::assume(guard_bits <= 64);
        kani::assume(depth >= guard_bits);
        kani::assume(depth <= 64);

        let val = extract_guard(address, depth, guard_bits);
        match guard_bits {
            0 => assert!(val == 0),
            64 => assert!(val == address),
            _ => assert!(val < (1u64 << guard_bits)),
        }
    }

    #[kani::proof]
    fn guard_mismatch_detected() {
        let phys = ObjPhys::new(0x1000);
        let generation = Generation::new(0);
        let store = SingleCNode::new(phys, generation);

        let result = resolve(&store, phys, generation, 0, 8, 1, 4);
        assert!(matches!(result, Err(KernelError::GuardMismatch)));
    }

    #[kani::proof]
    fn zero_guard_identity() {
        let phys = ObjPhys::new(0x1000);
        let generation = Generation::new(0);
        let store = SingleCNode::new(phys, generation);

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

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

    #[kani::proof]
    fn guard_plus_size_partitioning() {
        let guard_bits: u8 = kani::any();
        let size_bits: u8 = kani::any();

        kani::assume(guard_bits <= 8);
        kani::assume(size_bits >= 1 && size_bits <= 8);
        let total = guard_bits + size_bits;
        kani::assume(total <= 16);

        let guard_val: u64 = kani::any();
        let index_val: u64 = kani::any();
        kani::assume(guard_bits == 0 || guard_val < (1u64 << guard_bits));
        kani::assume(index_val < (1u64 << size_bits));

        let address = (guard_val << size_bits) | index_val;

        let got_guard = extract_guard(address, total, guard_bits);
        let got_index = extract_index(address, total - guard_bits, size_bits);

        assert!(got_guard == guard_val);
        assert!(got_index == index_val);
    }

    #[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))),
        }
    }
}