#[cfg(feature = "axum")]
use axum::{
    Router,
    extract::connect_info::ConnectInfo,
    extract::{FromRequestParts, Request, State},
    http::Extensions,
    http::StatusCode,
    http::request::Parts,
    response::{IntoResponse, Response},
    routing::{delete, get, post},
};
#[cfg(feature = "axum")]
use std::collections::{HashMap, HashSet};
#[cfg(feature = "axum")]
use std::net::SocketAddr;
#[cfg(feature = "axum")]
use std::sync::Arc;

#[cfg(feature = "axum")]
use crate::AthenaAuth;
#[cfg(feature = "axum")]
use athena_auth_core::SessionManager;
#[cfg(feature = "axum")]
use athena_auth_core::entity::AuthSession as AuthSessionTrait;
use athena_auth_core::{
    AuthError, AuthRequest, AuthResponse, DatabaseAdapter, HealthCheckResponse, HttpMethod,
    OkResponse, core_paths,
};

use athena_auth_core::SessionOps;
/// Integration trait for Axum web framework
#[cfg(feature = "axum")]
pub trait AxumIntegration<DB: DatabaseAdapter> {
    /// Create an Axum router with all authentication routes
    fn axum_router(self) -> Router<Arc<AthenaAuth<DB>>>;
}

#[cfg(feature = "axum")]
impl<DB: DatabaseAdapter> AxumIntegration<DB> for Arc<AthenaAuth<DB>> {
    fn axum_router(self) -> Router<Arc<AthenaAuth<DB>>> {
        // NOTE: disabled_paths is checked here at route-registration time so
        // that disabled routes are never mounted in Axum at all.  The core
        // handler (`handle_request_inner`) performs the same check at
        // request-dispatch time for non-Axum integrations (direct
        // `handle_request` callers).  The duplication is intentional.
        let disabled_paths: Vec<String> = self.config().disabled_paths.clone();

        let mut router: Router<Arc<AthenaAuth<DB>>> = Router::new();
        let mut registered_routes: HashSet<String> = HashSet::new();

        // Add status endpoints
        if !disabled_paths.contains(&core_paths::OK.to_string()) {
            router = register_route::<DB>(
                router,
                &mut registered_routes,
                "GET",
                core_paths::OK,
                get(ok_check),
            );
        }
        if !disabled_paths.contains(&core_paths::ERROR.to_string()) {
            router = register_route::<DB>(
                router,
                &mut registered_routes,
                "GET",
                core_paths::ERROR,
                get(error_check),
            );
        }

        // Add default health check route
        if !disabled_paths.contains(&core_paths::HEALTH.to_string()) {
            router = register_route::<DB>(
                router,
                &mut registered_routes,
                "GET",
                core_paths::HEALTH,
                get(health_check),
            );
        }

        // Add OpenAPI spec endpoint
        if !disabled_paths.contains(&core_paths::OPENAPI_SPEC.to_string()) {
            router = register_route::<DB>(
                router,
                &mut registered_routes,
                "GET",
                core_paths::OPENAPI_SPEC,
                get(create_plugin_handler::<DB>()),
            );
        }

        // Add core user management routes
        if !disabled_paths.contains(&core_paths::UPDATE_USER.to_string()) {
            router = register_route::<DB>(
                router,
                &mut registered_routes,
                "POST",
                core_paths::UPDATE_USER,
                post(create_plugin_handler::<DB>()),
            );
        }
        if !disabled_paths.contains(&core_paths::DELETE_USER.to_string()) {
            router = register_route::<DB>(
                router,
                &mut registered_routes,
                "POST",
                core_paths::DELETE_USER,
                post(create_plugin_handler::<DB>()),
            );
            router = register_route::<DB>(
                router,
                &mut registered_routes,
                "DELETE",
                core_paths::DELETE_USER,
                delete(create_plugin_handler::<DB>()),
            );
        }
        if !disabled_paths.contains(&core_paths::CHANGE_EMAIL.to_string()) {
            router = register_route::<DB>(
                router,
                &mut registered_routes,
                "POST",
                core_paths::CHANGE_EMAIL,
                post(create_plugin_handler::<DB>()),
            );
        }
        if !disabled_paths.contains(&core_paths::DELETE_USER_CALLBACK.to_string()) {
            router = register_route::<DB>(
                router,
                &mut registered_routes,
                "GET",
                core_paths::DELETE_USER_CALLBACK,
                get(create_plugin_handler::<DB>()),
            );
        }

        // Register plugin routes
        for plugin in self.plugins() {
            for route in plugin.routes() {
                // Skip disabled paths
                if disabled_paths.contains(&route.path) {
                    continue;
                }

                let handler_fn = create_plugin_handler::<DB>();
                match route.method {
                    HttpMethod::Get => {
                        router = register_route::<DB>(
                            router,
                            &mut registered_routes,
                            "GET",
                            &route.path,
                            get(handler_fn.clone()),
                        );
                    }
                    HttpMethod::Post => {
                        router = register_route::<DB>(
                            router,
                            &mut registered_routes,
                            "POST",
                            &route.path,
                            post(handler_fn.clone()),
                        );
                    }
                    HttpMethod::Put => {
                        router = register_route::<DB>(
                            router,
                            &mut registered_routes,
                            "PUT",
                            &route.path,
                            axum::routing::put(handler_fn.clone()),
                        );
                    }
                    HttpMethod::Delete => {
                        router = register_route::<DB>(
                            router,
                            &mut registered_routes,
                            "DELETE",
                            &route.path,
                            axum::routing::delete(handler_fn.clone()),
                        );
                    }
                    HttpMethod::Patch => {
                        router = register_route::<DB>(
                            router,
                            &mut registered_routes,
                            "PATCH",
                            &route.path,
                            axum::routing::patch(handler_fn.clone()),
                        );
                    }
                    _ => {} // Skip unsupported methods
                }
            }
        }

        router.with_state(self)
    }
}

