mod avl;

use std::rc::Rc;

use crate::flow::comparator::*;
use crate::func::*;

pub type KeyRc<'a, BT> = Rc<<BT as BinaryTrees<'a>>::Key>;

pub type Split<'a, BT> = (
    <BT as BinaryTrees<'a>>::Tree,
    <BT as BinaryTrees<'a>>::Tree,
    KeyRc<'a, BT>,
);
pub type KeySplit<'a, BT> = (<BT as BinaryTrees<'a>>::Tree, <BT as BinaryTrees<'a>>::Tree);

pub type BTWrap<'a, BT, A> = Wrap<'a, A, <BT as BinaryTrees<'a>>::T>;

pub trait BinaryTrees<'a>: 'a + Clone {
    type Node: 'a;
    type Reference: 'a;
    type Tree: 'a;
    type Key: 'a;
    type Comparator: Comparator<Self::Key>;

    type T: Monad<'a>;

    fn comparator(&self) -> &Self::Comparator;
    fn split(&self, node: &Self::Node) -> Split<'a, Self>;
    fn tree_of(&self, node: Self::Node) -> BTWrap<'a, Self, Self::Tree>;
    fn resolve(&self, reference: &Self::Reference) -> BTWrap<'a, Self, Self::Node>;
    fn equal(&self, rhs: &Self::Reference, lhs: &Self::Reference) -> bool;
    fn refer(&self, tree: &Self::Tree) -> Option<Self::Reference>;

    fn bind<A: 'a, B: 'a>(
        fa: BTWrap<'a, Self, A>,
        f: impl 'a + FnOnce(A) -> BTWrap<'a, Self, B>,
    ) -> BTWrap<'a, Self, B> {
        <Self::T as Monad>::bind(fa, f)
    }

    fn pure<A: 'a>(a: A) -> BTWrap<'a, Self, A> {
        <Self::T as Pure>::pure(a)
    }

    fn fmap<A: 'a, B: 'a>(
        fa: BTWrap<'a, Self, A>,
        f: impl 'a + FnOnce(A) -> B,
    ) -> BTWrap<'a, Self, B> {
        <Self::T as Functor>::fmap(fa, f)
    }

    fn fail<A: 'a, E: 'a>(e: E) -> BTWrap<'a, Self, A>
    where
        Self::T: Fail<'a, E>,
    {
        <Self::T as Fail<_>>::fail(e)
    }
}

pub trait BinaryTreesMutable<'a>: BinaryTrees<'a> {
    fn empty(&self) -> Self::Tree;

    fn join_key(
        self,
        tl: Self::Tree,
        key: KeyRc<'a, Self>,
        tr: Self::Tree,
    ) -> BTWrap<'a, Self, Self::Node>;

