//! Per-test database isolation using PostgreSQL template databases.
//!
//! A shared template database is created once (with all migrations) and each
//! test gets a cheap `CREATE DATABASE ... TEMPLATE` clone. Dropped
//! automatically when `TestDb` goes out of scope.

use sqlx::postgres::PgPoolOptions;
use sqlx::{Connection, Executor, PgConnection, PgPool};
use std::sync::{Once, OnceLock};
use std::time::Duration;
use uuid::Uuid;

/// Base name of the shared template database. The live name is suffixed with
/// the connecting role (see `template_name`).
const TEMPLATE_DB_BASE: &str = "mnw_test_template";

/// The live template name for this run, namespaced by the connecting postgres
/// role. Two suites run by different roles on a shared cluster (e.g. a
/// developer's `max` and the Sando gate's `sando`) get distinct templates, so
/// they never collide AND each role owns, and can therefore drop, its own.
/// This is what prevents a stale, foreign-owned `mnw_test_template` from
/// wedging a deploy's `cargo_test` gate. Set once by `ensure_template`.
static TEMPLATE_NAME: OnceLock<String> = OnceLock::new();

fn template_name() -> &'static str {
    TEMPLATE_NAME.get().map_or(TEMPLATE_DB_BASE, String::as_str)
}

/// Reduce a postgres role to a safe identifier suffix (`[a-z0-9_]`).
fn sanitize_role(role: &str) -> String {
    role.chars()
        .map(|c| {
            if c.is_ascii_alphanumeric() || c == '_' {
                c.to_ascii_lowercase()
            } else {
                '_'
            }
        })
        .collect()
}

/// Shared, deploy-stable role that owns all test databases when present. Both
/// the interactive (`max`) and Sando-gate (`sando`) logins are members, so
/// whoever runs the suite, every `mnw_test_*` DB ends up owned by this role and
/// any member can drop it, no superuser, no cross-user ownership clashes.
const SHARED_ROLE: &str = "mnw_test";

/// Best-effort `SET ROLE mnw_test` so subsequent `CREATE DATABASE`s are owned
/// by the shared role. Silently no-ops when the role is absent or the login
/// isn't a member (e.g. a fresh dev box without bootstrap), the per-role
/// template namespacing is the fallback in that case.
async fn assume_shared_role(conn: &mut PgConnection) {
    let _ = conn
        .execute(format!("SET ROLE \"{SHARED_ROLE}\"").as_str())
        .await;
}

/// Ensures template creation runs exactly once, across all threads and runtimes.
static TEMPLATE_INIT: Once = Once::new();

fn admin_url() -> String {
    std::env::var("TEST_DATABASE_URL")
        .unwrap_or_else(|_| "postgres://localhost/postgres".to_string())
}

/// Advisory-lock key serializing template setup across every process/connection
/// on the cluster (`std::sync::Once` only covers one process). Concurrent test
/// binaries must not both drop+recreate the shared template. Arbitrary, stable.
const TEMPLATE_LOCK_KEY: i64 = 0x6D6E_775F_7470_6C00; // "mnw_tpl\0"

/// Latest migration version embedded in this binary (head of `migrations/`).
///
/// `sqlx::migrate!` embeds every migration as one array literal, so
/// `large_stack_arrays` measures the whole `migrations/` directory. It crossed
/// the 16KB threshold when the mailing-list tables landed. Allowed rather than
/// worked around: the array is the point of the macro.
#[allow(
    clippy::large_stack_arrays,
    reason = "sqlx::migrate! embeds the whole directory"
)]
fn latest_migration_version() -> i64 {
    sqlx::migrate!("./migrations")
        .iter()
        .map(|m| m.version)
        .max()
        .unwrap_or(0)
}

