mod query_logger;

use std::{
    future::Future,
    sync::{
        Arc, OnceLock,
        atomic::{AtomicBool, Ordering},
    },
    time::Duration,
};

use diesel::{
    Connection, RunQueryDsl,
    connection::SimpleConnection,
    r2d2::{CustomizeConnection, Pool, PooledConnection},
};
use rocket::{
    Request,
    http::Status,
    request::{FromRequest, Outcome},
};
use tokio::{
    sync::{Mutex, OwnedSemaphorePermit, Semaphore},
    time::timeout,
};

use crate::{
    CONFIG,
    error::{Error, MapResult},
};

// These changes are based on Rocket 0.5-rc wrapper of Diesel: https://github.com/SergioBenitez/Rocket/blob/v0.5-rc/contrib/sync_db_pools
// A wrapper around spawn_blocking that propagates panics to the calling code.
pub async fn run_blocking<F, R>(job: F) -> R
where
    F: FnOnce() -> R + Send + 'static,
    R: Send + 'static,
{
    match tokio::task::spawn_blocking(job).await {
        Ok(ret) => ret,
        Err(e) => match e.try_into_panic() {
            Ok(panic) => std::panic::resume_unwind(panic),
            Err(_) => unreachable!("spawn_blocking tasks are never cancelled"),
        },
    }
}

// This is used to generate the main DbConn and DbPool enums, which contain one variant for each database supported
#[derive(diesel::MultiConnection)]
pub enum DbConnInner {
    #[cfg(mysql)]
    Mysql(diesel::mysql::MysqlConnection),
    #[cfg(postgresql)]
    Postgresql(diesel::pg::PgConnection),
    #[cfg(sqlite)]
    Sqlite(diesel::sqlite::SqliteConnection),
}

/// Custom connection manager that implements manual connection establishment
pub struct DbConnManager {
    database_url: String,
}

impl DbConnManager {
    pub fn new(database_url: &str) -> Self {
        Self {
            database_url: database_url.to_owned(),
        }
    }

    fn establish_connection(&self) -> Result<DbConnInner, diesel::r2d2::Error> {
        match DbConnType::from_url(&self.database_url) {
            #[cfg(mysql)]
            Ok(DbConnType::Mysql) => {
                let conn = diesel::mysql::MysqlConnection::establish(&self.database_url)?;
                Ok(DbConnInner::Mysql(conn))
            }
            #[cfg(postgresql)]
            Ok(DbConnType::Postgresql) => {
                let conn = diesel::pg::PgConnection::establish(&self.database_url)?;
                Ok(DbConnInner::Postgresql(conn))
            }
            #[cfg(sqlite)]
            Ok(DbConnType::Sqlite) => {
                let conn = diesel::sqlite::SqliteConnection::establish(&self.database_url)?;
                Ok(DbConnInner::Sqlite(conn))
            }

            Err(e) => Err(diesel::r2d2::Error::ConnectionError(diesel::ConnectionError::InvalidConnectionUrl(
                format!("Unable to estabilsh a connection: {e:?}"),
            ))),
        }
    }
}

impl diesel::r2d2::ManageConnection for DbConnManager {
    type Connection = DbConnInner;
    type Error = diesel::r2d2::Error;

    fn connect(&self) -> Result<Self::Connection, Self::Error> {
        self.establish_connection()
    }

    fn is_valid(&self, conn: &mut Self::Connection) -> Result<(), Self::Error> {
        use diesel::r2d2::R2D2Connection;
        conn.ping().map_err(diesel::r2d2::Error::QueryError)
    }

    fn has_broken(&self, conn: &mut Self::Connection) -> bool {
        use diesel::r2d2::R2D2Connection;
        conn.is_broken()
    }
}

#[derive(Eq, PartialEq)]
pub enum DbConnType {
    #[cfg(mysql)]
    Mysql,
    #[cfg(postgresql)]
    Postgresql,
    #[cfg(sqlite)]
    Sqlite,
}

pub static ACTIVE_DB_TYPE: OnceLock<DbConnType> = OnceLock::new();

pub struct DbConn {
    conn: Arc<Mutex<Option<PooledConnection<DbConnManager>>>>,
    permit: Option<OwnedSemaphorePermit>,
    transaction_open: AtomicBool,
}

