/*
 * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>
 *
 * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
 *
 * Modified by Coffey Labs in 2026 for INBUXA.
 */

use super::{PostgresStore, bounded, into_error};
use crate::{
    backend::postgres::{
        PsqlSearchField, into_pool_error,
        search::{PG_FALLBACK_LANG, PG_LANGS, PG_UNSTEMMED_LANG},
        tls::MakeRustlsConnect,
    },
    search::{
        CalendarSearchField, ContactSearchField, EmailSearchField, SearchableField,
        TracingSearchField,
    },
    *,
};
use ::registry::schema::{enums::PostgreSqlRecyclingMethod, structs};
use ahash::AHashSet;
use deadpool_postgres::{
    Config, ManagerConfig, Object, Pool, PoolConfig, RecyclingMethod, Runtime, Timeouts,
};
use std::time::Duration;
use tokio_postgres::NoTls;
use utils::tls::rustls_client_config;

/// inbuxa: how long a request waits for a pooled connection.
pub(crate) const POOL_WAIT_TIMEOUT: Duration = Duration::from_secs(30);
/// inbuxa: how long opening a connection may take when the store sets no
/// timeout of its own.
pub(crate) const POOL_CREATE_TIMEOUT: Duration = Duration::from_secs(15);
/// inbuxa: how long checking a pooled connection before reuse may take.
pub(crate) const POOL_RECYCLE_TIMEOUT: Duration = Duration::from_secs(10);
/// inbuxa: idle time before TCP keepalive probes start.
pub(crate) const POOL_KEEPALIVE_IDLE: Duration = Duration::from_secs(60);

/// inbuxa: the pool's timeouts. Opening a connection is bounded by the
/// store's own timeout when it has one; waiting for one covers at least that
/// long, so a slow connect isn't cut short by the wait.
pub(crate) fn pool_timeouts(connect_timeout: Option<Duration>) -> Timeouts {
    let create = connect_timeout.unwrap_or(POOL_CREATE_TIMEOUT);
    Timeouts {
        wait: POOL_WAIT_TIMEOUT.max(create).into(),
        create: create.into(),
        recycle: POOL_RECYCLE_TIMEOUT.into(),
    }
}

impl PostgresStore {
    pub async fn open(config: structs::PostgreSqlStore) -> Result<Store, String> {
        // inbuxa: ST-15: where the primary is, to tell a replica from it
        let primary_location = (config.host.clone(), config.port as u16, config.database.clone());
        let mut cfg = Config::new();
        cfg.dbname = config.database.into();
        cfg.host = config.host.into();
        cfg.user = config.auth_username;
        cfg.password = config.auth_secret.secret().await?.map(|v| v.into_owned());
        cfg.port = (config.port as u16).into();
        cfg.connect_timeout = config.timeout.map(|t| t.into_inner());
        cfg.options = config.options;
        cfg.manager = Some(ManagerConfig {
            recycling_method: match config.pool_recycling_method {
                PostgreSqlRecyclingMethod::Fast => RecyclingMethod::Fast,
                PostgreSqlRecyclingMethod::Verified => RecyclingMethod::Verified,
                PostgreSqlRecyclingMethod::Clean => RecyclingMethod::Clean,
            },
        });
        // inbuxa: upstream set no pool timeouts, so a request waited for a
        // free connection, or for one to be made or recycled, for as long as
        // it took: forever when the server stopped answering. A worker now
        // gets an error instead and the task or request is retried.
        let mut pool = config
            .pool_max_connections
            .map(|max_conn| PoolConfig::new(max_conn as usize))
            .unwrap_or_default();
        pool.timeouts = pool_timeouts(cfg.connect_timeout);
        cfg.pool = pool.into();
        // Notice a server that went away without closing the connection in
        // minutes rather than the system default of two hours
        cfg.keepalives = true.into();
        cfg.keepalives_idle = POOL_KEEPALIVE_IDLE.into();

        let primary_pool = if config.use_tls {
            cfg.create_pool(
                Some(Runtime::Tokio1),
                MakeRustlsConnect::new(rustls_client_config(config.allow_invalid_certs)?),
            )
        } else {
            cfg.create_pool(Some(Runtime::Tokio1), NoTls)
        }
        .map_err(|e| format!("Failed to create connection pool: {e}"))?;
        let ts_configs = discover_ts_configs(&primary_pool).await;

        // inbuxa: ST-5 to ST-15: each replica inherits the primary's settings
        // except where it is and how to sign in
        let mut replicas = vec![];
        for replica in config.read_replicas {
            let mut cfg = cfg.clone();
            cfg.dbname = replica.database.clone().into();
            cfg.host = replica.host.clone().into();
            cfg.user = replica.auth_username;
            cfg.password = replica.auth_secret.secret().await?.map(|v| v.into_owned());
            cfg.port = (replica.port as u16).into();
            cfg.options = replica.options;
            let pool = if config.use_tls {
                cfg.create_pool(
                    Some(Runtime::Tokio1),
                    MakeRustlsConnect::new(rustls_client_config(config.allow_invalid_certs)?),
                )
            } else {
                cfg.create_pool(Some(Runtime::Tokio1), NoTls)
            }
            .map_err(|e| format!("Failed to create connection pool: {e}"))?;
            replicas.push(crate::backend::scaleout::replica::Replica::new(
                Store::PostgreSQL(Arc::new(PostgresStore {
                    conn_pool: pool,
                    ts_configs: ts_configs.clone(),
                    timeouts: Default::default(),
                })),
                replica.host,
                replica.port as u16,
                replica.database,
            ));
        }

        let primary = Store::PostgreSQL(Arc::new(PostgresStore {
            conn_pool: primary_pool,
            ts_configs,
            timeouts: Default::default(),
        }));

        // ST-1: no replicas, no change
        if replicas.is_empty() {
            return Ok(primary);
        }
        Ok(Store::Replicated(
            crate::backend::scaleout::replica::ReplicatedStore::new(
                primary,
                primary_location,
                replicas,
                crate::backend::scaleout::replica::ReplicaKind::PostgreSql,
            ),
        ))
    }

