eric / rustwarden

//! Locale middleware for Axum
//!
//! This module provides middleware to inject locale into request extensions.

use axum::{extract::Request, response::Response};
use sea_orm::DatabaseConnection;
use std::{
    future::Future,
    pin::Pin,
    sync::Arc,
    task::{Context, Poll},
};
use tower::Service;

use crate::auth::jwt::Claims;
use crate::i18n::extractor::Locale;
use crate::i18n::service::I18nService;
use crate::services::user_preferences_service::UserPreferencesService;

/// Middleware service that injects locale into request extensions
///
/// Priority order for locale detection:
/// 1. User preference from database (if authenticated)
/// 2. Cookie `locale` (if present)
/// 3. Accept-Language header (only if not "fr", to respect French as default)
/// 4. Default: "fr" (FR-001: French is the default language)
#[derive(Clone)]
pub struct LocaleMiddleware<S> {
    inner: S,
    db: Arc<DatabaseConnection>,
    i18n_service: Arc<I18nService>,
}

impl<S> LocaleMiddleware<S> {
    /// Create a new `LocaleMiddleware` wrapping the given inner service.
    pub fn new(inner: S, db: Arc<DatabaseConnection>, i18n_service: Arc<I18nService>) -> Self {
        Self {
            inner,
            db,
            i18n_service,
        }
    }
}

impl<S, B> Service<Request<B>> for LocaleMiddleware<S>
where
    S: Service<Request<B>, Response = Response> + Clone + Send + 'static,
    S::Future: Send + 'static,
    B: Send + 'static,
{
    type Response = Response;
    type Error = S::Error;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, req: Request<B>) -> Self::Future {
        let db = self.db.clone();
        let i18n_service = self.i18n_service.clone();
        let mut inner = self.inner.clone();

        Box::pin(async move {
            // Extract locale from request
            let (parts, body) = Request::into_parts(req);
            let cookie_jar = axum_extra::extract::CookieJar::from_headers(&parts.headers);
            let has_locale_cookie = cookie_jar.get("locale").is_some();

            // Priority 1: Check user preference from database (if authenticated)
            let (locale, should_set_cookie) = if let Some(claims) = parts.extensions.get::<Claims>()
            {
                // User is authenticated, try to get preference from database
                match UserPreferencesService::get_preferences(&db, claims.sub).await {
                    Ok(Some(prefs))
                        if matches!(
                            prefs.language.as_str(),
                            "fr" | "en" | "de" | "es" | "it" | "pt"
                        ) =>
                    {
                        (Locale(prefs.language.clone()), false) // Don't set cookie for authenticated users
                    }
                    _ => {
                        // No preference or invalid, fall through to cookie/header
                        let (loc, should_set) =
                            extract_locale_from_request(&parts, has_locale_cookie);
                        (loc, should_set)
                    }
                }
            } else {
                // Not authenticated, use cookie/header
                // Set cookie if detected from Accept-Language header and no cookie exists
                let (loc, should_set) = extract_locale_from_request(&parts, has_locale_cookie);
                (loc, should_set)
            };

            // Inject locale and I18nService into extensions
            let mut req = Request::from_parts(parts, body);
            req.extensions_mut().insert(locale.clone());
            req.extensions_mut().insert(i18n_service);

            // Call inner service and modify response to add cookie if needed
            let response = inner.call(req).await?;

            // If we should set a cookie (non-authenticated user, detected from header, no cookie exists)
            if should_set_cookie {
                let (mut parts, body) = response.into_parts();

                // Create locale cookie (1 year expiration)
                use time::Duration as TimeDuration;
                let cookie =
                    axum_extra::extract::cookie::Cookie::build(("locale", locale.as_str()))
                        .http_only(false) // Allow JavaScript access for client-side language switching
                        .same_site(axum_extra::extract::cookie::SameSite::Lax)
                        .path("/")
                        .max_age(TimeDuration::seconds(31536000)); // 1 year

                // Add cookie to response headers
                if let Ok(header_value) = cookie.to_string().parse() {
                    parts
                        .headers
                        .insert(axum::http::header::SET_COOKIE, header_value);
                } else {
                    tracing::error!("Failed to create locale cookie header");
                }

                Ok(axum::response::Response::from_parts(parts, body))
            } else {
                Ok(response)
            }
        })
    }
}

/// Extract locale from cookie or Accept-Language header
///
/// Returns (Locale, should_set_cookie):
/// - should_set_cookie is true if locale was detected from Accept-Language header
///   and no cookie exists (for non-authenticated users)
fn extract_locale_from_request(
    parts: &axum::http::request::Parts,
    has_locale_cookie: bool,
) -> (Locale, bool) {
    // Priority 2: Check cookie
    if let Some(cookie) = axum_extra::extract::CookieJar::from_headers(&parts.headers).get("locale")
    {
        let locale_str = cookie.value().to_string();
        if matches!(locale_str.as_str(), "fr" | "en" | "de" | "es" | "it" | "pt") {
            return (Locale(locale_str), false); // Cookie exists, don't set it again
        }
    }

    // Priority 3: Check Accept-Language header
    // NOTE: We only use Accept-Language if it matches a supported language
    // If the browser sends "de" but user hasn't explicitly chosen it, we should default to "fr"
    // So we only use Accept-Language if there's no cookie AND it's not the default "fr"
    if let Some(accept_lang) = parts.headers.get("accept-language") {
        if let Ok(accept_lang_str) = accept_lang.to_str() {
            // Simple parsing: take first language code
            if let Some(locale_str) = accept_lang_str
                .split(',')
                .next()
                .and_then(|s| s.split(';').next())
                .map(|s| {
                    let trimmed = s.trim();
                    if trimmed.len() >= 2 {
                        trimmed[..2].to_string()
                    } else {
                        "fr".to_string() // Default to French
                    }
                })
                .filter(|s| matches!(s.as_str(), "fr" | "en" | "de" | "es" | "it" | "pt"))
            {
                // Only use Accept-Language if it's not the default "fr"
                // This ensures that users without explicit preference get French by default
                if locale_str != "fr" {
                    // Set cookie if detected from header and no cookie exists
                    let should_set = !has_locale_cookie;
                    return (Locale(locale_str), should_set);
                }
                // If Accept-Language is "fr", fall through to default (no cookie set)
            }
        }
    }

    // Priority 4: Default to French (FR-001: French is the default language)
    (Locale("fr".to_string()), false)
}