use std::{fmt::Display, marker::PhantomData, rc::Rc};

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

use super::*;

struct AvlN<A> {
    l: AvlT<A>,
    r: AvlT<A>,
    key: Rc<A>,
}

impl<A> Clone for AvlN<A> {
    fn clone(&self) -> Self {
        Self {
            l: self.l.clone(),
            r: self.r.clone(),
            key: self.key.clone(),
        }
    }
}

impl<A: Display> Display for AvlN<A> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_fmt(format_args!("({} {} {})", self.l, self.key, self.r))
    }
}

#[cfg(test)]
impl<A> AvlN<A> {
    fn balanced(&self) -> bool {
        let (lh, rh) = (self.l.height, self.r.height);
        std::cmp::max(lh, rh) - std::cmp::min(lh, rh) < 2 && self.l.balanced() && self.r.balanced()
    }
}

struct AvlR<A> {
    node: Rc<AvlN<A>>,
}

impl<A> Clone for AvlR<A> {
    fn clone(&self) -> Self {
        Self {
            node: self.node.clone(),
        }
    }
}

impl<A: Display> Display for AvlR<A> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_fmt(format_args!("{}", self.node))
    }
}

#[cfg(test)]
impl<A> AvlR<A> {
    fn balanced(&self) -> bool {
        self.node.balanced()
    }
}

struct AvlT<A> {
    reference: Option<AvlR<A>>,
    height: u64,
}

impl<A> Clone for AvlT<A> {
    fn clone(&self) -> Self {
        Self {
            reference: self.reference.clone(),
            height: self.height,
        }
    }
}

impl<A: Display> Display for AvlT<A> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match &self.reference {
            Some(reference) => f.write_fmt(format_args!("{}", reference)),
            None => f.write_fmt(format_args!("-")),
        }
    }
}

#[cfg(test)]
impl<A> AvlT<A> {
    fn balanced(&self) -> bool {
        match &self.reference {
            Some(reference) => reference.balanced(),
            None => true,
        }
    }
}

struct AvlTs<A> {
    _a: PhantomData<A>,
}

impl<A> AvlTs<A> {
    fn new() -> Self {
        Self { _a: PhantomData }
    }
}

impl<A> Clone for AvlTs<A> {
    fn clone(&self) -> Self {
        Self::new()
    }
}

impl<'a, A: 'a + PartialOrd> BinaryTrees<'a> for AvlTs<A> {
    type Node = AvlN<A>;

    type Reference = AvlR<A>;

    type Tree = AvlT<A>;

    type Key = A;

    type Comparator = DefaultComparator;

    type T = instances::solo::SoloInstance;

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

    fn split(&self, node: &Self::Node) -> Split<'a, Self> {
        (node.l.clone(), node.r.clone(), node.key.clone())
    }

    fn tree_of(&self, node: Self::Node) -> BTWrap<'a, Self, Self::Tree> {
        AvlT {
            height: std::cmp::max(node.l.height, node.r.height)
                .checked_add(1)
                .unwrap(),
            reference: Some(AvlR { node: node.into() }),
        }
    }

    fn resolve(&self, reference: &Self::Reference) -> BTWrap<'a, Self, Self::Node> {
        reference.node.as_ref().clone()
    }

    fn equal(&self, _rhs: &Self::Reference, _lhs: &Self::Reference) -> bool {
        false
    }

    fn refer(&self, tree: &Self::Tree) -> Option<Self::Reference> {
        tree.reference.clone()
    }
}

impl<'a, A: 'a + PartialOrd> BinaryTreesMutable<'a> for AvlTs<A> {
    fn empty(&self) -> Self::Tree {
        AvlT {
            reference: None,
            height: 0,
        }
    }

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

    fn split_key_empty(
        &self,
        _tree: Self::Tree,
        _key: KeyRc<'a, Self>,
    ) -> BTWrap<'a, Self, KeySplit<'a, Self>> {
        (self.empty(), self.empty())
    }
}

impl<'a, A: 'a + PartialOrd> BinaryTreesAvl<'a> for AvlTs<A> {
    fn height(&self, tree: &Self::Tree) -> u64 {
        tree.height
    }

    fn leaf_height_error<T: 'a>(&self, height: u64) -> BTWrap<'a, Self, T> {
        panic!("leaf height error: {height}.")
    }

    fn join_key_unbalanced(
        &self,
        tl: Self::Tree,
        key: KeyRc<'a, Self>,
        tr: Self::Tree,
    ) -> BTWrap<'a, Self, Self::Node> {
        AvlN { l: tl, r: tr, key }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test() {
        let trees = AvlTs::new();
        let mut tree = trees.empty();
        for i in [
            8, 3, 10, 17, 0, 13, 6, 1, 11, 5, 4, 7, 18, 14, 15, 9, 2, 19, 16, 12,
        ] {
            tree = trees.clone().add_tree(tree, i.into());
            assert!(tree.balanced());
            // println!("{} {}", tree.height, tree);
        }
        for i in [
            2, 9, 4, 7, 8, 10, 17, 1, 13, 15, 18, 12, 5, 0, 3, 6, 16, 19, 14, 11,
        ] {
            tree = trees.clone().remove(tree, i.into());
            assert!(tree.balanced());
            // println!("{} {}", tree.height, tree);
        }
        // assert!(false);
    }
}