use crate::func::context::*;

use super::*;

pub struct AvlReference<'a, Ctx: Context<'a>, A: Mentionable<'a, Ctx>> {
    node: Point<'a, Ctx, AvlNode<'a, Ctx, A>>,
    height: u64,
}

impl<'a, Ctx: Context<'a>, A: Mentionable<'a, Ctx>> PartialEq for AvlReference<'a, Ctx, A> {
    fn eq(&self, other: &Self) -> bool {
        self.node == other.node && self.height == other.height
    }
}

pub type AvlResult<'a, Ctx, A, T> = Result<T, AvlError<ParseError<'a, Ctx, Fctr<'a, Ctx, A>>>>;

pub type AvlNodeResult<'a, Ctx, A> = AvlResult<'a, Ctx, A, AvlNode<'a, Ctx, A>>;

pub type AvlTreeResult<'a, Ctx, A> = AvlResult<'a, Ctx, A, AvlTree<'a, Ctx, A>>;

impl<'a, Ctx: Context<'a>, A: Mentionable<'a, Ctx>> AvlNode<'a, Ctx, A> {
    fn parent_height(&self) -> AvlResult<'a, Ctx, A, u64> {
        std::cmp::max(self.l.height, self.r.height)
            .checked_add(1)
            .ok_or(AvlError::Balancing(BalancingError::HeightOverflow))
    }

    pub fn into_tree(self) -> AvlTreeResult<'a, Ctx, A> {
        let height = self.parent_height()?;
        let node = self.into();
        Ok(AvlTree { node, height })
    }

    pub fn new(
        l: AvlTree<'a, Ctx, A>,
        r: AvlTree<'a, Ctx, A>,
        key: A,
    ) -> AvlNodeResult<'a, Ctx, A> {
        balanced(l.height, r.height)?;
        Ok(Self { l, r, key })
    }

    fn matches_height(&self, parent_height: u64) -> AvlResult<'a, Ctx, A, ()> {
        let expected = self.parent_height()?;
        if expected != parent_height {
            Err(BalancingError::HeightMismatch {
                children: (self.l.height, self.r.height),
                parent: parent_height,
            })?
        }
        Ok(())
    }
}

impl<'a, Ctx: Context<'a>, A: Mentionable<'a, Ctx>> AvlTree<'a, Ctx, A> {
    pub fn empty(factory: Fctr<'a, Ctx, A>) -> Self {
        Self {
            node: Nullable::Null(AvlNodeFactory(factory)),
            height: 0,
        }
    }

    pub fn reference(&self) -> Option<AvlReference<'a, Ctx, A>> {
        match self.node {
            Nullable::Null(_) => None,
            Nullable::NotNull(ref point) => Some(AvlReference {
                node: point.clone(),
                height: self.height,
            }),
        }
    }
}

impl<'a, Ctx: Context<'a>, A: Mentionable<'a, Ctx>> AvlReference<'a, Ctx, A> {
    pub fn resolve(&self) -> Resolution<'a, Ctx, AvlNode<'a, Ctx, A>> {
        let parent_height = self.height;
        Ctx::fmap(self.node.resolve(), move |resolved| {
            let node = resolved?;
            node.matches_height(parent_height)
                .map_err(ResolutionError::Parse)?;
            Ok(node)
        })
    }

    pub fn height(&self) -> u64 {
        self.height
    }
}

impl<'a, Ctx: Context<'a>, A: Mentionable<'a, Ctx>> Clone for AvlTree<'a, Ctx, A> {
    fn clone(&self) -> Self {
        Self {
            node: self.node.clone(),
            height: self.height,
        }
    }
}

impl<'a, Ctx: Context<'a>, A: Mentionable<'a, Ctx> + Clone> Clone for AvlNode<'a, Ctx, A> {
    fn clone(&self) -> Self {
        Self {
            l: self.l.clone(),
            r: self.r.clone(),
            key: self.key.clone(),
        }
    }
}