pub mod binary;
pub mod bounds;
pub mod context;

use std::{error::Error, fmt::Display, rc::Rc};

use crate::rcore::*;
use crate::rstd::{
    atomic::{au64::*, *},
    nullable::*,
    point::*,
};

#[derive(Debug)]
pub enum AvlError<E> {
    Int(IntParseError),
    Point(PointParseError),
    Key(E),
    LeafHeight(u64),
    NodeHeight,
    Balance(u64, u64),
    HeightOverflow,
    HeightMismatch { child: (u64, u64), parent: u64 },
}

impl<E> From<IntParseError> for AvlError<E> {
    fn from(value: IntParseError) -> Self {
        Self::Int(value)
    }
}

impl<E> From<PointParseError> for AvlError<E> {
    fn from(value: PointParseError) -> Self {
        Self::Point(value)
    }
}

impl<E: Display> Display for AvlError<E> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Int(int_error) => {
                f.write_fmt(format_args!("failed to parse AVL tree height: {int_error}"))
            }
            Self::Point(point_error) => f.write_fmt(format_args!(
                "failed to parse AVL node reference: {point_error}"
            )),
            Self::Key(key_error) => {
                f.write_fmt(format_args!("failed to parse AVL node key: {key_error}"))
            }
            Self::NodeHeight => f.write_fmt(format_args!("invalid AVL non-leaf height: 0.")),
            Self::LeafHeight(height) => {
                f.write_fmt(format_args!("invalid AVL leaf height: {height}!=0."))
            }
            Self::Balance(hl, hr) => f.write_fmt(format_args!("unbalanced AVL node: {hl} {hr}.")),
            Self::HeightOverflow => f.write_fmt(format_args!("AVL tree height overflow")),
            Self::HeightMismatch { child, parent } => f.write_fmt(format_args!(
                "AVL child-parent height mismatch: {child:?}, {parent}"
            )),
        }
    }
}

impl<E: Error> Error for AvlError<E> {}

pub struct AvlNode<'a, Ctx: Context<'a>, A: Mentionable<'a, Ctx>> {
    l: AvlTree<'a, Ctx, A>,
    r: AvlTree<'a, Ctx, A>,
    key: A,
}

pub struct AvlTree<'a, Ctx: Context<'a>, A: Mentionable<'a, Ctx>> {
    node: Nullable<'a, Ctx, AvlNode<'a, Ctx, A>>,
    height: u64,
}

#[derive(Clone)]
pub struct AvlNodeFactory<F>(F);

#[derive(Clone)]
pub struct AvlTreeFactory<F>(NullableFactory<AvlNodeFactory<F>>);

impl<'a, Ctx: Context<'a>, A: Mentionable<'a, Ctx>> Serializable for AvlNode<'a, Ctx, A> {
    fn serialize(&self, serializer: &mut dyn Serializer) {
        self.l.serialize(serializer);
        self.r.serialize(serializer);
        self.key.serialize(serializer);
    }
}

impl<'a, Ctx: Context<'a>, A: Mentionable<'a, Ctx>> Serializable for AvlTree<'a, Ctx, A> {
    fn serialize(&self, serializer: &mut dyn Serializer) {
        self.height.serialize(serializer);
        self.node.serialize(serializer);
    }
}

impl<'a, Ctx: Context<'a>, A: Mentionable<'a, Ctx>> Mentionable<'a, Ctx> for AvlNode<'a, Ctx, A> {
    type Fctr = AvlNodeFactory<A::Fctr>;

    fn factory(&self) -> Self::Fctr {
        AvlNodeFactory(self.key.factory())
    }

    fn points_typed(&self, points: &mut impl TakesPoints<'a, Ctx>) {
        self.l.points_typed(points);
        self.r.points_typed(points);
        self.key.points_typed(points);
    }
}

impl<'a, Ctx: Context<'a>, A: Mentionable<'a, Ctx>> Mentionable<'a, Ctx> for AvlTree<'a, Ctx, A> {
    type Fctr = AvlTreeFactory<A::Fctr>;

    fn factory(&self) -> Self::Fctr {
        AvlTreeFactory(self.node.factory())
    }

    fn points_typed(&self, points: &mut impl TakesPoints<'a, Ctx>) {
        self.node.points_typed(points);
    }
}

fn balanced<E>(hl: u64, hr: u64) -> Result<(), AvlError<E>> {
    if std::cmp::max(hl, hr) - std::cmp::min(hl, hr) > 1 {
        return Err(AvlError::Balance(hl, hr));
    }
    Ok(())
}

impl<'a, Ctx: Context<'a>, F: Factory<'a, Ctx>> Factory<'a, Ctx> for AvlNodeFactory<F> {
    type Mtbl = AvlNode<'a, Ctx, F::Mtbl>;

    type ParseError = AvlError<F::ParseError>;

    fn deserialize(
        &self,
        deserializer: &mut dyn Deserializer,
        resolver: Rc<dyn Resolver<'a, Ctx>>,
        addresses: &mut Addresses,
    ) -> ParseResult<'a, Ctx, Self> {
        let tree_factory = AvlTreeFactory(NullableFactory::new(self.clone()));
        let l = tree_factory.deserialize(deserializer, resolver.clone(), addresses)?;
        let r = tree_factory.deserialize(deserializer, resolver.clone(), addresses)?;
        balanced(l.height, r.height)?;
        let key = self
            .0
            .deserialize(deserializer, resolver.clone(), addresses)
            .map_err(AvlError::Key)?;
        Ok(AvlNode { l, r, key })
    }

    fn unexpected_tail(&self, mentionable: Self::Mtbl, tail: &[u8]) -> Self::ParseError {
        AvlError::Key(self.0.unexpected_tail(mentionable.key, tail))
    }
}

impl<'a, Ctx: Context<'a>, F: Factory<'a, Ctx>> Factory<'a, Ctx> for AvlTreeFactory<F> {
    type Mtbl = AvlTree<'a, Ctx, F::Mtbl>;

    type ParseError = AvlError<F::ParseError>;

    fn deserialize(
        &self,
        deserializer: &mut dyn Deserializer,
        resolver: Rc<dyn Resolver<'a, Ctx>>,
        addresses: &mut Addresses,
    ) -> ParseResult<'a, Ctx, Self> {
        let node = self.0.deserialize(deserializer, resolver, addresses)?;
        let height = u64::a_deserialize(deserializer)?;
        if let Nullable::Null(_) = node {
            if height != 0 {
                return Err(AvlError::LeafHeight(height));
            }
        } else if height == 0 {
            return Err(AvlError::NodeHeight);
        }
        Ok(AvlTree { node, height })
    }

    fn unexpected_tail(&self, mentionable: Self::Mtbl, tail: &[u8]) -> Self::ParseError {
        u64::a_unexpected_tail(mentionable.height, tail).into()
    }
}