#[cfg(feature = "axum")]
fn register_route<DB: DatabaseAdapter>(
    router: Router<Arc<AthenaAuth<DB>>>,
    registered_routes: &mut HashSet<String>,
    method: &'static str,
    path: &str,
    method_router: axum::routing::MethodRouter<Arc<AthenaAuth<DB>>>,
) -> Router<Arc<AthenaAuth<DB>>> {
    let route_key = format!("{method}:{path}");
    if registered_routes.insert(route_key) {
        router.route(path, method_router)
    } else {
        router
    }
}

#[cfg(feature = "axum")]
async fn ok_check() -> impl IntoResponse {
    axum::Json(OkResponse { ok: true })
}

#[cfg(feature = "axum")]
async fn error_check() -> impl IntoResponse {
    axum::Json(OkResponse { ok: false })
}

#[cfg(feature = "axum")]
async fn health_check() -> impl IntoResponse {
    axum::Json(HealthCheckResponse {
        status: "ok",
        // Keep compatibility with existing test + consumer expectations for auth health.
        service: "athena-auth",
    })
}

#[cfg(feature = "axum")]
#[allow(clippy::type_complexity)]
fn create_plugin_handler<DB: DatabaseAdapter>() -> impl Fn(
    State<Arc<AthenaAuth<DB>>>,
    Request,
) -> std::pin::Pin<
    Box<dyn std::future::Future<Output = Response> + Send>,
> + Clone {
    |State(auth): State<Arc<AthenaAuth<DB>>>, req: Request| {
        Box::pin(async move { dispatch_auth_axum_request(auth.as_ref(), req).await })
    }
}

