//! Return-type parsing for `#[node]`.
//!
//! Extracts the update type `U` from either `-> NodeResult<U>` or
//! `-> Result<NodeResult<U>, PeError>` and reports which style
//! the user chose so the codegen can wrap accordingly.

use proc_macro2::TokenStream;
use quote::quote;
use syn::{Error, Result};

/// Whether the user's return type is bare `NodeResult<U>` or
/// `Result<NodeResult<U>, PeError>`.
pub(crate) enum ReturnStyle {
    /// `-> NodeResult<U>` -- body produces NodeResult directly
    Bare,
    /// `-> Result<NodeResult<U>, PeError>` -- body uses `?`, generated code unwraps
    Fallible,
}

/// Extract the update type `U` from `-> NodeResult<U>` or
/// `-> Result<NodeResult<U>, PeError>`.
///
/// Returns `(U_tokens, ReturnStyle)` so the caller knows whether to wrap the body.
pub(crate) fn extract_return_update_type(
    output: &syn::ReturnType,
) -> Result<(TokenStream, ReturnStyle)> {
    match output {
        syn::ReturnType::Type(_, ty) => {
            if let syn::Type::Path(type_path) = &**ty {
                let last_segment = type_path
                    .path
                    .segments
                    .last()
                    .ok_or_else(|| Error::new_spanned(ty, "empty path"))?;

                // Case 1: Result<NodeResult<U>, PeError>
                if last_segment.ident == "Result" {
                    return extract_result_wrapped(ty, last_segment);
                }

                // Case 2: NodeResult<U>
                if last_segment.ident == "NodeResult" {
                    let update_ty = extract_angle_bracket_arg(ty, last_segment)?;
                    return Ok((update_ty, ReturnStyle::Bare));
                }

                return Err(Error::new_spanned(
                    &last_segment.ident,
                    "#[node] return type must be `NodeResult<UpdateType>` \
                     or `Result<NodeResult<UpdateType>, PeError>`",
                ));
            }

            Err(Error::new_spanned(
                ty,
                "#[node] return type must be `NodeResult<UpdateType>` \
                 or `Result<NodeResult<UpdateType>, PeError>`",
            ))
        }
        syn::ReturnType::Default => Err(Error::new_spanned(
            output,
            "#[node] function must return `NodeResult<UpdateType>` \
             or `Result<NodeResult<UpdateType>, PeError>`",
        )),
    }
}

/// Extract `U` from `Result<NodeResult<U>, PeError>`.
fn extract_result_wrapped(
    full_ty: &syn::Type,
    segment: &syn::PathSegment,
) -> Result<(TokenStream, ReturnStyle)> {
    if let syn::PathArguments::AngleBracketed(args) = &segment.arguments {
        let first_arg = args
            .args
            .first()
            .ok_or_else(|| Error::new_spanned(full_ty, "Result must have type parameters"))?;

        if let syn::GenericArgument::Type(syn::Type::Path(inner_path)) = first_arg {
            let inner_seg = inner_path
                .path
                .segments
                .last()
                .ok_or_else(|| Error::new_spanned(full_ty, "empty path inside Result"))?;

            if inner_seg.ident == "NodeResult" {
                let update_ty = extract_angle_bracket_arg(full_ty, inner_seg)?;
                return Ok((update_ty, ReturnStyle::Fallible));
            }
        }
    }

    Err(Error::new_spanned(
        full_ty,
        "#[node] Result return type must be `Result<NodeResult<UpdateType>, PeError>`",
    ))
}

/// Extract the first generic type argument from a path segment: `Foo<T>` -> `T`.
fn extract_angle_bracket_arg(
    span_ty: &syn::Type,
    segment: &syn::PathSegment,
) -> Result<TokenStream> {
    if let syn::PathArguments::AngleBracketed(args) = &segment.arguments {
        let first_arg = args
            .args
            .first()
            .ok_or_else(|| Error::new_spanned(span_ty, "expected a type parameter"))?;

        if let syn::GenericArgument::Type(inner_ty) = first_arg {
            return Ok(quote! { #inner_ty });
        }
    }

    Err(Error::new_spanned(
        span_ty,
        "expected angle-bracketed type parameter",
    ))
}
