pub mod avl;
pub mod balancing;
pub mod bound;
pub mod bounds;

use std::fmt::Display;

use crate::flow::comparator::*;
use crate::func::{context::*, *};

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

pub type BTWrap<'a, BT, A> = WrapC<'a, A, BT>;

pub trait BinaryTrees<'a>: FunctorContext<'a, T = Self::_Tm> + Clone {
    type Node: 'a;
    type Reference: 'a;
    type Tree: 'a;
    type Key: 'a + Clone;
    type Comparator: Comparator<Self::Key>;

    type _Tm: Monad<'a>;

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

pub trait BinaryTreesTreeOf<'a>: BinaryTrees<'a> {
    fn tree_of(&self, node: Self::Node) -> BTWrap<'a, Self, Self::Tree>;

    fn tree_bind(self, fnode: BTWrap<'a, Self, Self::Node>) -> BTWrap<'a, Self, Self::Tree> {
        Self::bind(fnode, move |node| self.tree_of(node))
    }
}

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

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

pub trait BinaryTreesMutable<'a>: BinaryTreesEmpty<'a> + BinaryTreesTreeOf<'a> {
    fn join_key(
        self,
        tl: Self::Tree,
        key: Self::Key,
        tr: Self::Tree,
    ) -> BTWrap<'a, Self, Self::Node>;

    fn join_key_tree(
        self,
        tl: Self::Tree,
        key: Self::Key,
        tr: Self::Tree,
    ) -> BTWrap<'a, Self, Self::Tree> {
        self.clone().tree_bind(self.join_key(tl, key, tr))
    }

    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_node(
        self,
        node: Self::Node,
        key: Self::Key,
    ) -> BTWrap<'a, Self, KeySplit<'a, Self>> {
        let (tl, tr, k) = self.split(&node);
        match self.comparator().compare(&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: Self::Key) -> 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: Self::Key) -> 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: Self::Key) -> BTWrap<'a, Self, Self::Tree> {
        self.clone().tree_bind(self.add(tree, key))
    }

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

#[derive(Debug)]
pub enum HeightError {
    LeafHeight(u64),
    NodeHeight,
}

impl Display for HeightError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::NodeHeight => write!(f, "invalid node height: 0"),
            Self::LeafHeight(height) => {
                write!(f, "invalid leaf height: {height}!=0")
            }
        }
    }
}

pub trait BinaryTreesHeight<'a>: BinaryTrees<'a> {
    fn height(&self, tree: &Self::Tree) -> u64;
    fn height_error<T: 'a>(&self, error: HeightError) -> BTWrap<'a, Self, T>;
}

pub trait BinaryTreesTryJoin<'a>: BinaryTrees<'a> {
    fn try_join(
        &self,
        tl: Self::Tree,
        key: Self::Key,
        tr: Self::Tree,
    ) -> BTWrap<'a, Self, Self::Node>;
}