#[derive(Debug)]
pub struct DbConnOptions {
    pub init_stmts: String,
}

impl CustomizeConnection<DbConnInner, diesel::r2d2::Error> for DbConnOptions {
    fn on_acquire(&self, conn: &mut DbConnInner) -> Result<(), diesel::r2d2::Error> {
        if !self.init_stmts.is_empty() {
            conn.batch_execute(&self.init_stmts).map_err(diesel::r2d2::Error::QueryError)?;
        }
        Ok(())
    }
}

#[derive(Clone)]
pub struct DbPool {
    // This is an 'Option' so that we can drop the pool in a 'spawn_blocking'.
    pool: Option<Pool<DbConnManager>>,
    semaphore: Arc<Semaphore>,
}

impl Drop for DbConn {
    fn drop(&mut self) {
        let conn = Arc::clone(&self.conn);
        let permit = self.permit.take();
        let transaction_open = self.transaction_open.load(Ordering::Acquire);

        // Since connection can't be on the stack in an async fn during an
        // await, we have to spawn a new blocking-safe thread...
        tokio::task::spawn_blocking(move || {
            // And then re-enter the runtime to wait on the async mutex, but in a blocking fashion.
            let mut conn = tokio::runtime::Handle::current().block_on(conn.lock_owned());

            if let Some(mut conn) = conn.take() {
                if transaction_open {
                    // A cancelled request must not return an open transaction to the pool.
                    if let Err(error) = conn.batch_execute("ROLLBACK") {
                        error!("Failed to roll back a cancelled database transaction: {error}");
                    }
                }
                drop(conn);
            }

            // Drop permit after the connection is dropped
            drop(permit);
        });
    }
}

impl Drop for DbPool {
    fn drop(&mut self) {
        let pool = self.pool.take();
        // Only use spawn_blocking if the Tokio runtime is still available
        // Otherwise the pool will be dropped on the current thread
        if let Ok(handle) = tokio::runtime::Handle::try_current() {
            handle.spawn_blocking(move || drop(pool));
        }
    }
}

impl DbPool {
    // For the given database URL, guess its type, run migrations, create pool, and return it
    pub fn from_config() -> Result<Self, Error> {
        let db_url = CONFIG.database_url();
        let conn_type = DbConnType::from_url(&db_url)?;

        // Only set the default instrumentation if the log level is specifically set to either warn, info or debug
        if log_enabled!(target: "vaultwarden::db::query_logger", log::Level::Warn)
            || log_enabled!(target: "vaultwarden::db::query_logger", log::Level::Info)
            || log_enabled!(target: "vaultwarden::db::query_logger", log::Level::Debug)
        {
            drop(diesel::connection::set_default_instrumentation(query_logger::simple_logger));
        }

        match conn_type {
            #[cfg(mysql)]
            DbConnType::Mysql => {
                mysql_migrations::run_migrations(&db_url)?;
            }
            #[cfg(postgresql)]
            DbConnType::Postgresql => {
                postgresql_migrations::run_migrations(&db_url)?;
            }
            #[cfg(sqlite)]
            DbConnType::Sqlite => {
                sqlite_migrations::run_migrations(&db_url)?;
            }
        }

        let max_conns = CONFIG.database_max_conns();
        let manager = DbConnManager::new(&db_url);
        let pool = Pool::builder()
            .max_size(max_conns)
            .min_idle(Some(CONFIG.database_min_conns()))
            .idle_timeout(Some(Duration::from_secs(u64::from(CONFIG.database_idle_timeout()))))
            .connection_timeout(Duration::from_secs(u64::from(CONFIG.database_timeout())))
            .connection_customizer(Box::new(DbConnOptions {
                init_stmts: conn_type.get_init_stmts(),
            }))
            .build(manager)
            .map_res("Failed to create pool")?;

        // Set a global to determine the database more easily throughout the rest of the code
        if ACTIVE_DB_TYPE.set(conn_type).is_err() {
            error!("Tried to set the active database connection type more than once.");
        }

        Ok(DbPool {
            pool: Some(pool),
            semaphore: Arc::new(Semaphore::new(max_conns as usize)),
        })
    }