#[cfg(feature = "axum")]
const INTERNAL_AXUM_REMOTE_IP_HEADER: &str = "x-athena-client-ip";
#[cfg(feature = "axum")]
const EXTERNAL_CLIENT_IP_HEADERS: [&str; 7] = [
    "x-forwarded-for",
    "x-real-ip",
    "cf-connecting-ip",
    "true-client-ip",
    "x-client-ip",
    "x-cluster-client-ip",
    "forwarded",
];

#[cfg(feature = "axum")]
fn contains_external_client_ip_header(headers: &HashMap<String, String>) -> bool {
    headers.keys().any(|header_name| {
        EXTERNAL_CLIENT_IP_HEADERS
            .iter()
            .any(|known| header_name.eq_ignore_ascii_case(known))
    })
}

#[cfg(feature = "axum")]
fn remote_ip_from_connect_info(extensions: &Extensions) -> Option<String> {
    if let Some(connect_info) = extensions.get::<ConnectInfo<SocketAddr>>() {
        return Some(connect_info.0.ip().to_string());
    }
    if let Some(addr) = extensions.get::<SocketAddr>() {
        return Some(addr.ip().to_string());
    }
    None
}

#[cfg(feature = "axum")]
pub async fn dispatch_auth_axum_request<DB: DatabaseAdapter>(
    auth: &AthenaAuth<DB>,
    req: Request,
) -> Response {
    match convert_axum_request(req).await {
        Ok(auth_req) => match auth.handle_request(auth_req).await {
            Ok(auth_response) => convert_auth_response(auth_response),
            Err(err) => convert_auth_error(err),
        },
        Err(err) => convert_auth_error(err),
    }
}

#[cfg(feature = "axum")]
async fn convert_axum_request(req: Request) -> Result<AuthRequest, AuthError> {
    let (parts, body) = req.into_parts();

    // Convert method
    let method = match parts.method {
        axum::http::Method::GET => HttpMethod::Get,
        axum::http::Method::POST => HttpMethod::Post,
        axum::http::Method::PUT => HttpMethod::Put,
        axum::http::Method::DELETE => HttpMethod::Delete,
        axum::http::Method::PATCH => HttpMethod::Patch,
        axum::http::Method::OPTIONS => HttpMethod::Options,
        axum::http::Method::HEAD => HttpMethod::Head,
        _ => {
            return Err(AuthError::InvalidRequest(
                "Unsupported HTTP method".to_string(),
            ));
        }
    };

    // Convert headers
    let mut headers: HashMap<String, String> = HashMap::new();
    for (name, value) in parts.headers.iter() {
        if name
            .as_str()
            .eq_ignore_ascii_case(INTERNAL_AXUM_REMOTE_IP_HEADER)
        {
            continue;
        }
        if let Ok(value_str) = value.to_str() {
            headers.insert(name.to_string(), value_str.to_string());
        }
    }
    if !contains_external_client_ip_header(&headers)
        && let Some(remote_ip) = remote_ip_from_connect_info(&parts.extensions)
    {
        headers.insert(INTERNAL_AXUM_REMOTE_IP_HEADER.to_string(), remote_ip);
    }

    // Get path
    let path: String = parts.uri.path().to_string();

    // Convert query parameters
    let mut query: HashMap<String, String> = HashMap::new();
    if let Some(query_str) = parts.uri.query() {
        for (key, value) in url::form_urlencoded::parse(query_str.as_bytes()) {
            query.insert(key.to_string(), value.to_string());
        }
    }

    // Convert body
    let body_bytes: Option<Vec<u8>> = match axum::body::to_bytes(body, usize::MAX).await {
        Ok(bytes) => {
            if bytes.is_empty() {
                None
            } else {
                Some(bytes.to_vec())
            }
        }
        Err(_) => None,
    };

    Ok(AuthRequest::from_parts(
        method, path, headers, body_bytes, query,
    ))
}

