use super::avl::*;
use super::bounds::*;
use super::*;

#[derive(Clone)]
pub struct Bound<A, T> {
    bound: T,
    bounds: Bounds<A>,
}

impl<A, T> std::ops::Deref for Bound<A, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.bound
    }
}

impl<A, T> AsRef<T> for Bound<A, T> {
    fn as_ref(&self) -> &T {
        &self.bound
    }
}

pub struct BoundNode<A, T> {
    boundsl: Bounds<A>,
    boundsr: Bounds<A>,
    bounds: Bounds<A>,
    node: T,
}

#[derive(Clone)]
pub struct BoundTrees<BT>(BT);

impl<BT> BoundTrees<BT> {
    pub fn new(bt: BT) -> Self {
        Self(bt)
    }
}

impl<'a, BT: FunctorContext<'a>> FunctorContext<'a> for BoundTrees<BT> {
    type T = BT::T;
}

pub trait BinaryTreesBindable<'a>: MonadTrees<'a> {
    fn bounds_error<T: 'a + Send>(&self, error: BoundsError<Self::Key>) -> BTWrap<'a, Self, T>;

    fn bounds_bind<A: 'a + Send, B: 'a + Send>(
        &self,
        ra: Result<A, BoundsError<Self::Key>>,
        f: impl FnOnce(A) -> BTWrap<'a, Self, B>,
    ) -> BTWrap<'a, Self, B> {
        match ra {
            Ok(a) => f(a),
            Err(e) => self.bounds_error(e),
        }
    }
}

impl<BT: Keyed> Keyed for BoundTrees<BT> {
    type Key = BT::Key;
}

impl<BT: WithComparator> WithComparator for BoundTrees<BT> {
    type Comparator = BT::Comparator;

    fn comparator(&self) -> &Self::Comparator {
        self.0.comparator()
    }
}

impl<BT: WithComparator + BinaryTrees> BinaryTrees for BoundTrees<BT> {
    type Node = BoundNode<Self::Key, BT::Node>;
    type Reference = Bound<Self::Key, BT::Reference>;
    type Tree = Bound<Self::Key, BT::Tree>;

    fn split(&self, node: &Self::Node) -> Split<Self> {
        let (tl, tr, key) = self.0.split(&node.node);
        (
            Bound {
                bound: tl,
                bounds: node.boundsl.clone(),
            },
            Bound {
                bound: tr,
                bounds: node.boundsr.clone(),
            },
            key,
        )
    }

    fn equal(&self, rl: &Self::Reference, rr: &Self::Reference) -> bool {
        Bounds::equal(&rl.bounds, &rr.bounds, self.comparator())
            && self.0.equal(&rl.bound, &rr.bound)
    }

    fn refer(&self, tree: &Self::Tree) -> Option<Self::Reference> {
        Some(Bound {
            bound: self.0.refer(&tree.bound)?,
            bounds: tree.bounds.clone(),
        })
    }
}

impl<'a, BT> MonadTrees<'a> for BoundTrees<BT>
where
    BT: BinaryTreesBindable<'a> + WithComparator,
{
    fn resolve(&self, reference: &Self::Reference) -> BTWrap<'a, Self, Self::Node> {
        let ctx = self.0.clone();
        let bounds = reference.bounds.clone();
        Self::bind(self.0.resolve(&reference.bound), move |node| {
            let (_, _, key) = ctx.split(&node);
            ctx.bounds_bind(
                bounds.clone().split(&key, ctx.comparator()),
                |(boundsl, boundsr)| {
                    Self::pure(BoundNode {
                        boundsl,
                        boundsr,
                        bounds,
                        node,
                    })
                },
            )
        })
    }
}

impl<'a, BT> BinaryTreesTreeOf<'a> for BoundTrees<BT>
where
    BT: BinaryTreesBindable<'a> + BinaryTreesTreeOf<'a> + WithComparator,
{
    fn tree_of(&self, node: Self::Node) -> BTWrap<'a, Self, Self::Tree> {
        Self::fmap(self.0.tree_of(node.node), move |tree| Bound {
            bound: tree,
            bounds: node.bounds,
        })
    }
}

impl<'a, BT> BinaryTreesEmpty<'a> for BoundTrees<BT>
where
    BT: BinaryTreesBindable<'a> + BinaryTreesEmpty<'a> + WithComparator,
{
    fn empty(&self) -> Self::Tree {
        Bound {
            bound: self.0.empty(),
            bounds: Bounds::unbound(),
        }
    }

    fn split_key_empty(
        &self,
        tree: Self::Tree,
        key: Self::Key,
    ) -> BTWrap<'a, Self, KeySplit<Self>> {
        self.0.bounds_bind(
            tree.bounds.split(&key, self.comparator()),
            |(boundsl, boundsr)| {
                Self::bind(self.0.split_key_empty(tree.bound, key), |(tl, tr)| {
                    Self::pure((
                        Bound {
                            bound: tl,
                            bounds: boundsl,
                        },
                        Bound {
                            bound: tr,
                            bounds: boundsr,
                        },
                    ))
                })
            },
        )
    }
}

impl<BT: TreesHeight + WithComparator> TreesHeight for BoundTrees<BT> {
    fn height(&self, tree: &Self::Tree) -> u64 {
        self.0.height(&tree.bound)
    }
}

impl<'a, BT> TreesHeightError<'a> for BoundTrees<BT>
where
    BT: BinaryTreesBindable<'a> + TreesHeightError<'a> + WithComparator,
{
    fn height_error<T: 'a + Send>(&self, error: HeightError) -> BTWrap<'a, Self, T> {
        self.0.height_error(error)
    }
}

impl<'a, BT> BinaryTreesTryJoin<'a> for BoundTrees<BT>
where
    BT: BinaryTreesBindable<'a> + BinaryTreesTryJoin<'a> + WithComparator,
{
    fn try_join(
        &self,
        tl: Self::Tree,
        key: Self::Key,
        tr: Self::Tree,
    ) -> BTWrap<'a, Self, Self::Node> {
        let (boundsl, boundsr) = (tl.bounds, tr.bounds);
        self.0.bounds_bind(
            Bounds::join(boundsl.clone(), boundsr.clone(), &key, self.comparator()),
            |bounds| {
                Self::fmap(self.0.try_join(tl.bound, key, tr.bound), |node| BoundNode {
                    boundsl,
                    boundsr,
                    bounds,
                    node,
                })
            },
        )
    }
}

impl<'a, BT> BinaryTreesMutable<'a> for BoundTrees<BT>
where
    BT: BinaryTreesBindable<'a>
        + BinaryTreesTreeOf<'a>
        + BinaryTreesEmpty<'a>
        + TreesHeightError<'a>
        + BinaryTreesTryJoin<'a>
        + WithComparator,
{
    fn join_key(
        self,
        tl: Self::Tree,
        key: Self::Key,
        tr: Self::Tree,
    ) -> BTWrap<'a, Self, Self::Node> {
        self.join_key_balanced(tl, key, tr)
    }
}