/// True if `template` already exists and its applied-migration head matches the
/// code's latest migration, a clone of it is up to date, so we reuse it instead
/// of dropping + rebuilding. Any error / missing table => "not current" =>
/// caller rebuilds.
async fn template_is_current(admin_url: &str, template: &str) -> bool {
    let Ok(mut conn) = PgConnection::connect(admin_url).await else {
        return false;
    };
    let exists = sqlx::query_as::<_, (i32,)>("SELECT 1 FROM pg_database WHERE datname = $1")
        .bind(template)
        .fetch_optional(&mut conn)
        .await
        .ok()
        .flatten()
        .is_some();
    if !exists {
        return false;
    }
    let tpl_url = replace_db_name(admin_url, template);
    let Ok(mut tconn) = PgConnection::connect(&tpl_url).await else {
        return false;
    };
    matches!(
        sqlx::query_as::<_, (Option<i64>,)>("SELECT MAX(version) FROM _sqlx_migrations")
            .fetch_one(&mut tconn)
            .await,
        Ok((Some(v),)) if v == latest_migration_version()
    )
}

/// Create the template database with all migrations. Runs in a dedicated
/// single-threaded tokio runtime so it works from any context (including
/// inside `#[tokio::test]` and plain `#[test]`).
fn ensure_template() {
    TEMPLATE_INIT.call_once(|| {
        let rt = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .expect("build template setup runtime");

        rt.block_on(async {
            let t0 = std::time::Instant::now();
            let admin = admin_url();
            let mut conn = PgConnection::connect(&admin)
                .await
                .expect("connect to admin DB for template setup");

            // Adopt the shared role so the template (and every clone) is owned
            // by `mnw_test`, droppable by any member. No-ops on hosts without
            // the role, falling back to per-role namespacing below.
            assume_shared_role(&mut conn).await;

            // Namespace the template by the *effective* role: with the shared
            // role assumed this collapses to one template for everyone; without
            // it, each login role gets its own so a developer's `max` run and
            // the Sando gate's `sando` run never collide on a shared cluster.
            let (role,): (String,) = sqlx::query_as("SELECT current_user")
                .fetch_one(&mut conn)
                .await
                .expect("query current_user for template namespacing");
            let template = format!("{TEMPLATE_DB_BASE}_{}", sanitize_role(&role));
            let _ = TEMPLATE_NAME.set(template.clone());

            // Serialize template setup across ALL processes on this cluster, not
            // just this process's `Once`. Without it, concurrent test binaries
            // each unconditionally drop+recreate the SAME shared-role template,
            // so one drops the template another is mid-clone from, the mass
            // "template database ... does not exist" flake on the deploy gate.
            // Held until `conn` drops at the end of this block.
            sqlx::query("SELECT pg_advisory_lock($1)")
                .bind(TEMPLATE_LOCK_KEY)
                .execute(&mut conn)
                .await
                .expect("acquire template advisory lock");

            // Reuse the template when it's already migration-current. Dropping a
            // live template is what races concurrent clones, and a FORCE drop
            // also fails when another role holds connections we can't terminate
            // ("permission denied to terminate process"). Only rebuild when the
            // template is missing or its migration head is stale.
            if template_is_current(&admin, &template).await {
                eprintln!("[test-harness] Reusing current template DB {template}");
            } else {
                // We hold the cross-process lock, so no clone can be reading it.
                conn.execute(
                    format!("DROP DATABASE IF EXISTS \"{template}\" WITH (FORCE)").as_str(),
                )
                .await
                .unwrap_or_else(|e| {
                    panic!(
                        "drop stale template {template}: {e} \
                     (if owned by a different role, drop it as the postgres superuser)"
                    )
                });

                conn.execute(format!("CREATE DATABASE \"{template}\"").as_str())
                    .await
                    .expect("create template database");

                // Connect to the template and run all migrations
                let tpl_url = replace_db_name(&admin, &template);
                let tpl_pool = PgPoolOptions::new()
                    .max_connections(2)
                    .acquire_timeout(Duration::from_secs(10))
                    .connect(&tpl_url)
                    .await
                    .expect("connect to template database");

                let t_migrate = std::time::Instant::now();
                #[allow(clippy::large_stack_arrays, reason = "sqlx::migrate! embeds the whole directory")]
                sqlx::migrate!("./migrations")
                    .run(&tpl_pool)
                    .await
                    .expect("run migrations on template");
                let migrate_ms = t_migrate.elapsed().as_millis();

                // Also create the session store table
                let session_store = tower_sessions_sqlx_store::PostgresStore::new(tpl_pool.clone());
                session_store
                    .migrate()
                    .await
                    .expect("session store migration on template");

                tpl_pool.close().await;

                let total_ms = t0.elapsed().as_millis();
                eprintln!(
                    "[test-harness] Template DB created in {total_ms}ms (migrations: {migrate_ms}ms)"
                );
            }

            // Release explicitly (also released when `conn` drops); the lock was
            // held across the whole reuse-or-rebuild window above.
            let _ = sqlx::query("SELECT pg_advisory_unlock($1)")
                .bind(TEMPLATE_LOCK_KEY)
                .execute(&mut conn)
                .await;
        });
    });
}