    // Get a connection from the pool
    pub async fn get(&self) -> Result<DbConn, Error> {
        let duration = Duration::from_secs(u64::from(CONFIG.database_timeout()));
        let permit = match timeout(duration, Arc::clone(&self.semaphore).acquire_owned()).await {
            Ok(p) => p.expect("Semaphore should be open"),
            Err(_) => {
                err!("Timeout waiting for database connection");
            }
        };

        let p = self.pool.as_ref().expect("DbPool.pool should always be Some()");
        let pool = p.clone();
        let c =
            run_blocking(move || pool.get_timeout(duration)).await.map_res("Error retrieving connection from pool")?;
        Ok(DbConn {
            conn: Arc::new(Mutex::new(Some(c))),
            permit: Some(permit),
            transaction_open: AtomicBool::new(false),
        })
    }
}

impl DbConnType {
    pub fn from_url(url: &str) -> Result<Self, Error> {
        // Mysql
        if url.len() > 6 && &url[..6] == "mysql:" {
            #[cfg(mysql)]
            return Ok(DbConnType::Mysql);

            #[cfg(not(mysql))]
            err!("`DATABASE_URL` is a MySQL URL, but the 'mysql' feature is not enabled")

        // Postgresql
        } else if url.len() > 11 && (&url[..11] == "postgresql:" || &url[..9] == "postgres:") {
            #[cfg(postgresql)]
            return Ok(DbConnType::Postgresql);

            #[cfg(not(postgresql))]
            err!("`DATABASE_URL` is a PostgreSQL URL, but the 'postgresql' feature is not enabled")

        // Sqlite (explicit)
        } else if url.len() > 7 && &url[..7] == "sqlite:" {
            #[cfg(sqlite)]
            return Ok(DbConnType::Sqlite);

            #[cfg(not(sqlite))]
            err!("`DATABASE_URL` is a SQLite URL, but the 'sqlite' feature is not enabled")
        }

        // No recognized scheme — assume legacy bare-path SQLite, but the database file must already exist.
        // This prevents misconfigured URLs (typos, quoted strings) from silently creating a new empty SQLite database.
        #[cfg(sqlite)]
        {
            if std::path::Path::new(url).exists() {
                return Ok(DbConnType::Sqlite);
            }
            err!(format!(
                "`DATABASE_URL` does not match any known database scheme (mysql://, postgresql://, sqlite://) \
                    and no existing SQLite database was found at '{url}'. \
                    If you intend to use SQLite, use an explicit `sqlite://` scheme in your `DATABASE_URL`. \
                    Otherwise, check your DATABASE_URL for typos or quoting issues."
            ))
        }

        #[cfg(not(sqlite))]
        err!("`DATABASE_URL` does not match any known database scheme (mysql://, postgresql://, sqlite://)")
    }

    pub fn get_init_stmts(&self) -> String {
        let init_stmts = CONFIG.database_conn_init();
        if init_stmts.is_empty() {
            self.default_init_stmts()
        } else {
            init_stmts
        }
    }

    pub fn default_init_stmts(&self) -> String {
        match self {
            #[cfg(mysql)]
            Self::Mysql => String::new(),
            #[cfg(postgresql)]
            Self::Postgresql => String::new(),
            #[cfg(sqlite)]
            Self::Sqlite => "PRAGMA busy_timeout = 5000; PRAGMA synchronous = NORMAL;".to_owned(),
        }
    }
}

impl DbConn {
    /// Run asynchronous model calls on this request's pooled connection in one transaction.
    /// Send external notifications only after this returns successfully.
    pub async fn transaction<F, Fut, T>(&self, operation: F) -> Result<T, Error>
    where
        F: FnOnce() -> Fut,
        Fut: Future<Output = Result<T, Error>>,
    {
        self.run(|conn| conn.batch_execute("BEGIN")).await.map_res("Error starting database transaction")?;
        self.transaction_open.store(true, Ordering::Release);
        let result = operation().await;
        if result.is_ok() {
            let committed = self.run(|conn| conn.batch_execute("COMMIT")).await;
            let closed = committed.is_ok() || self.run(|conn| conn.batch_execute("ROLLBACK")).await.is_ok();
            if closed {
                self.transaction_open.store(false, Ordering::Release);
            }
            committed.map_res("Error committing database transaction")?;
        } else if self.run(|conn| conn.batch_execute("ROLLBACK")).await.is_ok() {
            self.transaction_open.store(false, Ordering::Release);
        }
        result
    }

