//! Extended type erasure support.
//!
//! The `AnyTrait` trait provides dynamic upcasting to trait objects.

pub mod indirect;
mod static_wrapper;
mod type_name_id;

use crate::hkt::{Invariant, Marker};
use core::marker::PhantomData;

pub use static_wrapper::*;
pub use type_name_id::*;

#[cfg(all(feature = "alloc", not(feature = "std")))]
use alloc::boxed::Box;

pub enum StaticType {}
pub enum NamedType {}
pub enum LifetimeType {}

#[allow(non_snake_case)]
pub mod TypeName {
    pub trait MemberTypeForLt<'a, 'ctx: 'a, B> {
        type T: ?Sized + LowerTypeWithBound<'a, 'ctx, &'a &'ctx (), Higher = Self>;
    }

    pub trait MemberType: 'static + for<'a, 'ctx> MemberTypeForLt<'a, 'ctx, &'a &'ctx ()> {}

    impl<T: ?Sized> MemberType for T where
        T: 'static + for<'a, 'ctx> MemberTypeForLt<'a, 'ctx, &'a &'ctx ()>
    {
    }

    pub trait LowerTypeWithBound<'a, 'ctx: 'a, B>: 'a + Send + Sync {
        type Higher: ?Sized + MemberTypeForLt<'a, 'ctx, &'a &'ctx (), T = Self> + MemberType;
    }

    pub trait LowerType<'a, 'ctx: 'a>: LowerTypeWithBound<'a, 'ctx, &'a &'ctx ()> {}

    impl<'a, 'ctx: 'a, T: ?Sized> LowerType<'a, 'ctx> for T where
        T: LowerTypeWithBound<'a, 'ctx, &'a &'ctx ()>
    {
    }

    pub type T<'a, 'ctx, __> = <__ as MemberTypeForLt<'a, 'ctx, &'a &'ctx ()>>::T;
    pub type HigherRanked<'a, 'ctx, __> =
        <__ as LowerTypeWithBound<'a, 'ctx, &'a &'ctx ()>>::Higher;
}

pub struct RefHrt<T: ?Sized>(Marker<T>);

impl<'a, 'ctx, T: ?Sized> TypeName::MemberTypeForLt<'a, 'ctx, &'a &'ctx ()> for RefHrt<T>
where
    T: TypeName::MemberType,
{
    type T = &'a TypeName::T<'a, 'ctx, T>;
}

impl<'a, 'ctx, T: ?Sized> TypeName::LowerTypeWithBound<'a, 'ctx, &'a &'ctx ()> for &'a T
where
    T: TypeName::LowerTypeWithBound<'a, 'ctx, &'a &'ctx ()>,
{
    type Higher = RefHrt<TypeName::HigherRanked<'a, 'ctx, T>>;
}

pub struct MutHrt<T: ?Sized>(Marker<T>);

impl<'a, 'ctx, T: ?Sized> TypeName::MemberTypeForLt<'a, 'ctx, &'a &'ctx ()> for MutHrt<T>
where
    T: TypeName::MemberType,
{
    type T = &'a mut TypeName::T<'a, 'ctx, T>;
}

impl<'a, 'ctx, T: ?Sized> TypeName::LowerTypeWithBound<'a, 'ctx, &'a &'ctx ()> for &'a mut T
where
    T: TypeName::LowerTypeWithBound<'a, 'ctx, &'a &'ctx ()>,
{
    type Higher = MutHrt<TypeName::HigherRanked<'a, 'ctx, T>>;
}

#[cfg(feature = "alloc")]
impl<'a, 'ctx, T: ?Sized> TypeName::MemberTypeForLt<'a, 'ctx, &'a &'ctx ()> for Box<T>
where
    T: TypeName::MemberType,
{
    type T = Box<TypeName::T<'a, 'ctx, T>>;
}

#[cfg(feature = "alloc")]
impl<'a, 'ctx, T: ?Sized> TypeName::LowerTypeWithBound<'a, 'ctx, &'a &'ctx ()> for Box<T>
where
    T: TypeName::LowerTypeWithBound<'a, 'ctx, &'a &'ctx ()>,
{
    type Higher = Box<TypeName::HigherRanked<'a, 'ctx, T>>;
}

