use proc_macro::TokenStream;
use proc_macro2::Span;
use quote::quote;
use syn::{
    parse::{Parse, ParseStream},
    parse_macro_input, Expr, FnArg, ItemFn, ItemStruct, LitStr, Pat, Token,
};

// ── argument parsers ──────────────────────────────────────────────────────────

struct RoleArg {
    role: LitStr,
}

impl Parse for RoleArg {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        Ok(RoleArg {
            role: input.parse()?,
        })
    }
}

struct AnyRoleArgs {
    roles: Vec<LitStr>,
}

impl Parse for AnyRoleArgs {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let mut roles = Vec::new();
        roles.push(input.parse::<LitStr>()?);
        while input.peek(Token![,]) {
            input.parse::<Token![,]>()?;
            if input.is_empty() {
                break;
            }
            roles.push(input.parse::<LitStr>()?);
        }
        Ok(AnyRoleArgs { roles })
    }
}

// ── claims-source detection ───────────────────────────────────────────────────

enum ClaimsSource {
    Claims(syn::Ident),
    Ctx(syn::Ident),
}

/// Find the first parameter typed as `Claims` or `Ctx` in the function signature.
fn find_claims_source(func: &ItemFn) -> Option<ClaimsSource> {
    for arg in &func.sig.inputs {
        let FnArg::Typed(pat_type) = arg else {
            continue;
        };
        let Pat::Ident(pat_ident) = pat_type.pat.as_ref() else {
            continue;
        };
        if let syn::Type::Path(tp) = pat_type.ty.as_ref() {
            match tp
                .path
                .segments
                .last()
                .map(|s| s.ident.to_string())
                .as_deref()
            {
                Some("Claims") => return Some(ClaimsSource::Claims(pat_ident.ident.clone())),
                Some("Ctx") => return Some(ClaimsSource::Ctx(pat_ident.ident.clone())),
                _ => {}
            }
        }
    }
    None
}

/// Find the identifier of the first `Claims`-typed parameter (legacy helper).
fn find_claims_ident(func: &ItemFn) -> Option<syn::Ident> {
    match find_claims_source(func)? {
        ClaimsSource::Claims(id) => Some(id),
        ClaimsSource::Ctx(_) => None,
    }
}

// ── argument parsers (ability variants reuse RoleArg / AnyRoleArgs) ──────────

// (Ability macros share the same parse structs as role macros.)

// ── response builders ─────────────────────────────────────────────────────────

