#![allow(stable_features)]
#![feature(if_let_guard)]
#![feature(let_chains)]
#![feature(never_type)]
#![feature(iterator_try_collect)]
#![feature(try_trait_v2)]
#![feature(try_trait_v2_residual)]

//! Provides a derive macro for [Try] & optionally [Try_ConvertResult] for interconversion with
//! `std::result::Result` and [Try_Iterator] for iterating over `IntoIterator` and collecting from
//! `FromIterator` analog to how `Result` & `Option` do this.
//! See ([try_trait_v2](https://rust-lang.github.io/rfcs/3058-try-trait-v2.html)) for more details
//! of the underlying trait.
//!
//! ## Requires:
//!
//!   - nightly
//!   - `#![feature(never_type)]`
//!   - `#![feature(try_trait_v2)]`
//!   - `#![feature(try_trait_v2_residual)]`
//!   - optionally: `#![feature(iterator_try_collect)]` (if using Try_Iterator)
//!
//! ## Limitations on the annotated type:
//!
//!   - must be an `enum`
//!   - must have _at least one_ generic type
//!   - the _first_ generic type must be the `Output` type (produced when not short circuiting)
//!   - the output variant (does not short-circuit) must be the _first_ variant and store the output
//!     type as the _only unnamed_ field
//!
//! See the individual documentation for [Try], [Try_ConvertResult] and [Try_Iterator] for specifics
//! on the generated code.
//!
//! ## Example Usage:
//!
//! ```rust
//! #![feature(never_type)]
//! #![feature(try_trait_v2)]
//! #![feature(try_trait_v2_residual)]
//! use try_v2::{Try, Try_ConvertResult};
//!
//! #[derive(Try, Try_ConvertResult)]
//! enum TestResult<T> {
//!     Ok(T),
//!     TestsFailed,
//!     OtherError(String)
//! }
//!
//! // Basic short circuiting thanks to `#[derive(Try)]`
//! fn run_tests() -> TestResult<()> {
//!     TestResult::OtherError("oops!".to_string())?; // <- Function short-circuits here ...
//!     TestResult::TestsFailed?;
//!     TestResult::Ok(())
//! }
//!
//! assert!(matches!(run_tests(), TestResult::OtherError(msg) if msg == "oops!"));
//!
//!
//! // Conversion from std::result::Result thanks to `#[derive(Try_ConvertResult)]`
//! struct TestFailure {}
//!
//! impl<T> From<TestFailure> for TestResult<T> {
//!     fn from(err: TestFailure) -> Self {
//!         TestResult::TestsFailed
//!     }
//! }
//!
//! fn run_more_tests() -> TestResult<()> {
//!     std::result::Result::Err(TestFailure{})?; // <- Function short-circuits here & converts to a TestResult...
//!     TestResult::Ok(())
//! }
//!
//! assert!(matches!(run_more_tests(), TestResult::TestsFailed));
//! ```
//!
//! ## MSRV
//!
//! 1.85.1, in case you are using a fixed version of nightly just to get access to specific unstable features.
//!
//! ## Currently untested (may work, may not ...):
//!
//!   - `where` clauses
//!   - storing `Fn`s in variants

use proc_macro::TokenStream as TokenStream1;
use proc_macro2::TokenStream as TokenStream2;
use proc_macro2_diagnostic::prelude::*;
use quote::{format_ident, quote};
use syn::{DeriveInput, GenericParam, parse_quote, spanned::Spanned};

mod parse;
use parse::TryEnum;

