//! # axum_thiserror
//! `axum_thiserror` is a library that offers a procedural macro to allow `thiserror` error types to be used as `axum` responses.
//! ## Usage
//! Add the library to your current project using Cargo:
//! ```bash
//! cargo add axum_thiserror
//! ```
//! Then you can create a basic `thiserror` error:
//! ```rust,ignore
//! #[derive(Error, Debug)]
//! pub enum UserCreateError {
//!   #[error("User {0} already exists")]
//!   UserAlreadyExists(String),
//! }
//! ```
//! Now you can use `axum_thiserror` to implement `IntoResponse` on your error. By
//! default (the `utoipa` feature), the response is the variant's status code with
//! `axum::Json(self)` as the body, so the enum must also derive `serde::Serialize`
//! and `utoipa::ToSchema` (add `serde` and `utoipa` to your dependencies) —
//! a missing implementation is a compile error:
//! ```rust,ignore
//! #[derive(Error, Debug, Serialize, ToSchema, ErrorStatus)]
//! pub enum UserCreateError {
//!   #[error("User {0} already exists")]
//!   #[status(StatusCode::CONFLICT)]
//!   UserAlreadyExists(String),
//! }
//! ```
//! ## Features
//! Three cargo features control the response body and the compile-time assertions
//! on the enum (a missing implementation is a compile error, not a silent fallback):
//! - `json` — respond with `axum::Json(self)`; the enum must implement
//!   `serde::Serialize`.
//! - `utoipa` (default, implies `json`) — additionally requires `utoipa::ToSchema`
//!   and derives `utoipa::IntoResponses` from the same `#[status(...)]` attributes.
//! - `schemars` (implies `json`) — additionally requires `schemars::JsonSchema`,
//!   giving a schema matching the response body exactly for OpenAPI tooling such
//!   as `aide`.
//!
//! `schemars` and `utoipa` compose — enable both while migrating between them.
//! ## OpenAPI responses with `utoipa`
//! With the `utoipa` feature, the derive also implements `utoipa::IntoResponses`:
//! one response per distinct status code (variants sharing a status are merged),
//! described by the variants' doc comments (one per line), with
//! `application/json` content using the enum's own schema:
//! ```rust,ignore
//! #[derive(Error, Debug, Serialize, ToSchema, ErrorStatus)]
//! pub enum UserCreateError {
//!   /// A user with this name already exists.
//!   #[error("User {0} already exists")]
//!   #[status(StatusCode::CONFLICT)]
//!   UserAlreadyExists(String),
//! }
//!
//! #[utoipa::path(post, path = "/users", responses(UserCreateError))]
//! async fn create_user() -> Result<(), UserCreateError> { /* ... */ }
//! ```
//! (With utoipa's `auto_into_responses` feature the error responses are picked up
//! from the handler's return type without listing them in `responses(...)`.)
//!
//! An enum missing `ToSchema` fails with a readable error:
#![cfg_attr(feature = "utoipa", doc = "```compile_fail")]
#![cfg_attr(not(feature = "utoipa"), doc = "```rust,ignore")]
//! use axum::http::StatusCode;
//! use axum_thiserror::ErrorStatus;
//!
//! #[derive(thiserror::Error, Debug, serde::Serialize, schemars::JsonSchema, ErrorStatus)]
//! pub enum UserCreateError {
//!     #[error("User {0} already exists")]
//!     #[status(StatusCode::CONFLICT)]
//!     UserAlreadyExists(String),
//! }
//! ```
//! ## Plain-text responses
//! Disable default features (`default-features = false`) to simply implement
//! `IntoResponse` with the variant's status code and the plain-text `#[error(...)]`
//! message as the body — no extra derives needed:
//! ```rust,ignore
//! #[derive(Error, Debug, ErrorStatus)]
//! pub enum UserCreateError {
//!   #[error("User {0} already exists")]
//!   #[status(StatusCode::CONFLICT)]
//!   UserAlreadyExists(String),
//! }
//! ```
//! ## License
//! This project is licensed under the [MIT License](LICENSE).

use proc_macro2::TokenStream;
use quote::{quote, quote_spanned};
use syn::{
    Data, DeriveInput, LitInt, Meta, Path, Variant, parse_macro_input, punctuated::Punctuated,
    spanned::Spanned, token::Comma,
};

