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

use super::{binary::*, *};

#[derive(Clone)]
struct Bounds<A> {
    l: Option<A>,
    r: Option<A>,
}

pub enum BoundsError<A> {
    BoundsViolated { l: A, r: A },
    CannotJoin(A),
}

pub enum BoundError<A, E> {
    Avl(AvlError<E>),
    Bounds(BoundsError<A>),
}

impl<A, E> From<AvlError<E>> for BoundError<A, E> {
    fn from(value: AvlError<E>) -> Self {
        Self::Avl(value)
    }
}

impl<A, E> From<BoundsError<A>> for BoundError<A, E> {
    fn from(value: BoundsError<A>) -> Self {
        Self::Bounds(value)
    }
}

pub type BoundErrorA<'a, Ctx, A> = BoundError<A, ParseError<'a, Ctx, Fctr<'a, Ctx, A>>>;

pub type BoundResult<'a, Ctx, A, T> = Result<T, BoundErrorA<'a, Ctx, A>>;

pub type BoundNodeResult<'a, Ctx, A> = BoundResult<'a, Ctx, A, BoundNode<'a, Ctx, A>>;

pub type BoundTreeResult<'a, Ctx, A> = BoundResult<'a, Ctx, A, BoundTree<'a, Ctx, A>>;

pub struct BoundTree<'a, Ctx: Context<'a>, A: Mentionable<'a, Ctx>> {
    tree: AvlTree<'a, Ctx, A>,
    bounds: Bounds<A>,
}

pub struct BoundNode<'a, Ctx: Context<'a>, A: Mentionable<'a, Ctx>> {
    l: BoundTree<'a, Ctx, A>,
    r: BoundTree<'a, Ctx, A>,
    key: A,
    bounds: Bounds<A>,
}

pub struct BoundReference<'a, Ctx: Context<'a>, A: Mentionable<'a, Ctx>> {
    reference: AvlReference<'a, Ctx, A>,
    bounds: Bounds<A>,
}

impl<A: Clone> Bounds<A> {
    fn unbound() -> Self {
        Bounds { l: None, r: None }
    }

    fn ordered(l: &A, r: &A, comparator: &impl Comparator<A>) -> Result<(), BoundsError<A>> {
        if let Comparison::R = comparator.pick_smaller(l, r) {
            Err(BoundsError::BoundsViolated {
                l: l.clone(),
                r: r.clone(),
            })
        } else {
            Ok(())
        }
    }

    fn new(
        l: Option<A>,
        r: Option<A>,
        comparator: &impl Comparator<A>,
    ) -> Result<Self, BoundsError<A>> {
        if let (Some(kl), Some(kr)) = (&l, &r) {
            Self::ordered(kl, kr, comparator)?
        }
        Ok(Bounds { l, r })
    }

    fn split(
        self,
        key: &A,
        comparator: &impl Comparator<A>,
    ) -> Result<(Self, Self), BoundsError<A>> {
        Ok((
            Self::new(self.l, Some(key.clone()), comparator)?,
            Self::new(Some(key.clone()), self.r, comparator)?,
        ))
    }

    fn join(
        l: Self,
        r: Self,
        key: &A,
        comparator: &impl Comparator<A>,
    ) -> Result<Self, BoundsError<A>> {
        if let Some(lr) = &l.r {
            Self::ordered(lr, key, comparator)?
        } else {
            Err(BoundsError::CannotJoin(key.clone()))?
        }
        if let Some(rl) = &r.l {
            Self::ordered(key, rl, comparator)?
        } else {
            Err(BoundsError::CannotJoin(key.clone()))?
        }
        Self::new(l.l, r.r, comparator)
    }

    pub fn equal_bound(l: &Option<A>, r: &Option<A>, comparator: &impl Comparator<A>) -> bool {
        match (l, r) {
            (None, None) => true,
            (Some(kl), Some(kr)) => matches!(comparator.pick_smaller(kl, kr), Comparison::E),
            _ => false,
        }
    }

    pub fn equal(bl: &Self, br: &Self, comparator: &impl Comparator<A>) -> bool {
        Self::equal_bound(&bl.l, &br.l, comparator) && Self::equal_bound(&bl.r, &br.r, comparator)
    }
}

pub type BoundSplit<'a, Ctx, A> = (BoundTree<'a, Ctx, A>, BoundTree<'a, Ctx, A>, A);