    pub async fn run<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&mut DbConnInner) -> R + Send,
        R: Send + 'static,
    {
        let conn = Arc::clone(&self.conn);
        let mut conn = conn.lock_owned().await;
        let conn = conn.as_mut().expect("Internal invariant broken: self.conn is Some");

        // Run blocking can't be used due to the 'static limitation, use block_in_place instead
        tokio::task::block_in_place(move || f(conn))
    }
}

#[macro_export]
macro_rules! db_run {
    ( $conn:ident: $body:block ) => {
        $conn.run(move |$conn| $body).await
    };

    ( $conn:ident: $( $($db:ident),+ $body:block )+ ) => {
        $conn.run(move |$conn| {
            match $conn {
                $($(
                #[cfg($db)]
                pastey::paste!(&mut $crate::db::DbConnInner::[<$db:camel>](ref mut $conn)) => {
                    $body
                },
            )+)+}
        }).await
    };
}

// Write all ToSql<Text, DB> and FromSql<Text, DB> given a serializable/deserializable type.
#[macro_export]
macro_rules! impl_FromToSqlText {
    ($name:ty) => {
        #[cfg(mysql)]
        impl ToSql<Text, diesel::mysql::Mysql> for $name {
            fn to_sql<'b>(&'b self, out: &mut Output<'b, '_, diesel::mysql::Mysql>) -> diesel::serialize::Result {
                serde_json::to_writer(out, self).map(|_| diesel::serialize::IsNull::No).map_err(Into::into)
            }
        }

        #[cfg(postgresql)]
        impl ToSql<Text, diesel::pg::Pg> for $name {
            fn to_sql<'b>(&'b self, out: &mut Output<'b, '_, diesel::pg::Pg>) -> diesel::serialize::Result {
                serde_json::to_writer(out, self).map(|_| diesel::serialize::IsNull::No).map_err(Into::into)
            }
        }

        #[cfg(sqlite)]
        impl ToSql<Text, diesel::sqlite::Sqlite> for $name {
            fn to_sql<'b>(&'b self, out: &mut Output<'b, '_, diesel::sqlite::Sqlite>) -> diesel::serialize::Result {
                serde_json::to_string(self).map_err(Into::into).map(|str| {
                    out.set_value(str);
                    diesel::serialize::IsNull::No
                })
            }
        }

        impl<DB: diesel::backend::Backend> FromSql<Text, DB> for $name
        where
            String: FromSql<Text, DB>,
        {
            fn from_sql(bytes: DB::RawValue<'_>) -> diesel::deserialize::Result<Self> {
                <String as FromSql<Text, DB>>::from_sql(bytes)
                    .and_then(|str| serde_json::from_str(&str).map_err(Into::into))
            }
        }
    };
}

pub mod schema;

// Reexport the models, needs to be after the macros are defined so it can access them
pub mod models;

/// Creates a back-up of the sqlite database
/// MySQL/MariaDB and PostgreSQL are not supported.
#[cfg(sqlite)]
pub fn backup_sqlite() -> Result<String, Error> {
    use diesel::Connection;

    let db_url = CONFIG.database_url();
    if DbConnType::from_url(&CONFIG.database_url()).is_ok_and(|t| t == DbConnType::Sqlite) {
        // Strip the sqlite:// prefix if present to get the raw file path
        let file_path = db_url.strip_prefix("sqlite://").unwrap_or(&db_url);
        // Open a read-only connection for the backup
        let mut conn = diesel::sqlite::SqliteConnection::establish(&format!("sqlite://{file_path}?mode=ro"))?;

        let db_path = std::path::Path::new(file_path).parent().unwrap();
        let backup_file = db_path
            .join(format!("db_{}.sqlite3", chrono::Utc::now().format("%Y%m%d_%H%M%S")))
            .to_string_lossy()
            .into_owned();

        diesel::sql_query("VACUUM INTO ?")
            .bind::<diesel::sql_types::Text, _>(&backup_file)
            .execute(&mut conn)
            .map(|_| ())
            .map_res("VACUUM INTO failed")?;

        Ok(backup_file)
    } else {
        err_silent!("The database type is not SQLite. Backups only works for SQLite databases")
    }
}

