use std::{collections::BTreeMap, ops::RangeInclusive};

//===========================================//

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AllocError {
    Failed,
}

pub type Addr = u8;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PinKind {
    DigitalIn,
    DigitalOut,
    PwmOut,
    AnalogIn,
}

#[derive(Debug)]
pub struct Board {
    capabilities: &'static [PinKind],
    filled: u8,
    capacity: u8,
}

impl Board {
    pub fn can_allocate(&self, kind: PinKind) -> bool {
        self.capabilities.contains(&kind)
    }

    fn allocate_one(&mut self) -> u8 {
        let index = self.filled;
        self.filled += 1;
        index
    }

    pub fn allocate(&mut self, kind: PinKind) -> Result<u8, AllocError> {
        if !self.can_allocate(kind) || self.filled >= self.capacity {
            return Err(AllocError::Failed);
        }
        Ok(self.allocate_one())
    }

    pub fn allocate_block(
        &mut self,
        kind: PinKind,
        count: u8,
    ) -> Result<RangeInclusive<u8>, AllocError> {
        if !self.can_allocate(kind) || count > self.capacity || self.capacity - self.filled < count
        {
            return Err(AllocError::Failed);
        }

        let start = self.filled;
        self.filled += count;
        Ok(start..=(self.filled - 1))
    }
}

//===========================================//

#[derive(Debug)]
pub struct PinBlock {
    pub addr: Addr,
    pub pins: RangeInclusive<u8>,
}

#[derive(Debug)]
pub struct Pin {
    pub addr: Addr,
    pub pin: u8,
}

//===========================================//

pub struct Allocator {
    boards: Vec<Board>,
    addresses: BTreeMap<Addr, usize>,
}

impl Allocator {
    pub fn new() -> Self {
        Self {
            boards: Vec::new(),
            addresses: BTreeMap::new(),
        }
    }

    pub fn register_board(
        &mut self,
        capacity: u8,
        capabilities: &'static [PinKind],
        range: RangeInclusive<Addr>,
    ) -> Result<Addr, AllocError> {
        let addr = range
            .into_iter()
            .find(|addr| !self.addresses.contains_key(addr))
            .ok_or(AllocError::Failed)?;

        let board = Board {
            capabilities,
            filled: 0,
            capacity,
        };

        self.boards.push(board);
        self.addresses.insert(addr, self.boards.len() - 1);
        Ok(addr)
    }

    pub fn allocate_pin(&mut self, kind: PinKind) -> Result<Pin, AllocError> {
        for (&addr, &idx) in self.addresses.iter() {
            if let Ok(pin) = self.boards[idx].allocate(kind) {
                return Ok(Pin { addr, pin });
            }
        }
        Err(AllocError::Failed)
    }

    pub fn allocate_pin_block(&mut self, kind: PinKind, count: u8) -> Result<PinBlock, AllocError> {
        for (&addr, &idx) in self.addresses.iter() {
            if let Ok(pins) = self.boards[idx].allocate_block(kind, count) {
                return Ok(PinBlock { addr, pins });
            }
        }
        Err(AllocError::Failed)
    }
}

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

//===========================================//
//================= TESTS ===================//
//===========================================//

#[cfg(test)]
mod tests {
    use super::*;
    use proptest::prelude::*;

    // ---------- helpers ----------

    fn basic_allocator(cap: u8, kinds: &'static [PinKind]) -> Allocator {
        let mut a = Allocator::new();
        a.register_board(cap, kinds, 0..=10).unwrap();
        a
    }

    // ---------- unit tests ----------

    #[test]
    fn basic_allocation_increments() {
        let mut a = basic_allocator(4, &[PinKind::DigitalOut]);

        let p0 = a.allocate_pin(PinKind::DigitalOut).unwrap();
        let p1 = a.allocate_pin(PinKind::DigitalOut).unwrap();

        assert_eq!(p0.pin, 0);
        assert_eq!(p1.pin, 1);
        assert_eq!(p0.addr, p1.addr);
    }