/// An isolated test database that cleans up after itself.
pub(crate) struct TestDb {
    pub pool: PgPool,
    db_name: String,
    admin_url: String,
    #[allow(dead_code)]
    test_url: String,
    /// Whether the session store table already exists (from template).
    pub session_migrated: bool,
}

impl TestDb {
    /// Create a fresh database cloned from the shared template.
    pub(crate) async fn new() -> Self {
        // ensure_template uses std::sync::Once + its own runtime, safe from any context.
        // When called from an async context, we run it on a blocking thread to avoid
        // nesting runtimes.
        tokio::task::spawn_blocking(ensure_template)
            .await
            .expect("template setup panicked");

        let t0 = std::time::Instant::now();
        let admin = admin_url();
        let db_name = format!("mnw_test_{}", Uuid::new_v4().simple());

        let mut admin_conn = PgConnection::connect(&admin)
            .await
            .expect("Failed to connect to admin database");

        // Own the clone via the shared role too, so cleanup (and any member)
        // can drop it. Matches the template's ownership.
        assume_shared_role(&mut admin_conn).await;

        // Clone from the template, retrying on transient contention. Postgres
        // serializes `CREATE DATABASE ... TEMPLATE t`: a concurrent clone of the
        // same template (under the full-suite parallel stampede) fails with
        // "source database ... is being accessed by other users". Retry rather
        // than fail the test on that transient.
        let create_sql = format!(
            "CREATE DATABASE \"{db_name}\" TEMPLATE \"{}\"",
            template_name()
        );
        let mut attempt = 0u32;
        loop {
            match admin_conn.execute(create_sql.as_str()).await {
                Ok(_) => break,
                Err(_) if attempt < 8 => {
                    attempt += 1;
                    tokio::time::sleep(Duration::from_millis(100 * attempt as u64)).await;
                }
                Err(e) => panic!(
                    "Failed to create test database from template after {attempt} retries: {e}"
                ),
            }
        }

        let test_url = replace_db_name(&admin, &db_name);

        let pool = PgPoolOptions::new()
            .max_connections(5)
            .acquire_timeout(Duration::from_secs(5))
            .connect(&test_url)
            .await
            .expect("Failed to connect to test database");

        let clone_ms = t0.elapsed().as_millis();
        if clone_ms > 500 {
            eprintln!("[test-harness] SLOW DB clone: {clone_ms}ms for {db_name}");
        }

        TestDb {
            pool,
            db_name,
            admin_url: admin,
            test_url,
            session_migrated: true,
        }
    }

    /// The connection URL for this test database.
    #[allow(dead_code)]
    pub(crate) fn url(&self) -> &str {
        &self.test_url
    }
}