#[proc_macro_derive(Try)]
/// Derives [try_trait_v2](https://rust-lang.github.io/rfcs/3058-try-trait-v2.html)
///
/// See the [crate level documentation](crate) for restrictions and detailed examples
///
/// ## Derived code
/// ```
/// # #![feature(never_type)]
/// # #![feature(try_trait_v2)]
/// # #![feature(try_trait_v2_residual)]
/// # use try_v2::Try;
/// #[derive(Try)]
/// enum TestResult<T, E> {
///     Ok(T),
///     TestsFailed,
///     OtherError(E)
/// }
/// ```
/// will result in code of the shape:
/// ```ignore
/// impl<T,E> Try for TestResult<T, E> {
///     type Output = T;
///     type Residual = TestResult<!,E>;
///
///     fn from_output(output: T) -> Self {
///         Self::Ok(output)
///     }
///
///     fn branch(self) -> ControlFlow<Self::Residual, Self::Output> {
///         Self::Ok(t) => Continue(t),
///         ... each failing variant => Break(failing variant) ...   
///     }
/// }
///
/// impl<T, E> FromResidual<TestResult<!,E>> for TestResult<T, E> {
///     fn from_residual(residual: TestResult<!,E>) -> Self {
///         match residual {
///             ... each failing variant => itself ...              
///         }
///     }
/// }
///
/// impl<T, E> Residual<T> for TestResult<!, E> {
///     type TryType = TestResult<T, E>;
/// }
/// ```
///
/// ## Things to note
///
/// This macro aims to reduce boilerpate for the most common implementations.
///
/// ### Residual
/// The `Residual` is generated by replacing the _first generic type_ with `!`. This means
///     - we rely on the unstable `#![feature(never_type)]` (on the assumption that it will be
///         stabilised on a similar timeframe to `#![feature(try_trait_v2)]`
///     - any short-circuiting variants which also store this type will have a "hole"
///     - the **invariant** _first generic type is stored by first variant_ is **strictly** enforced.
///         This is [by design](https://en.wikipedia.org/wiki/Poka-yoke) to guard against
///         accidental usage errors.
///
/// ### Output / Short-Circuiting
/// The _first variant_ is considered to represent the output case. This means
///     - there is only one continue (non-short-circuiting) case, all other variants will
///         short-circuit on `?`
///     - the first variant **must** store a generic type
///     - the **invariant** _first generic type is stored by first variant_ is **strictly** enforced.
///         This is [by design](https://en.wikipedia.org/wiki/Poka-yoke) to guard against
///         accidental usage errors.
///
/// ### Borrowed data
/// It is possible to derive Try for enums which store references to data, if your use case makes
/// this valuable.
///
/// ```
/// # #![feature(never_type)]
/// # #![feature(try_trait_v2)]
/// # #![feature(try_trait_v2_residual)]
/// # use try_v2::Try;
/// #[derive(Try)]
/// enum TestResult<'t, 'e, T, E> {
///     Ok(&'t T),
///     TestsFailed,
///     OtherError(&'e E)
/// }
/// ```
/// will have
/// ```ignore
///     type Output = &'t T;
///     type Residual = TestResult<'t, 'e, !, E>
/// ```
/// which _still includes all lifetimes_ in the Residual.
///
/// This is important to note if you are writing your own implementations of `FromResidual`.
///
/// Hint for matching on the Residual: `&!` itself will not be recognised as uninhabited by the
/// compiler but `&!` will dereference to `!` which will then coerce into any type or satisfy
/// `match ! {}`. Therefore, you should include a match arm `Ok(never) => *never` (doesn't guarantee
/// it's actually `&!`) or `Ok(&never) => match never {}` (more verbose but guarantees infallibility)
pub fn try_trait_v2_derive(input: TokenStream1) -> TokenStream1 {
    impl_derive(input.into()).into()
}

/// Parses & validates the input then quote!s the impl.  
fn impl_derive(input: TokenStream2) -> DiagnosticStream {
    let ast: DeriveInput = syn::parse2(input).expect("derive macro");

    let tryenum = TryEnum::parse(&ast)?;
    let (
        name,
        output_variant_name,
        output_type,
        _,
        residual_type,
        impl_generics,
        ty_generics,
        where_clause,
    ) = tryenum.split_for_impl();

    if !ast
        .attrs
        .iter()
        .any(|attr| attr.meta.path().is_ident("must_use"))
    {
        warn_spanned(
            (),
            ast.span(),
            "it is recommended to annotate try-types as `#[must_use]`",
        )?
    };

    let (branch_arms, residual_arms) = tryenum.generate_arms();

    let impl_try = quote! {
        impl #impl_generics std::ops::Try for #name #ty_generics #where_clause {
            type Output = #output_type;

            type Residual = #residual_type;

            #[inline]
            fn from_output(output: Self::Output) -> Self {
                Self::#output_variant_name(output)
            }

            #[inline]
            fn branch(self) -> std::ops::ControlFlow<Self::Residual, Self::Output> {
                match self {
                    #(#branch_arms)*
                }
            }
        }

        impl #impl_generics std::ops::FromResidual<#residual_type> for #name #ty_generics #where_clause {
            #[inline]
            #[track_caller]
            fn from_residual(residual: #residual_type) -> Self {
                match residual {
                    #(#residual_arms)*
                }
            }
        }

        impl #impl_generics std::ops::Residual<#output_type> for #residual_type #where_clause {
            type TryType = #name #ty_generics;
        }
    };
    Ok(impl_try)
}