#[cfg(not(sqlite))]
pub fn backup_sqlite() -> Result<String, Error> {
    err_silent!("The database type is not SQLite. Backups only works for SQLite databases")
}

/// Get the SQL Server version
pub async fn get_sql_server_version(conn: &DbConn) -> String {
    db_run! { conn:
        postgresql,mysql {
            diesel::select(diesel::dsl::sql::<diesel::sql_types::Text>("version();"))
            .get_result::<String>(conn)
            .unwrap_or_else(|_| "Unknown".to_owned())
        }
        sqlite {
            diesel::select(diesel::dsl::sql::<diesel::sql_types::Text>("sqlite_version();"))
            .get_result::<String>(conn)
            .unwrap_or_else(|_| "Unknown".to_owned())
        }
    }
}

/// Attempts to retrieve a single connection from the managed database pool. If
/// no pool is currently managed, fails with an `InternalServerError` status. If
/// no connections are available, fails with a `ServiceUnavailable` status.
#[rocket::async_trait]
impl<'r> FromRequest<'r> for DbConn {
    type Error = ();

    async fn from_request(request: &'r Request<'_>) -> Outcome<Self, Self::Error> {
        match request.rocket().state::<DbPool>() {
            Some(p) => match p.get().await {
                Ok(dbconn) => Outcome::Success(dbconn),
                _ => Outcome::Error((Status::ServiceUnavailable, ())),
            },
            None => Outcome::Error((Status::InternalServerError, ())),
        }
    }
}

// Embed the migrations from the migrations folder into the application
// This way, the program automatically migrates the database to the latest version
// https://docs.rs/diesel_migrations/*/diesel_migrations/macro.embed_migrations.html
#[cfg(sqlite)]
mod sqlite_migrations {
    use diesel::{Connection, RunQueryDsl};
    use diesel_migrations::{EmbeddedMigrations, MigrationHarness};
    pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("migrations/sqlite");

    pub fn run_migrations(db_url: &str) -> Result<(), super::Error> {
        // Establish a connection to the sqlite database (this will create a new one, if it does
        // not exist, and exit if there is an error).
        let mut connection = diesel::sqlite::SqliteConnection::establish(db_url)?;

        // Run the migrations after successfully establishing a connection
        // Disable Foreign Key Checks during migration
        // Scoped to a connection.
        diesel::sql_query("PRAGMA foreign_keys = OFF")
            .execute(&mut connection)
            .expect("Failed to disable Foreign Key Checks during migrations");

        // Turn on WAL in SQLite
        if crate::CONFIG.enable_db_wal() {
            diesel::sql_query("PRAGMA journal_mode=wal").execute(&mut connection).expect("Failed to turn on WAL");
        }

        connection.run_pending_migrations(MIGRATIONS).expect("Error running migrations");
        Ok(())
    }
}

#[cfg(mysql)]
mod mysql_migrations {
    use diesel::{Connection, RunQueryDsl};
    use diesel_migrations::{EmbeddedMigrations, MigrationHarness};
    pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("migrations/mysql");

    pub fn run_migrations(db_url: &str) -> Result<(), super::Error> {
        // Make sure the database is up to date (create if it doesn't exist, or run the migrations)
        let mut connection = diesel::mysql::MysqlConnection::establish(db_url)?;

        // Disable Foreign Key Checks during migration
        // Scoped to a connection/session.
        diesel::sql_query("SET FOREIGN_KEY_CHECKS = 0")
            .execute(&mut connection)
            .expect("Failed to disable Foreign Key Checks during migrations");

        connection.run_pending_migrations(MIGRATIONS).expect("Error running migrations");
        Ok(())
    }
}

#[cfg(postgresql)]
mod postgresql_migrations {
    use diesel::Connection;
    use diesel_migrations::{EmbeddedMigrations, MigrationHarness};
    pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("migrations/postgresql");

