use std::{net::Ipv4Addr, sync::Arc};

use arc_swap::ArcSwap;

#[derive(Default, Clone)]
pub struct TrieNode {
    child: [Option<Box<TrieNode>>; 2],
    prefix_len: u8,
    nexthop: Option<Ipv4Addr>,
}

#[derive(Default, Clone)]
struct IpTrie {
    root: TrieNode,
}

impl IpTrie {
    fn new() -> Self {
        Self {
            root: TrieNode::default()
        }
    }

    fn insert(&mut self, prefix: u32, prefix_len: u8, nexthop: Ipv4Addr) {
        if prefix_len > 32 {
            return;
        }

        let mut node = &mut self.root;
        for i in 0..prefix_len {
            let bit = ((prefix >> (31-i)) & 1) as usize;
            node = node.child[bit].get_or_insert_with(|| Box::new(TrieNode::default()));

        }
        if prefix_len > node.prefix_len {
            node.prefix_len = prefix_len;
            node.nexthop = Some(nexthop);
        }
    }

    fn lookup(&self, ip: u32) -> Option<(u8, Ipv4Addr)>{
        let mut node = &self.root;
        let mut best = None;

        for i in 0..32 {
            if node.nexthop.is_some() {
                best = Some((node.prefix_len, node.nexthop.unwrap()));
            }
            let bit = ((ip>>(31-i)) & 1) as usize;
            match &node.child[bit] {
                Some(child) => {
                    node = child;
                }
                None => {
                    break;
                }
            }
        }
        if node.nexthop.is_some() {
            best = Some((node.prefix_len, node.nexthop.unwrap()));
        }
        best
    }
}

pub struct RouteTableTrie {
    trie: ArcSwap<IpTrie>,
}

impl RouteTableTrie {
    pub fn new() -> Self {
        Self {
            trie: ArcSwap::new(Arc::new(IpTrie::default()))
        }
    }

    pub fn clear(&self)  {
        self.trie.store(Arc::new(IpTrie::default()));
    }

    pub fn lookup(&self, ip: u32) -> Option<(u8, Ipv4Addr)> {
        let trie = self.trie.load();
        trie.lookup(ip)
    }

    pub fn insert(&self, prefix: u32, prefix_len: u8, nexthop: Ipv4Addr) {
        let old = self.trie.load();
        let mut new_trie = (*(*old)).clone();

        new_trie.insert(prefix, prefix_len, nexthop);
        self.trie.store(Arc::new(new_trie));
    }
}