    pub(crate) async fn create_storage_tables(&self) -> trc::Result<()> {
        let conn = self.conn_pool.get().await.map_err(into_pool_error)?;
        let limit = self.timeouts.maintenance;
        let result = tokio::time::timeout(limit, async {
            for table in [
                SUBSPACE_ACL,
                SUBSPACE_TASK_QUEUE,
                SUBSPACE_DELETED_ITEMS,
                SUBSPACE_SPAM_SAMPLES,
                crate::SUBSPACE_INBUXA, // inbuxa: masked email
                SUBSPACE_BLOB_LINK,
                SUBSPACE_IN_MEMORY_VALUE,
                SUBSPACE_PROPERTY,
                SUBSPACE_REGISTRY,
                SUBSPACE_REGISTRY_PK,
                SUBSPACE_QUEUE_MESSAGE,
                SUBSPACE_QUEUE_EVENT,
                SUBSPACE_REPORT_OUT,
                SUBSPACE_REPORT_IN,
                SUBSPACE_LOGS,
                SUBSPACE_BLOBS,
                SUBSPACE_DIRECTORY,
                SUBSPACE_TELEMETRY_SPAN,
                SUBSPACE_TELEMETRY_METRIC,
            ] {
                let table = char::from(table);
                conn.execute(
                    &format!(
                        "CREATE TABLE IF NOT EXISTS {table} (
                            k BYTEA PRIMARY KEY,
                            v BYTEA NOT NULL
                        )"
                    ),
                    &[],
                )
                .await
                .map_err(into_error)?;
            }

            for table in [SUBSPACE_INDEXES, SUBSPACE_REGISTRY_IDX] {
                let table = char::from(table);
                conn.execute(
                    &format!(
                        "CREATE TABLE IF NOT EXISTS {table} (
                            k BYTEA PRIMARY KEY
                        )"
                    ),
                    &[],
                )
                .await
                .map_err(into_error)?;
            }

            for table in [SUBSPACE_COUNTER, SUBSPACE_QUOTA, SUBSPACE_IN_MEMORY_COUNTER] {
                conn.execute(
                    &format!(
                        "CREATE TABLE IF NOT EXISTS {} (
                        k BYTEA PRIMARY KEY,
                        v BIGINT NOT NULL DEFAULT 0
                    )",
                        char::from(table)
                    ),
                    &[],
                )
                .await
                .map_err(into_error)?;
            }

            Ok(())
        })
        .await;
        bounded(conn, result, limit)
    }

    pub(crate) async fn create_search_tables(&self) -> trc::Result<()> {
        let conn = self.conn_pool.get().await.map_err(into_pool_error)?;
        let limit = self.timeouts.maintenance;
        let result = tokio::time::timeout(limit, async {
            create_search_tables::<EmailSearchField>(&conn).await?;
            create_search_tables::<CalendarSearchField>(&conn).await?;
            create_search_tables::<ContactSearchField>(&conn).await?;
            //create_search_tables::<FileSearchField>(&conn).await?;
            create_search_tables::<TracingSearchField>(&conn).await?;

            Ok(())
        })
        .await;
        bounded(conn, result, limit)
    }
}