impl Drop for TestDb {
    fn drop(&mut self) {
        let admin_url = self.admin_url.clone();
        let db_name = self.db_name.clone();
        let pool = self.pool.clone();

        std::thread::spawn(move || {
            let rt = tokio::runtime::Builder::new_current_thread()
                .enable_all()
                .build()
                .expect("Failed to build cleanup runtime");

            rt.block_on(async {
                // Actually close the pool, and wait for it. Every one of its
                // live connections is a client of the database we're about to
                // drop, and `DROP DATABASE` refuses while any exist. Bounded:
                // `close()` waits for checked-out connections to come back, and
                // a test that leaked a guard would otherwise hang the suite
                // here forever. On timeout we fall through, the FORCE drop
                // below evicts whatever is left.
                if tokio::time::timeout(Duration::from_secs(5), pool.close())
                    .await
                    .is_err()
                {
                    eprintln!(
                        "[test-harness] pool for {db_name} did not close in 5s; forcing"
                    );
                }

                let Ok(mut conn) = PgConnection::connect(&admin_url).await else {
                    eprintln!("[test-harness] LEAKED {db_name}: admin connect failed");
                    return;
                };

                // Deliberately NO `SET ROLE mnw_test` here, unlike the create
                // path. Terminating a backend requires membership in the role
                // that *opened* it, and these were opened by our login role
                // (`max`, `sando`), not by `mnw_test`, which is a member of
                // neither. Assuming the shared role therefore turns both the
                // sweep below and FORCE's own eviction into "permission denied
                // to terminate process", leaking the clone. Dropping needs only
                // membership in the owning role (pg_has_role USAGE), which we
                // already have as the login role, so staying put satisfies both.
                let _ = conn
                    .execute(
                        format!(
                            "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = '{db_name}'"
                        )
                        .as_str(),
                    )
                    .await;

                // FORCE, and retry: a connection can still be mid-handshake
                // when we terminate, so it survives the sweep above and lands
                // in the split second before the drop. Bare `DROP DATABASE`
                // fails outright on that race; FORCE evicts it, and the retries
                // cover the case where another arrives.
                let drop_sql = format!("DROP DATABASE IF EXISTS \"{db_name}\" WITH (FORCE)");
                for attempt in 0..4u32 {
                    match conn.execute(drop_sql.as_str()).await {
                        Ok(_) => return,
                        Err(e) if attempt == 3 => {
                            // Never silent: a swallowed failure here is exactly
                            // how the cluster accumulated ~1000 stale clones.
                            eprintln!("[test-harness] LEAKED {db_name}: drop failed: {e}");
                        }
                        Err(_) => {
                            tokio::time::sleep(Duration::from_millis(50 * (attempt + 1) as u64))
                                .await;
                        }
                    }
                }
            });
        })
        .join()
        .ok();
    }
}

/// Replace the database name in a PostgreSQL connection URL.
fn replace_db_name(url: &str, new_db: &str) -> String {
    if let Some(pos) = url.rfind('/') {
        let base = &url[..pos];
        let query = url[pos + 1..].find('?').map_or("", |q| &url[pos + 1 + q..]);
        if query.is_empty() {
            format!("{base}/{new_db}")
        } else {
            format!("{base}/{new_db}{query}")
        }
    } else {
        panic!("Invalid database URL: no '/' found");
    }
}

#[cfg(test)]
pub(crate) mod tests {
    use super::*;

    #[test]
    fn replace_db_name_simple() {
        let result = replace_db_name("postgres://localhost/postgres", "test_db");
        assert_eq!(result, "postgres://localhost/test_db");
    }

    #[test]
    fn replace_db_name_with_auth() {
        let result = replace_db_name("postgres://user:pass@localhost:5432/mydb", "test_db");
        assert_eq!(result, "postgres://user:pass@localhost:5432/test_db");
    }

    #[test]
    fn replace_db_name_with_query() {
        let result = replace_db_name("postgres://localhost/postgres?sslmode=disable", "test_db");
        assert_eq!(result, "postgres://localhost/test_db?sslmode=disable");
    }
}