/// A derivation that implements the `IntoResponse` trait for a specific attribute.
#[proc_macro_derive(ErrorStatus, attributes(status))]
pub fn derive_error_status(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let ast: DeriveInput = parse_macro_input!(input);
    let enum_ident = ast.ident;
    let variants = match ast.data {
        Data::Enum(data) => data.variants,
        _ => panic!(
            "#[derive(ErrorStatus)] is only available for enums, other types are not supported."
        ),
    };
    let status_arms: Punctuated<TokenStream, Comma> =
        variants.iter().map(impl_enum_variant).collect();

    #[cfg(feature = "json")]
    let body = quote!(axum::Json(self));
    #[cfg(not(feature = "json"))]
    let body = quote!(format!("{}", self));

    #[cfg(feature = "schemars")]
    let bounds_assertion = quote_spanned! { enum_ident.span() =>
        const _: fn() = || {
            fn assert_json_error_response<T: ::serde::Serialize + ::schemars::JsonSchema>() {}
            assert_json_error_response::<#enum_ident>();
        };
    };
    #[cfg(all(feature = "json", not(feature = "schemars")))]
    let bounds_assertion = quote_spanned! { enum_ident.span() =>
        const _: fn() = || {
            fn assert_json_error_response<T: ::serde::Serialize>() {}
            assert_json_error_response::<#enum_ident>();
        };
    };
    #[cfg(not(feature = "json"))]
    let bounds_assertion = quote!();

    #[cfg(feature = "utoipa")]
    let into_responses_impl = impl_into_responses(&enum_ident, &variants);
    #[cfg(not(feature = "utoipa"))]
    let into_responses_impl = quote!();

    quote! {
        #bounds_assertion

        impl axum::response::IntoResponse for #enum_ident {
            fn into_response(self) -> axum::response::Response {
                let status: axum::http::StatusCode = match &self {
                    #status_arms
                };
                axum::response::IntoResponse::into_response((status, #body))
            }
        }

        #into_responses_impl
    }
    .into()
}

/// Implements `utoipa::IntoResponses` from the same `#[status(...)]` attributes:
/// one entry per distinct status code (variants sharing a status are merged),
/// described by the variants' doc comments, with an `application/json` content
/// whose schema is the enum's own `utoipa` schema.
#[cfg(feature = "utoipa")]
fn impl_into_responses(
    enum_ident: &syn::Ident,
    variants: &Punctuated<Variant, Comma>,
) -> TokenStream {
    let schema_assertion = quote_spanned! { enum_ident.span() =>
        const _: fn() = || {
            fn assert_utoipa_error_response<T: ::utoipa::ToSchema>() {}
            assert_utoipa_error_response::<#enum_ident>();
        };
    };
    let variant_entries: Vec<TokenStream> = variants
        .iter()
        .map(|variant| {
            let status_code = find_status_code(variant);
            let doc = doc_comment(variant);
            quote! {
                let status: axum::http::StatusCode = #status_code;
                let descriptions = statuses.entry(status.as_u16()).or_insert_with(Vec::new);
                if !#doc.is_empty() {
                    descriptions.push(#doc);
                }
            }
        })
        .collect();

    quote! {
        #schema_assertion

        impl utoipa::IntoResponses for #enum_ident {
            fn responses() -> ::std::collections::BTreeMap<
                ::std::string::String,
                utoipa::openapi::RefOr<utoipa::openapi::response::Response>,
            > {
                let mut statuses: ::std::collections::BTreeMap<u16, ::std::vec::Vec<&str>> =
                    ::std::collections::BTreeMap::new();
                #({ #variant_entries })*
                statuses
                    .into_iter()
                    .map(|(status, descriptions)| {
                        (
                            status.to_string(),
                            utoipa::openapi::RefOr::T(
                                utoipa::openapi::response::ResponseBuilder::new()
                                    .description(descriptions.join("\n"))
                                    .content(
                                        "application/json",
                                        utoipa::openapi::content::ContentBuilder::new()
                                            .schema(::std::option::Option::Some(
                                                <Self as utoipa::PartialSchema>::schema(),
                                            ))
                                            .build(),
                                    )
                                    .build(),
                            ),
                        )
                    })
                    .collect()
            }
        }
    }
}

/// The variant's doc comment, lines trimmed and joined with newlines.
#[cfg(feature = "utoipa")]
fn doc_comment(variant: &Variant) -> String {
    variant
        .attrs
        .iter()
        .filter(|attr| attr.path().is_ident("doc"))
        .filter_map(|attr| match &attr.meta {
            Meta::NameValue(name_value) => match &name_value.value {
                syn::Expr::Lit(syn::ExprLit {
                    lit: syn::Lit::Str(line),
                    ..
                }) => Some(line.value().trim().to_string()),
                _ => None,
            },
            _ => None,
        })
        .collect::<Vec<_>>()
        .join("\n")
}

fn impl_enum_variant(input: &Variant) -> TokenStream {
    let status_code = find_status_code(input);
    let case = if input.fields.is_empty() {
        case_empty_fields(input)
    } else if input.fields.iter().filter(|f| f.ident.is_none()).count() > 0 {
        case_unnamed_fields(input)
    } else {
        case_named_fields(input)
    };
    quote! {
        Self::#case => #status_code
    }
}

fn case_empty_fields(input: &Variant) -> TokenStream {
    let ident = &input.ident;
    quote!(#ident)
}

fn case_unnamed_fields(input: &Variant) -> TokenStream {
    let ident = &input.ident;
    quote!(#ident( .. ))
}

fn case_named_fields(input: &Variant) -> TokenStream {
    let ident = &input.ident;
    quote!(#ident { .. })
}

fn find_status_code(input: &Variant) -> TokenStream {
    match input
        .attrs
        .iter()
        .find(|attr| attr.path().is_ident("status"))
    {
        Some(attr) => match &attr.meta {
            Meta::List(l) => {
                if let Ok(number) = l.parse_args::<LitInt>() {
                    quote! {
                        axum::http::StatusCode::from_u16(#number as u16).unwrap()
                    }
                } else if let Ok(expr) = l.parse_args::<Path>() {
                    quote! {
                        #expr
                    }
                } else {
                    quote_spanned!(l.span() => compile_error!("Only #[status(StatusCode)] or #[status(u16)] syntaxe is supported"))
                }
            }
            _ => {
                quote_spanned! { attr.span() => compile_error!("Only #[status(StatusCode)] or #[status(u16)] syntaxe is supported") }
            }
        },
        None => {
            quote_spanned! { input.span() => compile_error!("Each enum variant should have a status code provided using the #[status()] attribute") }
        }
    }
}