    fn join_key_tree(
        self,
        tl: Self::Tree,
        key: KeyRc<'a, Self>,
        tr: Self::Tree,
    ) -> BTWrap<'a, Self, Self::Tree> {
        Self::bind(self.clone().join_key(tl, key, tr), move |n| self.tree_of(n))
    }

    fn join(self, tl: Self::Tree, tr: Self::Tree) -> BTWrap<'a, Self, Self::Tree> {
        let Some(rl) = self.refer(&tl) else {
            return Self::pure(tr)
        };
        let Some(rr) = self.refer(&tr) else {
            return Self::pure(tl)
        };
        Self::T::bind2(self.resolve(&rl), self.resolve(&rr), move |nl, nr| {
            let (tll, tlr, kl) = self.split(&nl);
            let (trl, trr, kr) = self.split(&nr);
            let ft = self.clone().join(tlr, trl);
            let ft = {
                let ctx = self.clone();
                Self::bind(ft, move |t| ctx.join_key_tree(tll, kl, t))
            };
            let ft = {
                let ctx = self.clone();
                Self::bind(ft, move |t| ctx.join_key_tree(t, kr, trr))
            };
            ft
        })
    }

    fn split_key_empty(
        &self,
        tree: Self::Tree,
        key: KeyRc<'a, Self>,
    ) -> BTWrap<'a, Self, KeySplit<'a, Self>>;

    fn split_key_node(
        self,
        node: Self::Node,
        key: KeyRc<'a, Self>,
    ) -> BTWrap<'a, Self, KeySplit<'a, Self>> {
        let (tl, tr, k) = self.split(&node);
        match self.comparator().pick_smaller(&key, &k) {
            Comparison::L => Self::bind(self.clone().split_key(tl, key), move |(tll, tlr)| {
                Self::fmap(self.join_key_tree(tlr, k, tr), |t| (tll, t))
            }),
            Comparison::E => Self::pure((tl, tr)),
            Comparison::R => Self::bind(self.clone().split_key(tr, key), move |(trl, trr)| {
                Self::fmap(self.join_key_tree(tl, k, trl), |t| (t, trr))
            }),
        }
    }

    fn split_key(
        self,
        tree: Self::Tree,
        key: KeyRc<'a, Self>,
    ) -> BTWrap<'a, Self, KeySplit<'a, Self>> {
        match self.refer(&tree) {
            Some(reference) => Self::bind(self.resolve(&reference), |node| {
                self.split_key_node(node, key)
            }),
            None => self.split_key_empty(tree, key),
        }
    }

    fn add(self, tree: Self::Tree, key: KeyRc<'a, Self>) -> BTWrap<'a, Self, Self::Node> {
        Self::bind(self.clone().split_key(tree, key.clone()), |(tl, tr)| {
            self.join_key(tl, key, tr)
        })
    }

    fn add_tree(self, tree: Self::Tree, key: KeyRc<'a, Self>) -> BTWrap<'a, Self, Self::Tree> {
        Self::bind(self.clone().add(tree, key), move |n| self.tree_of(n))
    }

    fn remove(self, tree: Self::Tree, key: KeyRc<'a, Self>) -> BTWrap<'a, Self, Self::Tree> {
        Self::bind(self.clone().split_key(tree, key.clone()), |(tl, tr)| {
            self.join(tl, tr)
        })
    }
}

pub trait BinaryTreesAvl<'a>: BinaryTrees<'a> {
    fn height(&self, tree: &Self::Tree) -> u64;

    fn leaf_height_error<T: 'a>(&self, height: u64) -> BTWrap<'a, Self, T>;

    fn assume_node(&self, tree: &Self::Tree) -> BTWrap<'a, Self, Self::Node> {
        match self.refer(tree) {
            Some(reference) => self.resolve(&reference),
            None => self.leaf_height_error(self.height(tree)),
        }
    }

    fn join_key_unbalanced(
        &self,
        tl: Self::Tree,
        key: KeyRc<'a, Self>,
        tr: Self::Tree,
    ) -> BTWrap<'a, Self, Self::Node>;

    fn join_key_balanced_tree(
        self,
        tl: Self::Tree,
        key: KeyRc<'a, Self>,
        tr: Self::Tree,
    ) -> BTWrap<'a, Self, Self::Tree> {
        Self::bind(self.clone().join_key_balanced(tl, key, tr), move |node| {
            self.tree_of(node)
        })
    }

    fn join_key_balanced(
        self,
        tl: Self::Tree,
        key: KeyRc<'a, Self>,
        tr: Self::Tree,
    ) -> BTWrap<'a, Self, Self::Node> {
        let (lh, rh) = (self.height(&tl), self.height(&tr));
        match (lh.saturating_sub(rh), rh.saturating_sub(lh)) {
            (0, 0) | (0, 1) | (1, 0) => self.join_key_unbalanced(tl, key, tr),
            (0, _) => Self::bind(self.assume_node(&tr), move |nr| {
                let (trl, trr, kr) = self.split(&nr);
                let (rlh, rrh) = (self.height(&trl), self.height(&trr));
                if rlh > rrh {
                    Self::bind(self.assume_node(&trl), move |nrl| {
                        let (trll, trlr, krl) = self.split(&nrl);
                        Self::T::bind2(
                            self.clone().join_key_balanced_tree(tl, key, trll),
                            self.clone().join_key_balanced_tree(trlr, kr, trr),
                            |ti, to| self.join_key_balanced(ti, krl, to),
                        )
                    })
                } else {
                    Self::bind(self.clone().join_key_balanced_tree(tl, key, trl), |t| {
                        self.join_key_balanced(t, kr, trr)
                    })
                }
            }),
            (_, 0) => Self::bind(self.assume_node(&tl), move |nl| {
                let (tll, tlr, kl) = self.split(&nl);
                let (llh, lrh) = (self.height(&tll), self.height(&tlr));
                if llh < lrh {
                    Self::bind(self.assume_node(&tlr), move |nlr| {
                        let (tlrl, tlrr, klr) = self.split(&nlr);
                        Self::T::bind2(
                            self.clone().join_key_balanced_tree(tll, kl, tlrl),
                            self.clone().join_key_balanced_tree(tlrr, key, tr),
                            |to, ti| self.join_key_balanced(to, klr, ti),
                        )
                    })
                } else {
                    Self::bind(self.clone().join_key_balanced_tree(tlr, key, tr), |t| {
                        self.join_key_balanced(tll, kl, t)
                    })
                }
            }),
            (_, _) => unreachable!(),
        }
    }
}