    pub fn run_migrations(db_url: &str) -> Result<(), super::Error> {
        // Make sure the database is up to date (create if it doesn't exist, or run the migrations)
        let mut connection = diesel::pg::PgConnection::establish(db_url)?;

        connection.run_pending_migrations(MIGRATIONS).expect("Error running migrations");
        Ok(())
    }
}

/// A throwaway SQLite database with the real schema, for tests that need to run queries.
#[cfg(all(test, sqlite))]
pub mod test_db {
    use super::{DbConn, DbConnManager, sqlite_migrations::MIGRATIONS};
    use diesel::{Connection, connection::SimpleConnection, r2d2::Pool, sqlite::SqliteConnection};
    use diesel_migrations::MigrationHarness;
    use std::sync::{Arc, atomic::AtomicBool};
    use tokio::sync::Mutex;

    pub struct TestDb {
        pool: Pool<DbConnManager>,
        // Keep the named in-memory database alive across pooled connections.
        _keeper: SqliteConnection,
    }

    impl TestDb {
        fn new() -> Self {
            let url = format!("sqlite://{}?mode=memory&cache=shared", crate::util::get_uuid());
            let mut keeper = SqliteConnection::establish(&url).unwrap();
            keeper.run_pending_migrations(MIGRATIONS).unwrap();
            let pool = Pool::builder().max_size(2).build(DbConnManager::new(&url)).unwrap();
            Self {
                pool,
                _keeper: keeper,
            }
        }

        pub fn conn(&self) -> DbConn {
            DbConn {
                conn: Arc::new(Mutex::new(Some(self.pool.get().unwrap()))),
                permit: None,
                transaction_open: AtomicBool::new(false),
            }
        }
    }

    // CONFIG must initialize in a child process: the DNS bootstrap test needs it uninitialized.
    pub fn run(test: impl AsyncFnOnce(&TestDb)) {
        if std::env::var_os("VAULTWARDEN_ISOLATED_CONFIG_TEST").is_none() {
            let output = std::process::Command::new(std::env::current_exe().unwrap())
                .args(["--exact", std::thread::current().name().unwrap(), "--nocapture"])
                .env("VAULTWARDEN_ISOLATED_CONFIG_TEST", "1")
                .env("DATABASE_URL", "sqlite://:memory:")
                .output()
                .unwrap();
            assert!(
                output.status.success() && String::from_utf8_lossy(&output.stdout).contains("1 passed; 0 failed"),
                "{}\n{}",
                String::from_utf8_lossy(&output.stdout),
                String::from_utf8_lossy(&output.stderr)
            );
            return;
        }
        // DbConn::run uses block_in_place, requiring a multi-threaded runtime.
        tokio::runtime::Builder::new_multi_thread().enable_all().build().unwrap().block_on(async {
            super::models::test_flag::CUSTOM_ROLES_ENABLED.set(Some(true));
            test(&TestDb::new()).await;
        });
    }