/// Dynamic trait lookup.
///
/// This trait allows looking up the trait object form of `self` for a
/// given trait object `id`. This is similar to upcasting to the trait given
/// by `id` if [`AnyTrait`] had every trait as a super bound.
///
// ```
// use treaty::any::{AnyTrait, any_trait, TypeName};
// use treaty::hkt::higher_ranked_type;
//
// // Create a test value.
// let my_num = MyNum(42);
//
// // Cast to be a AnyTrait trait object.
// // Now we don't know the type.
// let anything: &(dyn AnyTrait<'_> + Send + Sync) = &my_num;
//
// // We can still upcast to an impl of ToNum.
// let to_num_object: &dyn ToNum = anything.upcast::<DynToNum>().unwrap();
//
// assert_eq!(to_num_object.num(), 42);
//
// // === Type Setup ===
//
// // An example trait.
// trait ToNum {
//     fn num(&self) -> i32;
// }
//
// enum DynToNum {}
//
// higher_ranked_type! {
//     impl TypeName {
//         impl['a, 'ctx] type T['a, 'ctx] for DynToNum =
//             dyn ToNum + Send + Sync + 'a;
//
//         impl['a, 'ctx] type HigherRanked['a, 'ctx] for dyn ToNum + Send + Sync + 'a =
//             DynToNum;
//     }
// }
//
// // An example struct.
// struct MyNum(i32);
//
// // The example struct impls the example trait.
// impl ToNum for MyNum {
//     fn num(&self) -> i32 {
//         self.0
//     }
// }
//
// // Allow the example struct's trait impls to be looked up at runtime.
// // Here the only trait that can be looked up is ToNum as its the only
// // one in the list.
// any_trait! {
//     impl['ctx] MyNum = [DynToNum]
// }
// ```
pub trait AnyTrait<'ctx> {
    /// Upcast a borrow to the given trait object.
    ///
    /// Use the `<dyn AnyTrait>::upcast()` helper method instead, if possible.
    ///
    /// If `self` doesn't support upcasting to the requested type
    /// then `None` is returned. The returned trait object is type erased so this trait
    /// is object safe.
    fn upcast_to_id<'a>(
        &'a self,
        id: TypeNameId,
    ) -> Option<AnyTraitObject<'a, 'ctx, indirect::Ref>>
    where
        'ctx: 'a;

    /// Upcast a mutable borrow to the given trait object.
    ///
    /// Use the `<dyn AnyTrait>::upcast_mut()` helper method instead, if possible.
    ///
    /// If `self` doesn't support upcasting to the requested type
    /// then `None` is returned. The returned trait object is type erased so this trait
    /// is object safe.
    fn upcast_to_id_mut<'a>(
        &'a mut self,
        id: TypeNameId,
    ) -> Option<AnyTraitObject<'a, 'ctx, indirect::Mut>>
    where
        'ctx: 'a;
}

impl<'b, 'ctx: 'b> dyn AnyTrait<'ctx> + Send + Sync + 'b {
    /// Upcast a borrow to the given trait object type.
    ///
    /// This should be used instead of [`upcast_to_id`][AnyTrait::upcast_to_id]
    /// as it automatically downcasts the returned [`AnyTraitObject`].
    ///
    /// If the returned [`AnyTraitObject`] is the wrong type, then a panic happens.
    #[inline(always)]
    pub fn upcast<'a, Trait: ?Sized + TypeName::MemberType>(
        &'a self,
    ) -> Option<&'a TypeName::T<'a, 'ctx, Trait>> {
        self.upcast_to_id(TypeNameId::of::<Trait>())
            .map(|object| match object.downcast() {
                Ok(object) => object,
                Err(object) => panic!(
                    "Unexpected trait object. This means a bad impl of \
                    `upcast_to_id`. Expected: {:?}, Got {:?}",
                    TypeNameId::of::<Trait>(),
                    object.id()
                ),
            })
    }

    /// Upcast a mutable borrow to the given trait object type.
    ///
    /// This should be used instead of [`upcast_to_id_mut`][AnyTrait::upcast_to_id]
    /// as it automatically downcasts the returned [`AnyTraitObject`].
    ///
    /// If the returned [`AnyTraitObject`] is the wrong type, then a panic happens.
    #[inline(always)]
    pub fn upcast_mut<'a, Trait: ?Sized + TypeName::MemberType>(
        &'a mut self,
    ) -> Option<&'a mut TypeName::T<'a, 'ctx, Trait>> {
        self.upcast_to_id_mut(TypeNameId::of::<Trait>())
            .map(|object| match object.downcast() {
                Ok(object) => object,
                Err(object) => panic!(
                    "Unexpected trait object. This means a bad impl of \
                    `upcast_to_id_mut`. Expected: {:?}, Got {:?}",
                    TypeNameId::of::<Trait>(),
                    object.id()
                ),
            })
    }
}