#[proc_macro_derive(Try_ConvertResult)]
/// Derives ?-conversion from Result<T, E> and back where suitable implementations of From/Into exist.
///
/// ## Conversion from Result
/// For ?-conversion _from_ a `Result<T, SomeError>` `impl<T> From<SomeError> for MyTryEnum<T>`.
///
/// This will allow `?` on a call which returns `Result<T, SomeError>` in any function which returns
/// `MyTryEnum<T>`.
///
/// Type-hinting, type aliasing Result may be needed unless you have a
/// blanket From<E: Error> implementation.
///
/// ## Conversion to Result
/// For ?-conversion _to_ a `Result<_, SomeError>` `impl From<MyTryEnumResidual> for SomeError`.
///
/// Note that conversion must be defined between your _Residual_ and SomeError - this both avoids
/// triggering the orphan rule when your enum stores third-party / std types and requires
/// consideration and correct handling of each failure variant.
///
/// See the notes on [Try] for full details on identifying the correct Residual to use.
///
/// ## Derived Code
/// ```
/// # #![feature(never_type)]
/// # #![feature(try_trait_v2)]
/// # #![feature(try_trait_v2_residual)]
/// # use try_v2::{Try, Try_ConvertResult};
/// #[derive(Try, Try_ConvertResult)]
/// #[must_use]
/// enum TestResult<T, E> {
///     Ok(T),
///     TestsFailed,
///     OtherError(E)
/// }
/// ```
/// will generate:
/// ```ignore code-snippet
/// impl<T, E, RE> FromResidual<Result<Infallible, RE>> for TestResult<T, E>
/// where
///     RE: Into<TestResult<!,E>>
///
/// ... which calls Result::Err(e) => e.into(), ...
/// ```
/// and
/// ```ignore code-snippet
/// impl<E, RT, RE> FromResidual<TestResult<!,E>> for Result<RT, RE>
/// where
///     RE: From<TestResult<!,E>>
///
/// ... which calls Result::Err(residual.into()) ...
/// ```
///
/// ## Implementing [TryFrom]
///
/// TryFrom requires a [Result] to be returned. To handle this: use your residual
/// (e.g. `TestResult<!,E>` in the above example, or Eightball<!> in the one below) as the
/// `Error` type. Here's the example from the integration tests:
///
/// ```
/// #![feature(never_type)]
/// #![feature(try_trait_v2)]
/// #![feature(try_trait_v2_residual)]
///
/// use try_v2::{Try, Try_ConvertResult};
///
/// #[derive(Try, Try_ConvertResult)]
/// #[must_use]
/// enum Eightball<Y> {
///     Yes(Y),
///     No,
/// }
///
/// struct Even(i32);
///
/// impl TryFrom<i32> for Even {
///     type Error = Eightball<!>;
///
///     fn try_from(num: i32) -> Result<Even, Eightball<!>> {
///         if num % 2 == 0 {
///             Result::Ok(Even(num))
///         } else {
///             Result::Err(Eightball::No)
///         }
///     }
/// }
///
/// fn even_string(num: i32) -> Eightball<String> {
///     let n = Even::try_from(num)?;
///     let s = format!("{}", n.0);
///     Eightball::Yes(s)
/// }
///
/// assert!(matches!(even_string(2), Eightball::Yes(s) if s == "2"));
/// assert!(matches!(even_string(1), Eightball::No));
/// ```
///
pub fn try_trait_v2_convert_result(input: TokenStream1) -> TokenStream1 {
    impl_convert_result(input.into()).into()
}