#[cfg(feature = "axum")]
fn convert_auth_response(auth_response: AuthResponse) -> Response {
    let status: StatusCode =
        StatusCode::from_u16(auth_response.status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
    let mut response = Response::new(axum::body::Body::from(auth_response.body));
    *response.status_mut() = status;

    // Add headers. Invalid names/values are skipped to preserve response generation.
    for (name, value) in auth_response.headers {
        if let (Ok(header_name), Ok(header_value)) = (
            axum::http::HeaderName::from_bytes(name.as_bytes()),
            axum::http::HeaderValue::from_str(&value),
        ) {
            response.headers_mut().insert(header_name, header_value);
        }
    }

    response
}

#[cfg(feature = "axum")]
fn convert_auth_error(err: AuthError) -> Response {
    convert_auth_response(err.into_response())
}

// ---------------------------------------------------------------------------
// Axum extractors
// ---------------------------------------------------------------------------

/// Authenticated session extractor.
///
/// Extracts and validates the current user and session from the request.
/// Returns `401 Unauthorized` if no valid session is found.
///
/// Requires `State<Arc<AthenaAuth<DB>>>` to be present in the router.
///
/// # Example
///
/// ```rust,ignore
/// use athena_auth::handlers::axum::CurrentSession;
///
/// async fn profile(session: CurrentSession<MyDB>) -> impl IntoResponse {
///     let user = &session.user;
///     let session = &session.session;
///     axum::Json(serde_json::json!({ "id": user.id() }))
/// }
/// ```
#[cfg(feature = "axum")]
pub struct CurrentSession<DB: DatabaseAdapter> {
    pub user: DB::User,
    pub session: DB::Session,
}

/// Optional authenticated session extractor.
///
/// Like [`CurrentSession`] but returns `None` instead of a 401 error when
/// no valid session is found. Useful for routes that behave differently
/// for authenticated vs anonymous users.
///
/// # Example
///
/// ```rust,ignore
/// async fn home(session: OptionalSession<MyDB>) -> impl IntoResponse {
///     if let Some(session) = session.0 {
///         axum::Json(serde_json::json!({ "user": session.user.id() }))
///     } else {
///         axum::Json(serde_json::json!({ "user": null }))
///     }
/// }
/// ```
#[cfg(feature = "axum")]
pub struct OptionalSession<DB: DatabaseAdapter>(pub Option<CurrentSession<DB>>);

/// Extract a session token from the request parts.
///
/// Checks the `Authorization: Bearer <token>` header first, then falls
/// back to the configured session cookie.
#[cfg(feature = "axum")]
fn extract_token_from_parts(parts: &Parts, cookie_name: &str) -> Option<String> {
    // Try Bearer token first
    if let Some(auth_header) = parts.headers.get("authorization")
        && let Ok(auth_str) = auth_header.to_str()
        && let Some(token) = auth_str.strip_prefix("Bearer ")
    {
        return Some(token.to_string());
    }

    // Fall back to cookie
    if let Some(cookie_header) = parts.headers.get("cookie")
        && let Ok(cookie_str) = cookie_header.to_str()
    {
        for part in cookie_str.split(';') {
            let part = part.trim();
            if let Some(value) = part.strip_prefix(&format!("{}=", cookie_name))
                && !value.is_empty()
            {
                return Some(value.to_string());
            }
        }
    }

    None
}

#[cfg(feature = "axum")]
impl<DB: DatabaseAdapter> FromRequestParts<Arc<AthenaAuth<DB>>> for CurrentSession<DB> {
    type Rejection = Response;

    async fn from_request_parts(
        parts: &mut Parts,
        state: &Arc<AthenaAuth<DB>>,
    ) -> Result<Self, Self::Rejection> {
        let cookie_name: &String = &state.config().session.cookie_name;
        let token: String = extract_token_from_parts(parts, cookie_name)
            .ok_or_else(|| convert_auth_error(AuthError::Unauthenticated))?;

        let session_manager: SessionManager<DB> = state.session_manager().clone();

        let session: <DB as SessionOps>::Session = session_manager
            .get_session(&token)
            .await
            .map_err(convert_auth_error)?
            .ok_or_else(|| convert_auth_error(AuthError::SessionNotFound))?;

        let user = state
            .database()
            .get_user_by_id(session.user_id())
            .await
            .map_err(convert_auth_error)?
            .ok_or_else(|| convert_auth_error(AuthError::UserNotFound))?;

        Ok(CurrentSession { user, session })
    }
}

#[cfg(feature = "axum")]
impl<DB: DatabaseAdapter> FromRequestParts<Arc<AthenaAuth<DB>>> for OptionalSession<DB> {
    type Rejection = Response;

    async fn from_request_parts(
        parts: &mut Parts,
        state: &Arc<AthenaAuth<DB>>,
    ) -> Result<Self, Self::Rejection> {
        match CurrentSession::from_request_parts(parts, state).await {
            Ok(session) => Ok(OptionalSession(Some(session))),
            Err(_) => Ok(OptionalSession(None)),
        }
    }
}

#[cfg(all(test, feature = "axum"))]
mod tests {
    use super::*;
    use axum::body::Body;
    use axum::http::Request;
    use std::collections::HashMap;

    #[tokio::test]
    async fn convert_axum_request_does_not_override_existing_proxy_header() {
        let mut req = Request::builder()
            .method("POST")
            .uri("/api/auth/sign-up/email")
            .header("x-forwarded-for", "203.0.113.10")
            .body(Body::from("{}"))
            .expect("request must build");
        req.extensions_mut()
            .insert(ConnectInfo(SocketAddr::from(([198, 51, 100, 20], 443))));

        let converted = convert_axum_request(req)
            .await
            .expect("conversion must succeed");
        assert_eq!(
            converted.headers.get("x-forwarded-for").map(String::as_str),
            Some("203.0.113.10")
        );
        assert!(
            !converted
                .headers
                .contains_key(INTERNAL_AXUM_REMOTE_IP_HEADER),
            "internal fallback header must not be injected when external proxy header already exists"
        );
    }

    #[tokio::test]
    async fn convert_axum_request_injects_internal_fallback_ip_header_from_connect_info() {
        let mut req = Request::builder()
            .method("POST")
            .uri("/api/auth/sign-up/email")
            .body(Body::from("{}"))
            .expect("request must build");
        req.extensions_mut()
            .insert(ConnectInfo(SocketAddr::from(([198, 51, 100, 21], 443))));

        let converted = convert_axum_request(req)
            .await
            .expect("conversion must succeed");
        assert_eq!(
            converted
                .headers
                .get(INTERNAL_AXUM_REMOTE_IP_HEADER)
                .map(String::as_str),
            Some("198.51.100.21")
        );
    }

    #[test]
    fn convert_auth_response_falls_back_to_500_for_invalid_status_code() {
        let response = convert_auth_response(AuthResponse {
            status: 0,
            headers: HashMap::new(),
            body: b"payload".to_vec(),
        });

        assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
    }

    #[tokio::test]
    async fn convert_auth_response_ignores_invalid_header_entries() {
        let mut headers = HashMap::new();
        headers.insert("x-valid-header".to_string(), "ok".to_string());
        headers.insert("bad header".to_string(), "skip".to_string());
        headers.insert("x-invalid-value".to_string(), "\n".to_string());
        let response = convert_auth_response(AuthResponse {
            status: 200,
            headers,
            body: b"ok".to_vec(),
        });

        assert_eq!(response.status(), StatusCode::OK);
        assert_eq!(
            response
                .headers()
                .get("x-valid-header")
                .and_then(|value| value.to_str().ok()),
            Some("ok")
        );
        assert!(!response.headers().contains_key("bad header"));
        assert!(!response.headers().contains_key("x-invalid-value"));

        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
            .await
            .expect("body should be readable");
        assert_eq!(body.as_ref(), b"ok");
    }
}