    #[test]
    fn custom_role_migration_preserves_memberships_and_guards_downgrade() {
        use super::schema::users_organizations::dsl::*;
        use diesel::{migration::MigrationSource, prelude::*};
        let mut conn = SqliteConnection::establish(":memory:").unwrap();
        for migration in conn.pending_migrations(MIGRATIONS).unwrap() {
            if migration.name().version().to_string() == "20261010120000" {
                break;
            }
            conn.run_migration(&migration).unwrap();
        }
        conn.batch_execute(
            "PRAGMA foreign_keys=OFF;
            INSERT INTO organizations (uuid,name,billing_email) VALUES ('org','Org','a@b');
            INSERT INTO users_organizations (uuid,user_uuid,org_uuid,access_all,akey,status,atype) VALUES
            ('0','0','org',TRUE,'key',2,0), ('1','1','org',TRUE,'key',2,1),
            ('2','2','org',FALSE,'key',1,2), ('3','3','org',FALSE,'key',2,3), ('4','4','org',TRUE,'key',2,3);",
        )
        .unwrap();
        let rows = |conn: &mut SqliteConnection| {
            users_organizations
                .select((uuid, user_uuid, org_uuid, access_all, akey, status, atype))
                .order(uuid)
                .load::<(String, String, String, bool, String, i32, i32)>(conn)
                .unwrap()
        };
        let before = rows(&mut conn);
        conn.run_pending_migrations(MIGRATIONS).unwrap();
        assert_eq!(rows(&mut conn), before);
        let grants = |conn: &mut SqliteConnection| {
            diesel::select(diesel::dsl::sql::<diesel::sql_types::BigInt>(
                "(SELECT COUNT(*) FROM users_organizations WHERE manage_users OR manage_groups OR manage_policies
                OR create_new_collections OR edit_any_collection OR delete_any_collection
                OR access_event_logs OR access_import_export OR access_reports)",
            ))
            .get_result::<i64>(conn)
            .unwrap()
        };
        assert_eq!(grants(&mut conn), 0);
        let migration = MigrationSource::<diesel::sqlite::Sqlite>::migrations(&MIGRATIONS)
            .unwrap()
            .into_iter()
            .find(|m| m.name().version().to_string() == "20261010120000")
            .unwrap();
        for member_status in [1, 2] {
            conn.batch_execute(&format!(
                "UPDATE users_organizations SET atype=4,manage_users=TRUE,status={member_status} WHERE uuid='2'"
            ))
            .unwrap();
            let custom = rows(&mut conn);
            assert!(conn.revert_migration(&*migration).is_err());
            assert_eq!(rows(&mut conn), custom);
            assert!(users_organizations.filter(uuid.eq("2")).select(manage_users).first::<bool>(&mut conn).unwrap());
        }
        conn.batch_execute("UPDATE users_organizations SET atype=2,manage_users=FALSE,status=1 WHERE uuid='2'")
            .unwrap();
        conn.revert_migration(&*migration).unwrap();
        assert!(diesel::sql_query("SELECT manage_users FROM users_organizations").execute(&mut conn).is_err());
        conn.run_pending_migrations(MIGRATIONS).unwrap();
        assert_eq!(grants(&mut conn), 0);
        assert_eq!(rows(&mut conn), before);
    }

    use super::models::{Collection, CustomPermission, Membership, MembershipType, Organization, OrganizationId, User};

    pub async fn organization(conn: &DbConn) -> Organization {
        let org = Organization::new("Org".into(), "billing@example.com", None, None);
        org.save(conn).await.unwrap();
        org
    }

    pub async fn member(
        conn: &DbConn,
        org: &OrganizationId,
        role: MembershipType,
        permissions: &[CustomPermission],
    ) -> (User, Membership) {
        let mut user = User::new(&format!("{}@example.com", crate::util::get_uuid()), None);
        user.password_hash = b"existing-account".to_vec();
        user.save(conn).await.unwrap();
        let mut member = Membership::new(user.uuid.clone(), org.clone(), None);
        member.set_type(role);
        member.status = 2;
        member.set_custom_permissions(&permissions.iter().copied().collect());
        member.save(conn).await.unwrap();
        (user, member)
    }

    pub async fn custom_member(conn: &DbConn, permission: CustomPermission) -> (Organization, User, Membership) {
        let org = organization(conn).await;
        let (user, member) = member(conn, &org.uuid, MembershipType::Custom, &[permission]).await;
        (org, user, member)
    }

    pub async fn collection(conn: &DbConn, org: &OrganizationId) -> Collection {
        let collection = Collection::new(org.clone(), "Collection".into(), None);
        collection.save(conn).await.unwrap();
        collection
    }

    #[test]
    fn transactions_rollback_before_connection_reuse_and_commit_success() {
        use crate::error::MapResult;
        use diesel::RunQueryDsl;
        run(async |db| {
            let conn = db.conn();
            conn.run(|c| c.batch_execute("CREATE TABLE probe (id INTEGER PRIMARY KEY)")).await.unwrap();
            for (sql, expected) in [
                ("INSERT INTO probe VALUES (1); INSERT INTO probe VALUES (1)", 0_i64),
                ("INSERT INTO probe VALUES (1)", 1),
            ] {
                let result: Result<(), super::Error> = conn
                    .transaction(async || conn.run(move |c| c.batch_execute(sql)).await.map_res("Insert failed"))
                    .await;
                assert_eq!(result.is_ok(), expected == 1);
                let count: i64 = conn
                    .run(|c| {
                        diesel::select(diesel::dsl::sql::<diesel::sql_types::BigInt>("(SELECT COUNT(*) FROM probe)"))
                            .get_result(c)
                    })
                    .await
                    .unwrap();
                assert_eq!(count, expected);
                // The same live connection must accept another transaction, before DbConn::drop can help.
                conn.run(|c| c.batch_execute("BEGIN; ROLLBACK")).await.unwrap();
            }
        });
    }

    #[test]
    fn collection_and_cipher_access_requires_active_membership_and_enabled_groups() {
        use super::models::{Cipher, CipherAccessScope, CollectionCipher, CollectionUser, test_flag};
        use crate::{
            api::core::{CipherSyncData, CipherSyncType},
            auth::{CollectionAction, custom_member_may},
        };
        run(async |db| {
            let conn = db.conn();
            let (org, user, mut member) = custom_member(&conn, CustomPermission::EditAnyCollection).await;
            let collection = collection(&conn, &org.uuid).await;
            let mut cipher = Cipher::new(2);
            cipher.name = "Cipher".into();
            cipher.organization_uuid = Some(org.uuid.clone());
            cipher.save(&conn).await.unwrap();
            CollectionCipher::save(&cipher.uuid, &collection.uuid, &conn).await.unwrap();
            for (enabled, status, allowed) in [(true, 2, true), (false, 2, false), (true, 1, false)] {
                test_flag::CUSTOM_ROLES_ENABLED.set(Some(enabled));
                member.status = status;
                member.save(&conn).await.unwrap();
                assert!(!cipher.is_accessible_to_user(&user.uuid, CipherAccessScope::User, &conn).await);
                assert_eq!(
                    cipher.is_write_accessible_to_user(&user.uuid, CipherAccessScope::OrganizationAdmin, &conn).await,
                    allowed
                );
            }
            test_flag::CUSTOM_ROLES_ENABLED.set(Some(true));
            member.status = 2;
            member.save(&conn).await.unwrap();
            let foreign = organization(&conn).await;
            cipher.organization_uuid = Some(foreign.uuid);
            assert!(!cipher.is_accessible_to_user(&user.uuid, CipherAccessScope::OrganizationAdmin, &conn).await);
            cipher.organization_uuid = Some(org.uuid.clone());
            member.set_custom_permissions(&std::collections::HashSet::new());
            member.save(&conn).await.unwrap();
            let sql=format!("INSERT INTO groups (uuid,organizations_uuid,name,access_all,creation_date,revision_date)
                VALUES ('g','{}','Group',FALSE,CURRENT_TIMESTAMP,CURRENT_TIMESTAMP);
                INSERT INTO groups_users (groups_uuid,users_organizations_uuid) VALUES ('g','{}');
                INSERT INTO collections_groups (collections_uuid,groups_uuid,read_only,hide_passwords,manage) VALUES ('{}','g',FALSE,FALSE,TRUE);",
                org.uuid,member.uuid,collection.uuid);
            conn.run(move |c| c.batch_execute(&sql)).await.unwrap();
            for status in [2, 1, -3, -126] {
                member.status = status;
                member.save(&conn).await.unwrap();
                let allowed = status == 2 && crate::CONFIG.org_groups_enabled();
                assert_eq!(custom_member_may(CollectionAction::Edit, &member, &collection.uuid, &conn).await, allowed);
                let sync = CipherSyncData::new(&user.uuid, CipherSyncType::User, &conn).await;
                for cached in [None, Some(&sync)] {
                    assert_eq!(
                        cipher
                            .get_access_restrictions(&user.uuid, CipherAccessScope::User, cached, &conn)
                            .await
                            .is_some(),
                        allowed
                    );
                }
            }
            member.status = 2;
            member.save(&conn).await.unwrap();
            CollectionUser::save(&user.uuid, &collection.uuid, false, false, true, &conn).await.unwrap();
            assert!(custom_member_may(CollectionAction::Edit, &member, &collection.uuid, &conn).await);
        });
    }
}