neiam /dms

use axum::{
    extract::{FromRequestParts, OptionalFromRequestParts, Request, State},
    http::{header::{AUTHORIZATION, COOKIE}, request::Parts, StatusCode},
    middleware::Next,
    response::{IntoResponse, Response, Redirect},
};
use std::convert::Infallible;
use sea_orm::DatabaseConnection;
use std::sync::Arc;
use crate::auth;
use crate::entities::{users, user_sessions};

/// User context for authenticated requests
#[derive(Clone, Debug)]
pub struct UserContext {
    pub user: users::Model,
    pub session: user_sessions::Model,
}

impl<S> FromRequestParts<S> for UserContext
where
    S: Send + Sync,
{
    type Rejection = StatusCode;

    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
        parts
            .extensions
            .get::<UserContext>()
            .cloned()
            .ok_or(StatusCode::UNAUTHORIZED)
    }
}

impl<S> OptionalFromRequestParts<S> for UserContext
where
    S: Send + Sync,
{
    type Rejection = Infallible;

    async fn from_request_parts(
        parts: &mut Parts,
        _state: &S,
    ) -> Result<Option<Self>, Self::Rejection> {
        Ok(parts.extensions.get::<UserContext>().cloned())
    }
}

/// Extract session token from cookie
fn extract_session_token(request: &Request) -> Option<String> {
    let cookies = request
        .headers()
        .get(COOKIE)?
        .to_str()
        .ok()?;

    // Parse cookies to find session_token
    for cookie in cookies.split(';') {
        let cookie = cookie.trim();
        if let Some(session_token) = cookie.strip_prefix("session_token=") {
            return Some(session_token.to_string());
        }
    }
    None
}

/// Extract Bearer token from Authorization header.
fn extract_bearer_token(request: &Request) -> Option<String> {
    let header = request.headers().get(AUTHORIZATION)?.to_str().ok()?;
    let stripped = header.strip_prefix("Bearer ").or_else(|| header.strip_prefix("bearer "))?;
    let token = stripped.trim();
    if token.is_empty() { None } else { Some(token.to_string()) }
}

/// Authentication middleware. Accepts either a session cookie (browser) or
/// an `Authorization: Bearer <api_token>` header (Android, scripts).
pub async fn auth_middleware(
    State(db): State<Arc<DatabaseConnection>>,
    mut request: Request,
    next: Next,
) -> Result<Response, StatusCode> {
    if let Some(bearer) = extract_bearer_token(&request) {
        match auth::validate_api_token(&db, &bearer).await {
            Ok(Some(user)) => {
                // Synthesize a session-shaped context so handlers don't need
                // to care which credential type was used.
                let synthetic = synthetic_session_for(&user);
                let user_context = UserContext { user, session: synthetic };
                request.extensions_mut().insert(user_context);
                return Ok(next.run(request).await);
            }
            Ok(None) => return Err(StatusCode::UNAUTHORIZED),
            Err(_) => return Err(StatusCode::INTERNAL_SERVER_ERROR),
        }
    }

    let session_token = extract_session_token(&request);
    if let Some(token) = session_token {
        match auth::validate_session(&db, &token).await {
            Ok(Some((user, session))) => {
                let user_context = UserContext { user, session };
                request.extensions_mut().insert(user_context);
                Ok(next.run(request).await)
            }
            Ok(None) => Err(StatusCode::UNAUTHORIZED),
            Err(_) => Err(StatusCode::INTERNAL_SERVER_ERROR),
        }
    } else {
        Err(StatusCode::UNAUTHORIZED)
    }
}

/// Build a placeholder user_sessions::Model for Bearer-authed requests so
/// downstream code that reads UserContext.session works uniformly. The fake
/// session is never persisted.
fn synthetic_session_for(user: &users::Model) -> user_sessions::Model {
    let now: chrono::DateTime<chrono::FixedOffset> = chrono::Utc::now().into();
    user_sessions::Model {
        id: uuid::Uuid::nil(),
        user_id: user.id,
        session_token: String::from("api-token"),
        expires_at: now,
        created_at: now,
        last_used_at: now,
    }
}

/// Optional authentication middleware - doesn't fail if no auth provided
pub async fn optional_auth_middleware(
    State(db): State<Arc<DatabaseConnection>>,
    mut request: Request,
    next: Next,
) -> Response {
    let session_token = extract_session_token(&request);

    if let Some(token) = session_token
        && let Ok(Some((user, session))) = auth::validate_session(&db, &token).await
    {
        let user_context = UserContext { user, session };
        request.extensions_mut().insert(user_context);
    }

    next.run(request).await
}

/// Redirect to login middleware for web routes
pub async fn web_auth_middleware(
    State(db): State<Arc<DatabaseConnection>>,
    mut request: Request,
    next: Next,
) -> Response {
    let session_token = extract_session_token(&request);

    if let Some(token) = session_token {
        match auth::validate_session(&db, &token).await {
            Ok(Some((user, session))) => {
                let user_context = UserContext { user, session };
                request.extensions_mut().insert(user_context);
                next.run(request).await
            }
            Ok(None) | Err(_) => {
                // Redirect to login page
                Redirect::temporary("/login").into_response()
            }
        }
    } else {
        // Redirect to login page
        Redirect::temporary("/login").into_response()
    }
}

/// Extension trait to get user context from request
pub trait RequestUserExt {
    fn user_context(&self) -> Option<&UserContext>;
    fn current_user(&self) -> Option<&users::Model>;
    fn current_session(&self) -> Option<&user_sessions::Model>;
}

impl RequestUserExt for Request {
    fn user_context(&self) -> Option<&UserContext> {
        self.extensions().get::<UserContext>()
    }

    fn current_user(&self) -> Option<&users::Model> {
        self.user_context().map(|ctx| &ctx.user)
    }

    fn current_session(&self) -> Option<&user_sessions::Model> {
        self.user_context().map(|ctx| &ctx.session)
    }
}