/// Implement [`AnyTrait`] for a type.
///
/// This allows looking up trait objects from the provided list.
/// See [`AnyTrait`] for an example.
#[doc(hidden)]
#[macro_export]
macro_rules! any_trait {
    {
        impl[$lt:lifetime $($generic:tt)*] $name:ty = [$($protocol:ty),* $(,)?]
        else ref {
            let ($this:ident, $id:ident);
            $($fallback:tt)*
        }
        else mut {
            let ($mut_this:ident, $mut_id:ident);
            $($mut_fallback:tt)*
        }
        $(where $($bound:tt)*)?
    } => {
        impl<$lt $($generic)*> $crate::any::AnyTrait<$lt> for $name
        $(where $($bound)*)?
        {
            #[inline(always)]
            fn upcast_to_id<'__>(
                &'__ self,
                id: $crate::any::TypeNameId
            ) -> ::core::option::Option<$crate::any::AnyTraitObject<'__, $lt, $crate::any::indirect::Ref>>
            where
                $lt: '__
            {
                // This match should be optimized well by llvm.
                match id {
                    $(id if id == $crate::any::TypeNameId::of::<$protocol>()
                        => ::core::option::Option::Some($crate::any::AnyTraitObject::<'__, $lt, _>::new::<
                            $crate::any::TypeName::T<'__, $lt, $protocol>
                        >(self as _)),)*
                    $id => {
                        let $this = self;
                        $($fallback)*
                    }
                }
            }

            #[inline(always)]
            fn upcast_to_id_mut<'__>(
                &'__ mut self,
                id: $crate::any::TypeNameId
            ) -> ::core::option::Option<$crate::any::AnyTraitObject<'__, $lt, $crate::any::indirect::Mut>>
            where
                $lt: '__
            {
                // This match should be optimized well by llvm.
                match id {
                    $(id if id == $crate::any::TypeNameId::of::<$protocol>()
                        => ::core::option::Option::Some($crate::any::AnyTraitObject::<'__, $lt, _>::new::<
                            $crate::any::TypeName::T<'__, $lt, $protocol>
                        >(self as _)),)*
                    $mut_id => {
                        let $mut_this = self;
                        $($mut_fallback)*
                    }
                }
            }
        }
    };
    {
        impl[$lt:lifetime $($generic:tt)*] $name:ty = [$($protocol:ty),* $(,)?]
        $(where $($bound:tt)*)?
    } => {
        $crate::any::any_trait! {
            impl[$lt $($generic)*] $name = [$($protocol),*]
            else ref {
                // Always answer no in the fallback branch if no fallback was given.
                let (_this, _id);
                ::core::option::Option::None
            } else mut {
                // Always answer no in the fallback branch if no fallback was given.
                let (_this, _id);
                ::core::option::Option::None
            } $(where $($bound)*)?
        }
    }
}
#[doc(inline)]
pub use any_trait;

use self::indirect::{sealed::RawIndirect, Indirect};

/// A double fat pointer.
///
/// This struct wraps a possibly fat pointer of type described by `I`, and adds
/// an additional vtable to allow downcasting to a specific fat pointer.
/// This type is similar to if `&dyn Any` was allowed to itself store unsized types like trait
/// objects.
///
/// The `'a` lifetime is the lifetime of the pointer, and the `'ctx` lifetime is a context lifetime
/// the inner type can use. This type is always invariant over both lifetimes.
///
/// `&'a dyn MyTrait<ctx>` becomes `AnyTraitObject<'a, 'ctx, Ref>`.
///
/// The `I` generic is the flavor if pointer being used. It can be [`Ref`][indirect::Ref] or [`Mut`][indirect::Mut].
#[must_use]
pub struct AnyTraitObject<'a, 'ctx: 'a, I: Indirect<'a>> {
    /// The extra vtable pointer.
    ///
    /// The TypeNameId gives the TypeId of the T's type name.
    /// This is unique per T minus the 'a and 'ctx lifetimes.
    /// which means a `dyn Trait<'ctx> + 'a` can have a TypeNameId.
    ///
    /// The unsafe function is the drop impl for the pointer.
    /// It must only be called once per RawIndirect value, and the value must not be used after the
    /// call. Only a RawIndirect of the correct I type must be passed.
    info: fn() -> TypeNameId,

    /// The inner pointer value.
    ///
    /// This is some form of &T where T may be sized or not.
    indirect: RawIndirect<'a, I>,

    _lifetime: Invariant<'ctx>,
    _not_send_sync: PhantomData<*const ()>,
}