fn forbidden_response(role_display: proc_macro2::TokenStream) -> proc_macro2::TokenStream {
    quote! {
        {
            use ::axum::response::IntoResponse;
            return (
                ::axum::http::StatusCode::FORBIDDEN,
                ::axum::Json(::serde_json::json!({
                    "error": "forbidden",
                    "message": format!("role required: {}", #role_display)
                }))
            ).into_response();
        }
    }
}

fn ability_forbidden_response(
    ability_display: proc_macro2::TokenStream,
) -> proc_macro2::TokenStream {
    quote! {
        {
            use ::axum::response::IntoResponse;
            return (
                ::axum::http::StatusCode::FORBIDDEN,
                ::axum::Json(::serde_json::json!({
                    "error": "forbidden",
                    "message": format!("ability required: {}", #ability_display)
                }))
            ).into_response();
        }
    }
}

fn unauthorized_response() -> proc_macro2::TokenStream {
    quote! {
        {
            use ::axum::response::IntoResponse;
            return (
                ::axum::http::StatusCode::UNAUTHORIZED,
                ::axum::Json(::serde_json::json!({
                    "error": "unauthorized",
                    "message": "authentication required"
                }))
            ).into_response();
        }
    }
}

// ── shared rewrite helper ─────────────────────────────────────────────────────

fn rewrite_fn_with_guard(func: &mut ItemFn, guard: proc_macro2::TokenStream) {
    func.sig.output =
        syn::parse2(quote! { -> ::axum::response::Response }).expect("return type parse failed");

    let original_stmts = func.block.stmts.drain(..).collect::<Vec<_>>();

    *func.block = syn::parse2(quote! {
        {
            #guard
            let __body_result = { #(#original_stmts)* };
            ::axum::response::IntoResponse::into_response(__body_result)
        }
    })
    .expect("block rewrite failed");
}

// ── #[controller] ─────────────────────────────────────────────────────────────

/// Marks a struct as a controller — derives `Clone` and `Debug`.
///
/// # Example
///
/// ```rust,ignore
/// #[controller]
/// pub struct UserController;
/// ```
#[proc_macro_attribute]
pub fn controller(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let input = parse_macro_input!(item as ItemStruct);
    quote! {
        #[derive(Clone, Debug)]
        #input
    }
    .into()
}

// ── #[require_role("role")] ───────────────────────────────────────────────────

/// Injects a role-check before the handler body.
///
/// Works with both `claims: Claims` and `ctx: Ctx` parameters:
/// - `Claims` → checks `claims.has_role(role)`, returns 403 if not found.
/// - `Ctx` → returns 401 if unauthenticated, 403 if role absent.
///
/// # Example
///
/// ```rust,ignore
/// #[require_role("admin")]
/// async fn admin_only(claims: Claims) -> impl IntoResponse { ... }
///
/// #[require_role("admin")]
/// async fn admin_ctx(ctx: Ctx) -> impl IntoResponse { ... }
/// ```
#[proc_macro_attribute]
pub fn require_role(attr: TokenStream, item: TokenStream) -> TokenStream {
    let RoleArg { role } = parse_macro_input!(attr as RoleArg);
    let mut func = parse_macro_input!(item as ItemFn);

    let role_val = role.value();
    let forbidden = forbidden_response(quote! { #role_val });

    let guard = match find_claims_source(&func) {
        Some(ClaimsSource::Claims(ident)) => quote! {
            if !#ident.has_role(#role) {
                #forbidden
            }
        },
        Some(ClaimsSource::Ctx(ident)) => {
            let unauthorized = unauthorized_response();
            quote! {
                match #ident.user.as_ref() {
                    None => { #unauthorized }
                    Some(__claims) if !__claims.has_role(#role) => { #forbidden }
                    _ => {}
                }
            }
        }
        None => {
            return syn::Error::new(
                Span::call_site(),
                "#[require_role] requires a `claims: Claims` or `ctx: Ctx` parameter",
            )
            .to_compile_error()
            .into();
        }
    };

    rewrite_fn_with_guard(&mut func, guard);
    quote! { #func }.into()
}

// ── #[require_any_role("a", "b")] ────────────────────────────────────────────

/// Injects an any-of-roles check before the handler body.
///
/// Works with `claims: Claims` parameters. Returns 403 if none of the roles
/// are present.
///
/// # Example
///
/// ```rust,ignore
/// #[require_any_role("admin", "moderator")]
/// async fn mod_panel(claims: Claims) -> impl IntoResponse { ... }
/// ```
#[proc_macro_attribute]
pub fn require_any_role(attr: TokenStream, item: TokenStream) -> TokenStream {
    let AnyRoleArgs { roles } = parse_macro_input!(attr as AnyRoleArgs);
    let mut func = parse_macro_input!(item as ItemFn);

    let claims_ident = match find_claims_ident(&func) {
        Some(id) => id,
        None => {
            return syn::Error::new(
                Span::call_site(),
                "#[require_any_role] requires a `claims: Claims` parameter in the function signature",
            )
            .to_compile_error()
            .into();
        }
    };

    let roles_display = roles
        .iter()
        .map(|r| r.value())
        .collect::<Vec<_>>()
        .join(" | ");
    let forbidden = forbidden_response(quote! { #roles_display });

    let guard = quote! {
        if !#claims_ident.has_any_role(&[#(#roles),*]) {
            #forbidden
        }
    };

    rewrite_fn_with_guard(&mut func, guard);
    quote! { #func }.into()
}

// ── #[require_ability("ability")] ────────────────────────────────────────────

/// Injects an ability check before the handler body.
///
/// Requires a `ctx: Ctx` parameter.  Returns 401 if unauthenticated, 403 if
/// the current token does not grant the named ability.
///
/// # Example
///
/// ```rust,ignore
/// #[require_ability("posts:write")]
/// async fn create_post(ctx: Ctx) -> impl IntoResponse { ... }
/// ```
#[proc_macro_attribute]
pub fn require_ability(attr: TokenStream, item: TokenStream) -> TokenStream {
    let RoleArg { role } = parse_macro_input!(attr as RoleArg);
    let mut func = parse_macro_input!(item as ItemFn);

    let ability_val = role.value();
    let forbidden = ability_forbidden_response(quote! { #ability_val });

    let guard = match find_claims_source(&func) {
        Some(ClaimsSource::Ctx(ident)) => {
            let unauthorized = unauthorized_response();
            quote! {
                if #ident.user.is_none() {
                    #unauthorized
                }
                if #ident.token_cant(#role) {
                    #forbidden
                }
            }
        }
        _ => {
            return syn::Error::new(
                Span::call_site(),
                "#[require_ability] requires a `ctx: Ctx` parameter",
            )
            .to_compile_error()
            .into();
        }
    };

    rewrite_fn_with_guard(&mut func, guard);
    quote! { #func }.into()
}

// ── #[require_any_ability("a", "b")] ─────────────────────────────────────────

/// Injects an any-of-abilities check before the handler body.
///
/// Requires a `ctx: Ctx` parameter.  Returns 401 if unauthenticated, 403 if
/// the token does not grant at least one of the listed abilities.
///
/// # Example
///
/// ```rust,ignore
/// #[require_any_ability("admin", "moderator")]
/// async fn mod_panel(ctx: Ctx) -> impl IntoResponse { ... }
/// ```
#[proc_macro_attribute]
pub fn require_any_ability(attr: TokenStream, item: TokenStream) -> TokenStream {
    let AnyRoleArgs { roles } = parse_macro_input!(attr as AnyRoleArgs);
    let mut func = parse_macro_input!(item as ItemFn);

    let ctx_ident = match find_claims_source(&func) {
        Some(ClaimsSource::Ctx(id)) => id,
        _ => {
            return syn::Error::new(
                Span::call_site(),
                "#[require_any_ability] requires a `ctx: Ctx` parameter",
            )
            .to_compile_error()
            .into();
        }
    };

    let abilities_display = roles
        .iter()
        .map(|r| r.value())
        .collect::<Vec<_>>()
        .join(" | ");
    let forbidden = ability_forbidden_response(quote! { #abilities_display });
    let unauthorized = unauthorized_response();

    let guard = quote! {
        if #ctx_ident.user.is_none() {
            #unauthorized
        }
        if ![#(#roles),*].iter().any(|a| #ctx_ident.token_can(a)) {
            #forbidden
        }
    };

    rewrite_fn_with_guard(&mut func, guard);
    quote! { #func }.into()
}

// ── #[authorize(PolicyType::action)] / #[authorize(PolicyType::action, model)] ──

struct AuthorizeArg {
    policy_type: syn::Path,
    action: String,
    model: Option<Expr>,
}

impl Parse for AuthorizeArg {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let full_path: syn::Path = input.parse()?;
        let syn::Path {
            leading_colon,
            segments,
        } = full_path;

        let mut segs: Vec<syn::PathSegment> = segments.into_iter().collect();

        if segs.len() < 2 {
            return Err(syn::Error::new(
                Span::call_site(),
                "#[authorize] expects `PolicyType::action` (e.g. `PostPolicy::view`)",
            ));
        }

        let action = segs.pop().unwrap().ident.to_string();

        let policy_type = syn::Path {
            leading_colon,
            segments: segs.into_iter().collect(),
        };

        let model = if input.peek(Token![,]) {
            input.parse::<Token![,]>()?;
            if input.is_empty() {
                None
            } else {
                Some(input.parse::<Expr>()?)
            }
        } else {
            None
        };

        Ok(AuthorizeArg {
            policy_type,
            action,
            model,
        })
    }
}

/// Injects a policy authorization check before the handler body.
///
/// Requires a `ctx: Ctx` parameter.  Returns 401 if unauthenticated, 403 on
/// deny, or 404 when the policy uses `DenyAsNotFound`.
///
/// # Example
///
/// ```rust,ignore
/// // check view permission on an existing resource
/// #[authorize(PostPolicy::view, post)]
/// async fn show(ctx: Ctx, ...) -> impl IntoResponse { ... }
///
/// // check create permission (no existing model)
/// #[authorize(PostPolicy::create)]
/// async fn store(ctx: Ctx) -> impl IntoResponse { ... }
/// ```
#[proc_macro_attribute]
pub fn authorize(attr: TokenStream, item: TokenStream) -> TokenStream {
    let AuthorizeArg {
        policy_type,
        action,
        model,
    } = parse_macro_input!(attr as AuthorizeArg);
    let mut func = parse_macro_input!(item as ItemFn);

    let ctx_ident = match find_claims_source(&func) {
        Some(ClaimsSource::Ctx(id)) => id,
        _ => {
            return syn::Error::new(
                Span::call_site(),
                "#[authorize] requires a `ctx: Ctx` parameter",
            )
            .to_compile_error()
            .into();
        }
    };

    let model_expr = match model {
        Some(m) => quote! { &(#m) },
        None => quote! { &() },
    };

    let action_str = action.as_str();

    let guard = quote! {
        {
            use ::rok_bouncer::CtxBounceExt as _;
            use ::axum::response::IntoResponse as _;
            match #ctx_ident.authorize(&#policy_type, #action_str, #model_expr).await {
                Ok(()) => {}
                Err(::rok_bouncer::AuthzError::Unauthenticated) => {
                    return (
                        ::axum::http::StatusCode::UNAUTHORIZED,
                        ::axum::Json(::serde_json::json!({
                            "error": "unauthorized",
                            "message": "authentication required"
                        }))
                    ).into_response();
                }
                Err(::rok_bouncer::AuthzError::Forbidden) => {
                    return (
                        ::axum::http::StatusCode::FORBIDDEN,
                        ::axum::Json(::serde_json::json!({
                            "error": "forbidden",
                            "message": "authorization failed"
                        }))
                    ).into_response();
                }
                Err(::rok_bouncer::AuthzError::NotFound) => {
                    return (
                        ::axum::http::StatusCode::NOT_FOUND,
                        ::axum::Json(::serde_json::json!({
                            "error": "not found"
                        }))
                    ).into_response();
                }
            }
        }
    };

    rewrite_fn_with_guard(&mut func, guard);
    quote! { #func }.into()
}

// ── #[throttle(N, per = "unit"[, by = "ip"|"user"])] ─────────────────────────

struct ThrottleArgs {
    requests: u64,
    window_secs: u64,
    by: ThrottleBy,
}

enum ThrottleBy {
    Global,
    User,
}

impl Parse for ThrottleArgs {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let count: syn::LitInt = input.parse()?;
        let requests: u64 = count.base10_parse()?;

        input.parse::<Token![,]>()?;

        let per_ident: syn::Ident = input.parse()?;
        if per_ident != "per" {
            return Err(syn::Error::new(per_ident.span(), "expected `per`"));
        }
        input.parse::<Token![=]>()?;
        let unit: LitStr = input.parse()?;
        let window_secs = match unit.value().as_str() {
            "second" => 1u64,
            "minute" => 60,
            "hour" => 3_600,
            "day" => 86_400,
            other => {
                return Err(syn::Error::new(
                    unit.span(),
                    format!("unknown period `{other}`; use second, minute, hour, or day"),
                ));
            }
        };

        let by = if input.peek(Token![,]) {
            input.parse::<Token![,]>()?;
            if input.is_empty() {
                ThrottleBy::Global
            } else {
                let by_ident: syn::Ident = input.parse()?;
                if by_ident != "by" {
                    return Err(syn::Error::new(by_ident.span(), "expected `by`"));
                }
                input.parse::<Token![=]>()?;
                let by_val: LitStr = input.parse()?;
                match by_val.value().as_str() {
                    "user" => ThrottleBy::User,
                    "global" => ThrottleBy::Global,
                    other => {
                        return Err(syn::Error::new(
                            by_val.span(),
                            format!("unknown `by` value `{other}`; use `user` or `global`"),
                        ));
                    }
                }
            }
        } else {
            ThrottleBy::Global
        };

        Ok(ThrottleArgs {
            requests,
            window_secs,
            by,
        })
    }
}

/// Injects a rate-limit check before the handler body.
///
/// Uses the process-global `rok_rate_limit::Limiter` (in-memory, fixed-window).
///
/// # Syntax
///
/// ```rust,ignore
/// // Global handler-level limit (100 req/min shared across all callers)
/// #[throttle(100, per = "minute")]
/// async fn list_posts(ctx: Ctx) -> impl IntoResponse { ... }
///
/// // Per-user limit (10 req/sec per authenticated user)
/// #[throttle(10, per = "second", by = "user")]
/// async fn create_post(ctx: Ctx) -> impl IntoResponse { ... }
/// ```
///
/// For IP-based limiting use `rok_rate_limit::ThrottleLayer::by_ip` as a
/// router-level Tower middleware instead.
///
/// # Requirements
///
/// - Add `rok-rate-limit` to your `Cargo.toml`.
/// - `by = "user"` requires a `ctx: Ctx` parameter in the handler signature.
#[proc_macro_attribute]
pub fn throttle(attr: TokenStream, item: TokenStream) -> TokenStream {
    let ThrottleArgs {
        requests,
        window_secs,
        by,
    } = parse_macro_input!(attr as ThrottleArgs);
    let mut func = parse_macro_input!(item as ItemFn);

    let fn_name = func.sig.ident.to_string();

    let key_expr: proc_macro2::TokenStream = match by {
        ThrottleBy::Global => {
            quote! { format!("throttle:{}", #fn_name) }
        }
        ThrottleBy::User => {
            let ctx_ident = match find_claims_source(&func) {
                Some(ClaimsSource::Ctx(id)) => id,
                _ => {
                    return syn::Error::new(
                        Span::call_site(),
                        "#[throttle(..., by = \"user\")] requires a `ctx: Ctx` parameter",
                    )
                    .to_compile_error()
                    .into();
                }
            };
            quote! {
                {
                    let __sub = #ctx_ident.user.as_ref()
                        .map(|c| c.sub.as_str())
                        .unwrap_or("anonymous");
                    format!("throttle:{}:user:{}", #fn_name, __sub)
                }
            }
        }
    };

    let guard = quote! {
        {
            let __throttle_key = #key_expr;
            let __throttle_result = ::rok_rate_limit::global_limiter()
                .for_key(__throttle_key)
                .requests(#requests)
                .per(::std::time::Duration::from_secs(#window_secs))
                .check();
            if let ::rok_rate_limit::LimitResult::Exceeded { retry_after_secs: __retry } =
                __throttle_result
            {
                use ::axum::response::IntoResponse as _;
                return (
                    ::axum::http::StatusCode::TOO_MANY_REQUESTS,
                    [(
                        ::axum::http::header::RETRY_AFTER,
                        ::axum::http::HeaderValue::from_str(&__retry.to_string())
                            .unwrap_or(::axum::http::HeaderValue::from_static("1")),
                    )],
                    ::axum::Json(::serde_json::json!({
                        "error":       "too_many_requests",
                        "message":     "rate limit exceeded",
                        "retry_after": __retry,
                    }))
                ).into_response();
            }
        }
    };

    rewrite_fn_with_guard(&mut func, guard);
    quote! { #func }.into()
}
