#![allow(dead_code)]
use crate::common::ip::Block;
use crate::common::ip::Meta;
use crate::common::ip::BITMASK;

pub struct Trie {
    root: Node,
}

#[derive(Debug)]
struct Node {
    block: Option<Block>,
    left: Option<Box<Node>>,
    right: Option<Box<Node>>,
}

impl Node {
    pub const NULL: Node = Self {
        block: None,
        left: None,
        right: None,
    };

    fn new(block: Option<Block>) -> Self {
        Node {
            block,
            left: None,
            right: None,
        }
    }

    fn is_leaf(&self) -> bool {
        self.left.is_none() && self.right.is_none()
    }
}

impl Trie {
    pub fn new() -> Self {
        Trie { root: Node::NULL }
    }

    pub fn insert(&mut self, block: Block, _meta: Meta) {
        let mut node = &mut self.root;
        let mut addr = block.addr;

        for _ in 0..block.mlen {
            if addr & BITMASK == 0 {
                if node.left.is_none() {
                    node.left = Some(Box::new(Node::NULL));
                }
                node = node.left.as_deref_mut().unwrap();
            } else {
                if node.right.is_none() {
                    node.right = Some(Box::new(Node::NULL));
                }
                node = node.right.as_deref_mut().unwrap();
            }
            addr <<= 1;
        }

        node.block = Some(block);
    }

    pub fn find(&self, needle: &Block) -> Option<Block> {
        self.find_block_node(needle)?.block
    }

    pub fn parent(&self, needle: &Block) -> Option<Block> {
        let mut node = &self.root;
        let mut addr = needle.addr;
        let mut parent = None;

        for _ in 0..needle.mlen {
            if node.block.is_some() {
                parent = node.block;
            }

            if addr & BITMASK == 0 {
                if node.left.is_some() {
                    node = node.left.as_deref().unwrap();
                } else {
                    break;
                }
            } else if node.right.is_some() {
                node = node.right.as_deref().unwrap();
            } else {
                break;
            }
            addr <<= 1;
        }

        parent
    }

    pub fn children(&self, parent: &Block) -> Vec<Block> {
        let root: &Node = match self.find_block_node(parent) {
            Some(root) => root,
            None => return vec![],
        };

        let mut children: Vec<Block> = Vec::new();
        let mut stack: Vec<&Node> = vec![root];

        while let Some(node) = stack.pop() {
            if node.is_leaf() && node.block.is_some() && node.block.unwrap() != root.block.unwrap()
            {
                children.push(node.block.unwrap());
                continue;
            }

            if node.left.is_some() {
                stack.push(node.left.as_deref().unwrap());
            }
            if node.right.is_some() {
                stack.push(node.right.as_deref().unwrap());
            }
        }

        children
    }

    fn find_block_node(&self, needle: &Block) -> Option<&Node> {
        let mut node = &self.root;
        let mut addr = needle.addr;
        let mut match_ = None;

        for _ in 0..(needle.mlen + 1) {
            if node.block.is_some() {
                match_ = Some(node);
            }

            if addr & BITMASK == 0 {
                if node.left.is_some() {
                    node = node.left.as_deref().unwrap();
                } else {
                    break;
                }
            } else if node.right.is_some() {
                node = node.right.as_deref().unwrap();
            } else {
                break;
            }
            addr <<= 1;
        }

        match match_?.block {
            Some(block) if block == *needle => match_,
            _ => None,
        }
    }
}