//! Row-Level Security (RLS) middleware for PostgreSQL.

use axum::{
    extract::FromRequestParts,
    http::{request::Parts, Request, StatusCode},
    response::{IntoResponse, Response},
    Json,
};
use sqlx::{pool::PoolConnection, postgres::PgConnection, PgPool, Postgres};
use std::{
    future::Future,
    ops::{Deref, DerefMut},
    pin::Pin,
    sync::Arc,
    task::{Context, Poll},
};
use tokio::sync::{Mutex, MutexGuard};
use tower::{Layer, Service};
use tower_sessions::Session;
use uuid::Uuid;

use super::session::{SESSION_GUEST_ID_KEY, SESSION_USER_ID_KEY};
use chattyness_error::ErrorResponse;

// =============================================================================
// RLS Connection Wrapper
// =============================================================================

struct RlsConnectionInner {
    conn: Option<PoolConnection<Postgres>>,
    pool: PgPool,
}

impl Drop for RlsConnectionInner {
    fn drop(&mut self) {
        if let Some(mut conn) = self.conn.take() {
            let pool = self.pool.clone();
            tokio::spawn(async move {
                let _ = sqlx::query("SELECT public.set_current_user_id(NULL)")
                    .execute(&mut *conn)
                    .await;
                let _ = sqlx::query("SELECT public.set_current_guest_session_id(NULL)")
                    .execute(&mut *conn)
                    .await;
                drop(conn);
                drop(pool);
            });
        }
    }
}

/// A database connection with RLS user ID already set.
#[derive(Clone)]
pub struct RlsConnection {
    inner: Arc<Mutex<RlsConnectionInner>>,
}

impl RlsConnection {
    fn new(conn: PoolConnection<Postgres>, pool: PgPool) -> Self {
        Self {
            inner: Arc::new(Mutex::new(RlsConnectionInner {
                conn: Some(conn),
                pool,
            })),
        }
    }

    /// Acquire exclusive access to the RLS connection.
    pub async fn acquire(&self) -> RlsGuard<'_> {
        RlsGuard {
            guard: self.inner.lock().await,
        }
    }

    /// Set the current user ID on the RLS connection.
    /// Use this after creating a new user to set the context for subsequent operations.
    pub async fn set_user_id(&self, user_id: Uuid) -> Result<(), sqlx::Error> {
        let mut guard = self.inner.lock().await;
        let conn = guard.conn.as_mut().expect("RlsConnection already consumed");
        sqlx::query("SELECT public.set_current_user_id($1)")
            .bind(user_id)
            .execute(&mut **conn)
            .await?;
        Ok(())
    }
}

/// A guard providing mutable access to an RLS-configured database connection.
pub struct RlsGuard<'a> {
    guard: MutexGuard<'a, RlsConnectionInner>,
}

impl Deref for RlsGuard<'_> {
    type Target = PgConnection;

    fn deref(&self) -> &Self::Target {
        self.guard
            .conn
            .as_ref()
            .expect("RlsConnection already consumed")
            .deref()
    }
}

impl DerefMut for RlsGuard<'_> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.guard
            .conn
            .as_mut()
            .expect("RlsConnection already consumed")
            .deref_mut()
    }
}

// =============================================================================
// RLS Connection Extractor
// =============================================================================

/// Extractor for an RLS-enabled database connection.
pub struct RlsConn(pub RlsConnection);

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

    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
        parts
            .extensions
            .remove::<RlsConnection>()
            .map(RlsConn)
            .ok_or(RlsError::NoConnection)
    }
}

/// Errors related to RLS connection handling.
#[derive(Debug)]
pub enum RlsError {
    NoConnection,
    DatabaseError(String),
}

impl IntoResponse for RlsError {
    fn into_response(self) -> Response {
        let (status, message) = match self {
            RlsError::NoConnection => (
                StatusCode::INTERNAL_SERVER_ERROR,
                "RLS connection not available",
            ),
            RlsError::DatabaseError(msg) => {
                (StatusCode::INTERNAL_SERVER_ERROR, msg.leak() as &'static str)
            }
        };

        let body = ErrorResponse {
            error: message.to_string(),
            code: Some("RLS_ERROR".to_string()),
        };

        (status, Json(body)).into_response()
    }
}

// =============================================================================
// RLS Middleware Layer
// =============================================================================

/// Layer that provides RLS-enabled database connections per request.
#[derive(Clone)]
pub struct RlsLayer {
    pool: PgPool,
}

impl RlsLayer {
    pub fn new(pool: PgPool) -> Self {
        Self { pool }
    }
}

impl<S> Layer<S> for RlsLayer {
    type Service = RlsMiddleware<S>;

    fn layer(&self, inner: S) -> Self::Service {
        RlsMiddleware {
            inner,
            pool: self.pool.clone(),
        }
    }
}

/// Middleware that sets up RLS connections per request.
#[derive(Clone)]
pub struct RlsMiddleware<S> {
    inner: S,
    pool: PgPool,
}

impl<S, B> Service<Request<B>> for RlsMiddleware<S>
where
    S: Service<Request<B>, Response = Response> + Clone + Send + 'static,
    S::Future: Send,
    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, mut request: Request<B>) -> Self::Future {
        let pool = self.pool.clone();
        let mut inner = self.inner.clone();

        let session = request.extensions().get::<Session>().cloned();

        Box::pin(async move {
            let (user_id, guest_session_id) = get_session_ids(session).await;

            match acquire_rls_connection(&pool, user_id, guest_session_id).await {
                Ok(rls_conn) => {
                    request.extensions_mut().insert(rls_conn);
                    inner.call(request).await
                }
                Err(e) => {
                    tracing::error!("Failed to acquire RLS connection: {}", e);
                    Ok(RlsError::DatabaseError(e.to_string()).into_response())
                }
            }
        })
    }
}

async fn get_session_ids(session: Option<Session>) -> (Option<Uuid>, Option<Uuid>) {
    let Some(session) = session else {
        return (None, None);
    };

    let user_id = session
        .get::<Uuid>(SESSION_USER_ID_KEY)
        .await
        .ok()
        .flatten();
    let guest_session_id = session
        .get::<Uuid>(SESSION_GUEST_ID_KEY)
        .await
        .ok()
        .flatten();

    (user_id, guest_session_id)
}

async fn acquire_rls_connection(
    pool: &PgPool,
    user_id: Option<Uuid>,
    guest_session_id: Option<Uuid>,
) -> Result<RlsConnection, sqlx::Error> {
    let mut conn = pool.acquire().await?;

    if user_id.is_some() {
        sqlx::query("SELECT public.set_current_user_id($1)")
            .bind(user_id)
            .execute(&mut *conn)
            .await?;
    } else if guest_session_id.is_some() {
        sqlx::query("SELECT public.set_current_user_id(NULL)")
            .execute(&mut *conn)
            .await?;
        sqlx::query("SELECT public.set_current_guest_session_id($1)")
            .bind(guest_session_id)
            .execute(&mut *conn)
            .await?;
    } else {
        sqlx::query("SELECT public.set_current_user_id(NULL)")
            .execute(&mut *conn)
            .await?;
    }

    Ok(RlsConnection::new(conn, pool.clone()))
}

impl std::ops::Deref for RlsConn {
    type Target = RlsConnection;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl std::ops::DerefMut for RlsConn {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}