impl<'a, Ctx: Context<'a>, A: Mentionable<'a, Ctx> + Clone> BoundNode<'a, Ctx, A> {
    pub fn into_tree(self) -> BoundTreeResult<'a, Ctx, A> {
        let tree = AvlNode::new(self.l.tree, self.r.tree, self.key)?.into_tree()?;
        let bounds = self.bounds;
        Ok(BoundTree { tree, bounds })
    }

    pub fn into_tree_resolution(self) -> BrTree<'a, Ctx, A> {
        match self.into_tree() {
            Ok(tree) => Ctx::pure(Ok(tree)),
            Err(e) => Ctx::pure(Err(ResolutionError::Parse(e))),
        }
    }

    pub fn new(
        l: BoundTree<'a, Ctx, A>,
        r: BoundTree<'a, Ctx, A>,
        key: A,
        comparator: &impl Comparator<A>,
    ) -> BoundNodeResult<'a, Ctx, A> {
        let bounds = Bounds::join(l.bounds.clone(), r.bounds.clone(), &key, comparator)?;
        Ok(Self { l, r, key, bounds })
    }

    pub fn split(&self) -> BoundSplit<'a, Ctx, A> {
        (self.l.clone(), self.r.clone(), self.key.clone())
    }
}

pub type BoundKeySplit<'a, Ctx, A> = (BoundTree<'a, Ctx, A>, BoundTree<'a, Ctx, A>);

impl<'a, Ctx: Context<'a>, A: Mentionable<'a, Ctx> + Clone> BoundTree<'a, Ctx, A> {
    pub fn height(&self) -> u64 {
        self.tree.height
    }

    pub fn empty(factory: Fctr<'a, Ctx, A>) -> Self {
        Self {
            tree: AvlTree::empty(factory),
            bounds: Bounds::unbound(),
        }
    }

    fn empty_bound(factory: Fctr<'a, Ctx, A>, bounds: Bounds<A>) -> Self {
        Self {
            tree: AvlTree::empty(factory),
            bounds,
        }
    }

    pub fn reference(&self) -> Option<BoundReference<'a, Ctx, A>> {
        let reference = self.tree.reference()?;
        let bounds = self.bounds.clone();
        Some(BoundReference { reference, bounds })
    }

    fn split_empty(
        bounds: Bounds<A>,
        key: &A,
        factory: &Fctr<'a, Ctx, A>,
        comparator: &impl Comparator<A>,
    ) -> Result<(Self, Self), BoundsError<A>> {
        let (bl, br) = bounds.split(key, comparator)?;
        Ok((
            Self::empty_bound(factory.clone(), bl),
            Self::empty_bound(factory.clone(), br),
        ))
    }

    pub fn split_empty_res(
        self,
        key: A,
        factory: Fctr<'a, Ctx, A>,
        comparator: Rc<impl 'a + Comparator<A>>,
    ) -> BrKeySplit<'a, Ctx, A> {
        Ctx::pure(
            Self::split_empty(self.bounds, &key, &factory, comparator.as_ref())
                .map_err(BoundError::Bounds)
                .map_err(ResolutionError::Parse),
        )
    }
}

impl<'a, Ctx: Context<'a>, A: Mentionable<'a, Ctx> + Clone> Clone for BoundTree<'a, Ctx, A> {
    fn clone(&self) -> Self {
        Self {
            tree: self.tree.clone(),
            bounds: self.bounds.clone(),
        }
    }
}

pub type BoundResolutionError<'a, Ctx, A> =
    ResolutionError<<Ctx as Context<'a>>::LookupError, BoundErrorA<'a, Ctx, A>>;

pub type BoundResolutionResult<'a, Ctx, A, T> = Result<T, BoundResolutionError<'a, Ctx, A>>;

pub type BoundResolution<'a, Ctx, A, T> = Wrapped<'a, Ctx, BoundResolutionResult<'a, Ctx, A, T>>;

pub type BrNode<'a, Ctx, A> = BoundResolution<'a, Ctx, A, BoundNode<'a, Ctx, A>>;

pub type BrTree<'a, Ctx, A> = BoundResolution<'a, Ctx, A, BoundTree<'a, Ctx, A>>;

pub type BrKeySplit<'a, Ctx, A> = BoundResolution<'a, Ctx, A, BoundKeySplit<'a, Ctx, A>>;

impl<'a, Ctx: Context<'a>, A: Mentionable<'a, Ctx> + Clone> BoundReference<'a, Ctx, A> {
    pub fn resolve(&self, comparator: Rc<impl 'a + Comparator<A>>) -> BrNode<'a, Ctx, A> {
        let bounds = self.bounds.clone();
        Ctx::fmap(self.reference.resolve(), move |resolved| {
            let node = resolved
                .map_err(|e| e.map_parse(BoundError::Avl))?
                .as_ref()
                .clone();
            let (bl, br) = bounds
                .clone()
                .split(&node.key, comparator.as_ref())
                .map_err(BoundError::Bounds)
                .map_err(ResolutionError::Parse)?;
            Ok(BoundNode {
                l: BoundTree {
                    tree: node.l,
                    bounds: bl,
                },
                r: BoundTree {
                    tree: node.r,
                    bounds: br,
                },
                key: node.key,
                bounds,
            })
        })
    }

    pub fn equal(rl: &Self, rr: &Self, comparator: &impl Comparator<A>) -> bool {
        rl.reference == rr.reference && Bounds::equal(&rl.bounds, &rr.bounds, comparator)
    }
}