pub mod api_tokens;
pub mod capability_guard;
pub mod capability_sets;
pub mod jwt;
pub mod passkeys;
pub mod password;

use std::sync::{Arc, Mutex};

use axum::{async_trait, extract::FromRequestParts, http::request::Parts};
use axum_extra::headers::{authorization::Bearer, Authorization};
use axum_extra::TypedHeader;
use serde::{Deserialize, Serialize};
use utoipa::ToSchema;

use crate::{
    auth::capability_sets::{load_capabilities_for_set, load_capability_set},
    error::AppError,
    models::ApiCapability,
    state::{AppState, PgPooledConnection},
};
use uuid::Uuid;

use crate::auth::jwt::PrincipalKind;

#[derive(Clone)]
pub struct TenantConnectionHolder {
    inner: Arc<Mutex<Option<PgPooledConnection>>>,
}

impl TenantConnectionHolder {
    pub fn new(conn: PgPooledConnection) -> Self {
        Self {
            inner: Arc::new(Mutex::new(Some(conn))),
        }
    }

    pub fn into_conn(self) -> Option<PgPooledConnection> {
        self.inner.lock().ok()?.take()
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
pub struct AuthenticatedUser {
    pub user_id: uuid::Uuid,
    pub username: String,
    pub tenant_id: uuid::Uuid,
    pub principal_kind: PrincipalKind,
    pub principal_id: Uuid,
    pub capability_set_id: Uuid,
    pub cap_version: i32,
    pub capabilities: Vec<ApiCapability>,
}

#[async_trait]
impl FromRequestParts<AppState> for AuthenticatedUser {
    type Rejection = AppError;

    async fn from_request_parts(
        parts: &mut Parts,
        state: &AppState,
    ) -> Result<Self, Self::Rejection> {
        if let Some(user) = parts.extensions.get::<AuthenticatedUser>() {
            return Ok(user.clone());
        }

        let TypedHeader(Authorization(bearer)) =
            TypedHeader::<Authorization<Bearer>>::from_request_parts(parts, state)
                .await
                .map_err(|_| AppError::unauthorized())?;

        let claims = state
            .jwt
            .verify_token(bearer.token())
            .map_err(|_| AppError::unauthorized())?;

        let mut tenant_conn = state.db_for_tenant(claims.tenant_id)?;
        let capability_set = load_capability_set(&mut tenant_conn, claims.capability_set_id)
            .map_err(|_| AppError::unauthorized())?;

        if capability_set.cap_version != claims.cap_version {
            return Err(AppError::unauthorized());
        }

        let capabilities = load_capabilities_for_set(&mut tenant_conn, capability_set.id)
            .map_err(|_| AppError::unauthorized())?;

        let user = AuthenticatedUser {
            user_id: claims.sub,
            username: claims.username,
            tenant_id: claims.tenant_id,
            principal_kind: claims.principal_kind,
            principal_id: claims.principal_id,
            capability_set_id: claims.capability_set_id,
            cap_version: claims.cap_version,
            capabilities,
        };

        parts.extensions.insert(user.clone());
        parts
            .extensions
            .insert(TenantConnectionHolder::new(tenant_conn));

        Ok(user)
    }
}

pub struct TenantScopedConn {
    pub conn: PgPooledConnection,
    pub tenant_id: Uuid,
    pub user_id: Uuid,
    pub user: AuthenticatedUser,
}

impl TenantScopedConn {
    pub fn conn(&mut self) -> &mut PgPooledConnection {
        &mut self.conn
    }
}

#[async_trait]
impl FromRequestParts<AppState> for TenantScopedConn {
    type Rejection = AppError;

    async fn from_request_parts(
        parts: &mut Parts,
        state: &AppState,
    ) -> Result<Self, Self::Rejection> {
        let user = AuthenticatedUser::from_request_parts(parts, state).await?;
        let tenant_id = user.tenant_id;
        let conn = if let Some(holder) = parts.extensions.remove::<TenantConnectionHolder>() {
            holder
                .into_conn()
                .ok_or_else(|| AppError::internal("tenant connection unavailable"))?
        } else {
            state.db_for_tenant(tenant_id)?
        };

        Ok(Self {
            conn,
            tenant_id,
            user_id: user.user_id,
            user,
        })
    }
}