use std::sync::Arc;

use diesel::{
    pg::PgConnection,
    r2d2::{ConnectionManager, PooledConnection},
};
use uuid::Uuid;

use crate::{
    auth::jwt::JwtService,
    config::AppConfig,
    db::PgPool,
    error::{AppError, AppResult},
    storage::{ObjectStorage, TenantStorage},
    tenants::{apply_tenant_guc, TenantService},
};

pub type PgPooledConnection = PooledConnection<ConnectionManager<PgConnection>>;

#[derive(Clone)]
pub struct AppState {
    pub pool: PgPool,
    pub config: Arc<AppConfig>,
    storage: Arc<dyn ObjectStorage>,
    pub jwt: JwtService,
    pub tenants: TenantService,
}

impl AppState {
    pub async fn initialize(
        config: AppConfig,
        pool_size_override: Option<u32>,
    ) -> anyhow::Result<Self> {
        let pool_size = pool_size_override.unwrap_or(config.database_max_pool_size);
        let pool = crate::db::init_pool_with_size(&config.database_url, pool_size)?;
        let s3_client = crate::s3::build_client(&config).await?;
        let storage = Arc::new(crate::storage::S3Storage::new(
            s3_client,
            config.s3_bucket.clone(),
        ));
        let jwt = crate::auth::jwt::JwtService::from_config(&config)?;

        Ok(Self::new(pool, config, storage, jwt))
    }

    pub fn new(
        pool: PgPool,
        config: AppConfig,
        storage: Arc<dyn ObjectStorage>,
        jwt: JwtService,
    ) -> Self {
        let config = Arc::new(config);
        let tenants = TenantService::new(pool.clone());

        Self {
            pool,
            config,
            storage,
            jwt,
            tenants,
        }
    }

    pub fn db_for_tenant(&self, tenant_id: Uuid) -> AppResult<PgPooledConnection> {
        debug_assert!(!tenant_id.is_nil(), "nil tenant_id passed to db_for_tenant");
        let mut conn = self.db_unscoped()?;
        apply_tenant_guc(&mut conn, tenant_id)?;
        Ok(conn)
    }

    pub(crate) fn db_unscoped(&self) -> AppResult<PgPooledConnection> {
        self.pool
            .get()
            .map_err(|err| AppError::internal(format!("database pool error: {err}")))
    }

    pub fn storage_for_tenant(&self, tenant_id: Uuid) -> AppResult<TenantStorage> {
        let tenant = self.tenants.get_by_id(tenant_id)?;
        TenantStorage::new(self.storage.clone(), &tenant)
            .map_err(|err| AppError::internal(format!("tenant storage error: {err}")))
    }
}