fn impl_convert_result(input: TokenStream2) -> DiagnosticStream {
    let ast: DeriveInput = syn::parse2(input).expect("derive macro");

    let tryenum = TryEnum::parse(&ast)?;
    let (name, _, _, output_type_name, residual_type, _, ty_generics, where_clause) =
        tryenum.split_for_impl();

    let result_e = format_ident!("Derive_TryConvert_ResultE");
    let result_t = format_ident!("Derive_TryConvert_ResultT");

    let from_result_generics = tryenum.generics(|g| {
        g.params
            .push(parse_quote! {#result_e: Into<#residual_type>})
    });
    let (from_result_impl_generics, _, _) = from_result_generics.split_for_impl();

    let mut impl_convert = quote! {
        impl #from_result_impl_generics std::ops::FromResidual<std::result::Result<std::convert::Infallible, #result_e>> for #name #ty_generics #where_clause
        {
            #[inline]
            #[track_caller]
            fn from_residual(residual: std::result::Result<std::convert::Infallible, #result_e>) -> Self {
                match residual {
                    std::result::Result::Err(e) => {
                        let bang: #residual_type = e.into();
                        Self::from_residual(bang)
                    }
                }
            }
        }
    };

    let to_result_generics = tryenum.generics_with_params(|p| {
        p
            //remove output type
            .filter(|p| !matches!(p, GenericParam::Type(t) if t.ident == *output_type_name))
            // add result types
            .chain([
                parse_quote! {#result_t},
                parse_quote! {#result_e: From<#residual_type>},
            ])
    });

    let (to_result_impl_generics, _, _) = to_result_generics.split_for_impl();

    impl_convert.extend(quote! {
        impl #to_result_impl_generics std::ops::FromResidual<#residual_type> for std::result::Result<#result_t, #result_e>
        {
            #[inline]
            #[track_caller]
            fn from_residual(residual: #residual_type) -> Self {
                std::result::Result::Err(residual.into())
            }
        }
    });
    Ok(impl_convert)
}

#[proc_macro_derive(Try_Iterator)]
/// Derives `IntoIterator` and `FromIterator` analog to `Result` & `Option`.
///
/// - Vec<TryEnum>::collect() -> TryEnum<Vec>.
/// - TryEnum.into_iter() -> yields _one_ value if Ok, else empty.
///
/// ## Example
/// ```
/// # #![feature(never_type)]
/// # #![feature(try_trait_v2)]
/// # #![feature(try_trait_v2_residual)]
/// # #![feature(iterator_try_collect)]
/// # use try_v2::{Try, Try_Iterator};
/// # use TestResult::{Ok, TestsFailed, OtherError};
/// #[derive(Try, Try_Iterator)]
/// #[must_use]
/// enum TestResult<T, E> {
///     Ok(T),
///     TestsFailed,
///     OtherError(E),
/// }
///
/// # fn main() {
/// let tests: Vec<TestResult<i32, &'static str>> = vec![Ok(1), TestsFailed, Ok(2), OtherError("something wierd"), Ok(3), Ok(4)];
///
/// let first_results: TestResult<Vec<i32>, &'static str> = tests.into_iter().collect();
/// assert!(matches!(first_results, TestsFailed));
///
/// let mut test: TestResult<i32, &'static str> = Ok(4);
/// let borrowed_result: &i32 = test.iter().next().unwrap();
/// assert_eq!(borrowed_result, &4);
/// match test.iter_mut().next() {
///     Some(v) => *v = 5,
///     None => {},
/// }
/// assert!(matches!(test, TestResult::Ok(v) if v == 5));
/// let result = test.into_iter().next();
/// assert_eq!(result, Some(5));
/// # }
/// ```
pub fn iterator_traits(input: TokenStream1) -> TokenStream1 {
    impl_iterator_traits(input.into()).into()
}

fn impl_iterator_traits(input: TokenStream2) -> DiagnosticStream {
    let ast: DeriveInput = syn::parse2(input).expect("derive macro");

    let tryenum = TryEnum::parse(&ast)?;
    let (
        name,
        output_variant_name,
        output_type,
        output_type_name,
        _,
        impl_generics,
        ty_generics,
        where_clause,
    ) = tryenum.split_for_impl();

    // Standing on the shoulders of giants & blatanty (ab)using `std::option`'s work
    let mut impl_traits = quote! {
        impl #impl_generics std::iter::IntoIterator for #name #ty_generics #where_clause {
            type Item = #output_type;
            type IntoIter = std::option::IntoIter<#output_type>;


            fn into_iter(self) -> Self::IntoIter {
                let opt = match self {
                    Self::#output_variant_name(v) => Some(v),
                    _ => None,
                };
                opt.into_iter()
            }
        }

        impl #impl_generics #name #ty_generics {
            pub fn iter(&self) -> std::option::IntoIter<&#output_type> {
                let opt = match self {
                    Self::#output_variant_name(v) => Some(v),
                    _ => None,
                };
                opt.into_iter()
            }

            pub fn iter_mut(&mut self) -> std::option::IntoIter<&mut #output_type> {
                let opt = match self {
                    Self::#output_variant_name(v) => Some(v),
                    _ => None,
                };
                opt.into_iter()
            }
        }
    };

    let defined_type = quote! {#name #ty_generics};
    let vec_ish = format_ident!("Derive_TryIterator_V");

    let full_generics = tryenum.generics(|g| {
        g.params
            .push(parse_quote! {#vec_ish: FromIterator<#output_type>})
    });

    let (full_impl_generics, _, full_where_clause) = full_generics.split_for_impl();

    let returned_generics = tryenum.generics(|g| {
        for param in g.type_params_mut() {
            if param.ident == *output_type_name {
                *param = parse_quote! {#vec_ish};
                break;
            }
        }
    });
    let (_, ret_ty_generics, _) = returned_generics.split_for_impl();

    impl_traits.extend(quote! {
        impl #full_impl_generics std::iter::FromIterator<#defined_type> for #name #ret_ty_generics #full_where_clause
        {
            fn from_iter<I: IntoIterator<Item=#defined_type>>(iter: I) -> Self {
                iter.into_iter().try_collect()
            }
        }
    });

    Ok(impl_traits)
}

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

    #[test]
    fn derive() {
        let original: TokenStream2 = quote! {
            #[derive(Try)]
            #[must_use]
            enum Exit<T: Termination> {
                Ok(T),
                TestsFailed,
                OtherError(String),
                NamedError{err: String, text: String},
            }
        };

        let derived_impl: TokenStream2 = quote! {
            impl<T: Termination> std::ops::Try for Exit<T> {
                type Output = T;

                type Residual = Exit<!>;

                #[inline]
                fn from_output(output: Self::Output) -> Self {
                    Self::Ok(output)
                }

                #[inline]
                fn branch(self) -> std::ops::ControlFlow<Self::Residual, Self::Output> {
                    match self {
                        Self::Ok(v0) => std::ops::ControlFlow::Continue(v0),
                        Self::TestsFailed => std::ops::ControlFlow::Break(Exit::TestsFailed),
                        Self::OtherError(v0) => std::ops::ControlFlow::Break(Exit::OtherError(v0)),
                        Self::NamedError{err, text} => std::ops::ControlFlow::Break(Exit::NamedError{err, text}),
                    }
                }
            }

            impl<T: Termination> std::ops::FromResidual<Exit<!> > for Exit<T> {
                #[inline]
                #[track_caller]
                fn from_residual(residual: Exit<!>) -> Self {
                    match residual {
                        Exit::TestsFailed => Exit::TestsFailed,
                        Exit::OtherError(v0) => Exit::OtherError(v0),
                        Exit::NamedError{err, text} => Exit::NamedError{err, text},
                    }
                }
            }

            impl<T: Termination> std::ops::Residual<T> for Exit<!> {
                type TryType = Exit<T>;
            }
        };
        assert_eq!(
            derived_impl.to_string(),
            impl_derive(original).unwrap().to_string()
        )
    }
    #[test]
    fn convert_result() {
        let original: TokenStream2 = quote! {
            #[derive(Try_ConvertResult)]
            enum Exit<T: Termination, E> {
                Ok(T),
                TestsFailed,
                OtherError(E),
            }
        };

        let expected_impl: TokenStream2 = quote! {
            impl<T: Termination, E, Derive_TryConvert_ResultE: Into< Exit<!, E> > > std::ops::FromResidual<std::result::Result<std::convert::Infallible, Derive_TryConvert_ResultE>> for Exit<T, E>
            {
                #[inline]
                #[track_caller]
                fn from_residual(residual: std::result::Result<std::convert::Infallible, Derive_TryConvert_ResultE>) -> Self {
                    match residual {
                        std::result::Result::Err(e) => {
                            let bang: Exit<!, E> = e.into();
                            Self::from_residual(bang)
                        }
                    }
                }
            }

            impl<E, Derive_TryConvert_ResultT, Derive_TryConvert_ResultE: From<Exit<!, E> > > std::ops::FromResidual<Exit<!, E> > for std::result::Result<Derive_TryConvert_ResultT, Derive_TryConvert_ResultE>
            {
                #[inline]
                #[track_caller]
                fn from_residual(residual: Exit<!, E>) -> Self {
                    std::result::Result::Err(residual.into())
                }
            }
        };

        assert_eq!(
            expected_impl.to_string(),
            impl_convert_result(original).unwrap().to_string()
        )
    }
}