    #[test]
    fn unsupported_kind_fails() {
        let mut a = basic_allocator(4, &[PinKind::DigitalOut]);

        let err = a.allocate_pin(PinKind::AnalogIn).unwrap_err();
        assert_eq!(err, AllocError::Failed);
    }

    #[test]
    fn out_of_space_single() {
        let mut a = basic_allocator(2, &[PinKind::DigitalOut]);

        assert!(a.allocate_pin(PinKind::DigitalOut).is_ok());
        assert!(a.allocate_pin(PinKind::DigitalOut).is_ok());

        let err = a.allocate_pin(PinKind::DigitalOut).unwrap_err();
        assert_eq!(err, AllocError::Failed);
    }

    #[test]
    fn block_allocation_success_and_fail() {
        let mut a = basic_allocator(5, &[PinKind::PwmOut]);

        let block = a.allocate_pin_block(PinKind::PwmOut, 3).unwrap();
        assert_eq!(block.pins, 0..=2);

        let err = a.allocate_pin_block(PinKind::PwmOut, 3).unwrap_err();
        assert_eq!(err, AllocError::Failed);
    }

    #[test]
    fn block_too_large() {
        let mut a = basic_allocator(4, &[PinKind::PwmOut]);

        let err = a.allocate_pin_block(PinKind::PwmOut, 10).unwrap_err();
        assert_eq!(err, AllocError::Failed);
    }

    #[test]
    fn address_exhaustion() {
        let mut a = Allocator::new();

        assert!(a.register_board(2, &[PinKind::DigitalOut], 0..=1).is_ok());
        assert!(a.register_board(2, &[PinKind::DigitalOut], 0..=1).is_ok());

        let err = a
            .register_board(2, &[PinKind::DigitalOut], 0..=1)
            .unwrap_err();

        assert_eq!(err, AllocError::Failed);
    }

    #[test]
    fn multi_board_fallback() {
        let mut a = Allocator::new();

        a.register_board(1, &[PinKind::DigitalOut], 0..=10).unwrap();
        a.register_board(3, &[PinKind::DigitalOut], 0..=10).unwrap();

        let p0 = a.allocate_pin(PinKind::DigitalOut).unwrap();
        let p1 = a.allocate_pin(PinKind::DigitalOut).unwrap();

        assert_ne!(p0.addr, p1.addr);
    }

    // ---------- property tests ----------

    proptest! {

        #[test]
        fn never_over_allocates(cap in 1u8..20) {
            let mut a = basic_allocator(cap, &[PinKind::DigitalOut]);

            let mut success = 0;

            for _ in 0..50 {
                if a.allocate_pin(PinKind::DigitalOut).is_ok() {
                    success += 1;
                }
            }

            prop_assert!(success <= cap);
        }

        #[test]
        fn no_duplicate_pins(cap in 1u8..20) {
            let mut a = basic_allocator(cap, &[PinKind::DigitalOut]);

            let mut seen = std::collections::BTreeSet::new();

            for _ in 0..cap {
                let p = a.allocate_pin(PinKind::DigitalOut).unwrap();
                prop_assert!(seen.insert((p.addr, p.pin)));
            }
        }

        #[test]
        fn block_is_contiguous(cap in 1u8..20, size in 1u8..10) {
            let mut a = basic_allocator(cap, &[PinKind::PwmOut]);

            let res = a.allocate_pin_block(PinKind::PwmOut, size);

            match res {
                Ok(block) => {
                    let range = block.pins;
                    let len = *range.end() - *range.start() + 1;
                    prop_assert_eq!(len, size);
                }
                Err(_) => {
                    prop_assert!(true); // any failure is valid in simplified model
                }
            }
        }

        #[test]
        fn unsupported_always_fails(cap in 1u8..20) {
            let mut a = basic_allocator(cap, &[PinKind::DigitalOut]);

            let res = a.allocate_pin(PinKind::AnalogIn);
            prop_assert!(res.is_err());
        }
    }
}