impl<'a, 'ctx, I: Indirect<'a>> AnyTraitObject<'a, 'ctx, I> {
    /// Type erase a pointer.
    ///
    /// `T` doesn't need to be [`Sized`]. As such, a fat pointer can be passed to this function.
    pub fn new<T: ?Sized + TypeName::LowerType<'a, 'ctx>>(indirect: I::ForT<T>) -> Self {
        Self {
            info: TypeNameId::of_lower::<T>,
            indirect: RawIndirect::new(indirect),
            _lifetime: Default::default(),
            _not_send_sync: PhantomData,
        }
    }

    /// Downcast to an indirection with a given `T` type.
    ///
    /// If the type of the stored value is different, then `self` is
    /// returned as is.
    pub fn downcast<T: ?Sized + TypeName::LowerType<'a, 'ctx>>(self) -> Result<I::ForT<T>, Self> {
        if self.id() == TypeNameId::of_lower::<T>() {
            // SAFETY: We know that the type name type is unique per T because it is bijective.
            // A self is only made in Self::new where the info is taken from T.
            // If the check above passes then we know T must be the same minus the lifetimes.
            // We know the lifetime 'ctx is correct because Self is invariant over it.
            // RawIndirect makes sure that the 'a is correct by being invariant over it.
            //
            // See the tests at the bottom of the file for a proof that the type name is bijective
            // to T.
            Ok(unsafe { self.indirect.into_inner::<T>() })
        } else {
            Err(self)
        }
    }

    /// Type ID of the stored value's `T`.
    pub fn id(&self) -> TypeNameId {
        (self.info)()
    }
}

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

    // #[test]
    // fn any_trait_macro_implements_the_trait() {
    //     trait Z<'ctx> {
    //         fn get(&self) -> i32;
    //     }
    //
    //     struct DynZ;
    //
    //     higher_ranked_type! {
    //         impl TypeName {
    //             impl['a, 'ctx] type T['a, 'ctx] for DynZ =
    //                 dyn Z<'ctx> + Send + Sync + 'a;
    //
    //             impl['a, 'ctx] type HigherRanked['a, 'ctx] for dyn Z<'ctx> + Send + Sync + 'a =
    //                 DynZ;
    //         }
    //     }
    //
    //     struct X<'ctx>(&'ctx i32);
    //
    //     impl<'ctx> Z<'ctx> for X<'ctx> {
    //         fn get(&self) -> i32 {
    //             *self.0
    //         }
    //     }
    //
    //     any_trait! {
    //         impl['ctx] X<'ctx> = [
    //             DynZ
    //         ]
    //     }
    //
    //     let z = 42;
    //     let x = X(&z);
    //     let y = (&x as &(dyn AnyTrait<'_> + Send + Sync))
    //         .upcast::<DynZ>()
    //         .unwrap();
    //     assert_eq!(y.get(), 42);
    // }

    // The following proves that the higher ranked types are bijective using the type system.
    //
    // We have the type tower: T<'a, 'ctx> <-> DynT<'ctx> <-> NameT
    // We want every T, DynT, NameT set to be unique.
    //
    // Assume there was a U that tried to use NameT in it's type tower:
    // U<'a, 'ctx> <-> DynU<'ctx> <-> NameT
    //
    // If we traverse the type tower in this order: T -r-> A -r-> B -l-> C -l-> D
    // where -r-> is a raise and -l-> is a lower, then if D is always T we know that no sequence
    // U -r-> A2 -r-> B -l-> C2 -l-> D can exist exept where T == U. This is because B cannot
    // have information about where it came from and still be the same B in the type system.
    // The following makes sure that a U could never become a T by the raise then lower process.

    // This proves that the bijective type names are really bijective.
    fn _is_bijective_raise<'a, 'ctx: 'a, T>(
        x: &TypeName::T<'a, 'ctx, TypeName::HigherRanked<'a, 'ctx, T>>,
    ) where
        T: TypeName::LowerType<'a, 'ctx>,
    {
        // If C -> B -> A -> B -> C (shown by this assignment), then C and A must be bijective.
        let _y: &T = x;
    }

    // This proves that the bijective type names are really bijective.
    fn _is_bijective_lower<'a, 'ctx: 'a, U>(
        x: &TypeName::HigherRanked<'a, 'ctx, TypeName::T<'a, 'ctx, U>>,
    ) where
        U: TypeName::MemberType,
    {
        // If A -> B -> C -> B -> A (shown by this assignment), then A and C must be bijective.
        let _y: &U = x;
    }

    // fn _is_bijective_raise2<'a, 'ctx: 'a, 'b, 'c: 'b, T>(
    //     x: &TypeName::T<'b, 'c, TypeName::HigherRanked<'a, 'ctx, T>>,
    // ) where
    //     T: TypeName::LowerType<'a, 'ctx>,
    // {
    //     // If C -> B -> A -> B -> C (shown by this assignment), then C and A must be bijective.
    //     let _y: &T = x;
    // }
}