async fn create_search_tables<T: SearchableField + PsqlSearchField + 'static>(
    conn: &Object,
) -> trc::Result<()> {
    let table_name = T::index().psql_table();
    let mut query = format!("CREATE TABLE IF NOT EXISTS {} (", table_name);

    // Add primary key columns
    let pkeys = T::primary_keys();
    for pkey in pkeys {
        query.push_str(&format!("{} {}, ", pkey.column(), pkey.column_type()));
    }

    // Add other columns
    for field in T::all_fields() {
        query.push_str(&format!("{} {}", field.column(), field.column_type()));
        if let Some(sort_type) = field.sort_column_type() {
            query.push_str(&format!(", {} {}", field.sort_column().unwrap(), sort_type));
        }
        query.push_str(", ");
    }

    // Add primary key constraint
    query.push_str("PRIMARY KEY (");
    for (i, pkey) in pkeys.iter().enumerate() {
        if i > 0 {
            query.push_str(", ");
        }
        query.push_str(pkey.column());
    }
    query.push_str("))");

    conn.execute(&query, &[]).await.map_err(into_error)?;

    // Create indexes
    for field in T::all_fields() {
        if field.is_text() || field.is_json() {
            let column_name = field.column();
            // inbuxa: with GIN's default fastupdate=on, new entries wait in
            // an unindexed pending list that every search scans in full
            // until a VACUUM (or 4 MB of backlog) merges it. On a mailbox
            // taking steady mail that list never drains and searches slow
            // from milliseconds to hundreds of them. Pay the index update
            // at insert time instead.
            let index_name = format!("gin_{table_name}_{column_name}");
            let create_index_query = format!(
                "CREATE INDEX IF NOT EXISTS {index_name} ON {table_name} USING GIN({column_name}) WITH (fastupdate = off)",
            );
            conn.execute(&create_index_query, &[])
                .await
                .map_err(into_error)?;
            // Indexes made before this change keep fastupdate=on
            disable_gin_fastupdate(conn, &index_name).await;
        }

        if field.is_indexed() {
            let column_name = field.sort_column().unwrap_or(field.column());
            let create_index_query = format!(
                "CREATE INDEX IF NOT EXISTS idx_{table_name}_{column_name} ON {table_name}({column_name})",
            );
            conn.execute(&create_index_query, &[])
                .await
                .map_err(into_error)?;
        }
    }

    Ok(())
}

/// inbuxa: turns fastupdate off on a GIN index made with the default and
/// merges the pending list it has built up. Idempotent: an index that already
/// has the option is left alone, so this costs one catalog read per index at
/// startup. A failure is logged and startup goes on, since search still works,
/// only slower.
async fn disable_gin_fastupdate(conn: &Object, index_name: &str) {
    if let Err(err) = try_disable_gin_fastupdate(conn, index_name).await {
        trc::event!(
            Store(trc::StoreEvent::PostgresqlError),
            Details = format!("Failed to turn off fastupdate on search index {index_name}"),
            Reason = err.to_string(),
        );
    }
}

async fn try_disable_gin_fastupdate(conn: &Object, index_name: &str) -> trc::Result<()> {
    let options = conn
        .query_opt(
            "SELECT COALESCE(reloptions, '{}')::text[] FROM pg_class WHERE oid = to_regclass($1)",
            &[&index_name],
        )
        .await
        .map_err(into_error)?
        .map(|row| row.try_get::<_, Vec<String>>(0))
        .transpose()
        .map_err(into_error)?;
    let Some(options) = options else {
        return Ok(());
    };
    if gin_fastupdate_is_off(&options) {
        return Ok(());
    }
    // SET (fastupdate) takes a SHARE UPDATE EXCLUSIVE lock, which doesn't
    // block reads or writes. Turning it off stops new entries going to the
    // pending list but doesn't flush the entries already there.
    conn.execute(
        &format!("ALTER INDEX {index_name} SET (fastupdate = off)"),
        &[],
    )
    .await
    .map_err(into_error)?;
    conn.query_one(
        "SELECT gin_clean_pending_list($1::text::regclass)",
        &[&index_name],
    )
    .await
    .map_err(into_error)?;
    Ok(())
}

/// Whether a relation's reloptions turn GIN's fastupdate off.
fn gin_fastupdate_is_off(options: &[String]) -> bool {
    options.iter().any(|option| {
        option.split_once('=').is_some_and(|(name, value)| {
            name.trim().eq_ignore_ascii_case("fastupdate")
                && matches!(
                    value.trim().to_ascii_lowercase().as_str(),
                    "off" | "false" | "no" | "0" | "f" | "n"
                )
        })
    })
}

async fn discover_ts_configs(pool: &Pool) -> AHashSet<&'static str> {
    let mut ts_configs = AHashSet::from_iter([PG_FALLBACK_LANG, PG_UNSTEMMED_LANG]);

    match probe_ts_configs(pool).await {
        Ok(available) => {
            for name in available {
                if let Some(config) = PG_LANGS.iter().copied().find(|config| *config == name) {
                    ts_configs.insert(config);
                }
            }
        }
        Err(err) => {
            trc::event!(
                Store(trc::StoreEvent::PostgresqlError),
                Details = "Failed to query pg_ts_config, assuming english only",
                Reason = err.to_string(),
            );
        }
    }

    ts_configs
}

async fn probe_ts_configs(pool: &Pool) -> trc::Result<Vec<String>> {
    let conn = pool.get().await.map_err(into_pool_error)?;

    conn.query("SELECT cfgname::text FROM pg_ts_config", &[])
        .await
        .map_err(into_error)?
        .into_iter()
        .map(|row| row.try_get::<_, String>(0).map_err(into_error))
        .collect()
}