drizzle 0.3.0

use std::{
    borrow::Cow,
    marker::PhantomData,
    sync::{
        Arc, Mutex,
        atomic::{AtomicU64, Ordering},
    },
};

use drizzle_core::{
    param::{OwnedParam, Param},
    prepared::{
        OwnedPreparedStatement as CoreOwnedPreparedStatement,
        PreparedStatement as CorePreparedStatement,
    },
    traits::ToSQL,
};
use drizzle_postgres::values::{OwnedPostgresValue, PostgresValue};
use postgres::{
    Client, Row, Statement,
    types::{ToSql, Type},
};

use crate::builder::postgres::prepared_common::postgres_prepared_sync_impl;

const STATEMENT_CACHE_CAP: usize = 32;

static NEXT_CLIENT_ID: AtomicU64 = AtomicU64::new(1);
static REGISTERED_CLIENTS: Mutex<Vec<ClientRegistration>> = Mutex::new(Vec::new());

#[derive(Clone, Copy)]
struct ClientRegistration {
    client_key: usize,
    client_id: u64,
}

pub(crate) fn next_client_id() -> u64 {
    NEXT_CLIENT_ID.fetch_add(1, Ordering::Relaxed)
}

pub(crate) fn register_client(client: &Client, client_id: u64) {
    let client_key = client as *const Client as usize;
    let mut registrations = REGISTERED_CLIENTS
        .lock()
        .unwrap_or_else(|err| err.into_inner());

    registrations.retain(|entry| entry.client_id == client_id || entry.client_key != client_key);
    if let Some(entry) = registrations
        .iter_mut()
        .find(|entry| entry.client_id == client_id)
    {
        entry.client_key = client_key;
    } else {
        registrations.push(ClientRegistration {
            client_key,
            client_id,
        });
    }
}

pub(crate) fn unregister_client(client_id: u64) {
    let mut registrations = REGISTERED_CLIENTS
        .lock()
        .unwrap_or_else(|err| err.into_inner());
    registrations.retain(|entry| entry.client_id != client_id);
}

fn registered_client_id(client: &Client) -> Option<u64> {
    let client_key = client as *const Client as usize;
    let registrations = REGISTERED_CLIENTS
        .lock()
        .unwrap_or_else(|err| err.into_inner());
    registrations
        .iter()
        .find(|entry| entry.client_key == client_key)
        .map(|entry| entry.client_id)
}

/// A query rendered once, ready to run many times with new placeholder values.
///
/// Made by `.prepare()` on a query builder. Run it with `execute`, `all`, or
/// `get`, passing the client and the placeholder bindings. It borrows the
/// values embedded in the query; call [`into_owned`](Self::into_owned) to
/// store it.
#[derive(Debug, Clone)]
pub struct PreparedStatement<'a, Marker = (), DecodedRow = ()> {
    pub(crate) inner: CorePreparedStatement<'a, PostgresValue<'a>>,
    pub(crate) statement_cache: StatementCache,
    pub(crate) marker: PhantomData<(Marker, DecodedRow)>,
}

#[derive(Clone, Default)]
pub(crate) struct StatementCache(Arc<Mutex<Vec<CachedStatement>>>);

struct CachedStatement {
    client_id: u64,
    sql: Box<str>,
    param_types: Box<[Type]>,
    statement: Statement,
}

impl std::fmt::Debug for StatementCache {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("StatementCache").finish_non_exhaustive()
    }
}

/// Whether `error` rejected a cached statement that no longer fits its
/// session: the schema changed under it ("cached plan must not change result
/// type") or the connection behind it was replaced ("prepared statement ...
/// does not exist"). Both come from planning, before the statement runs.
pub(crate) fn is_stale_statement(error: &postgres::Error) -> bool {
    use postgres::error::SqlState;
    match error.code() {
        Some(code) if *code == SqlState::INVALID_SQL_STATEMENT_NAME => true,
        Some(code) if *code == SqlState::FEATURE_NOT_SUPPORTED => {
            error.as_db_error().is_some_and(|db| {
                db.message()
                    .contains("cached plan must not change result type")
            })
        }
        _ => false,
    }
}

impl StatementCache {
    /// Drops the statement cached for `sql` on `client`, so the next use
    /// prepares it again.
    pub(crate) fn evict_for(&self, client: &Client, sql: &str, param_types: &[Type]) {
        let Some(client_id) = registered_client_id(client) else {
            return;
        };
        let mut cache = self.0.lock().unwrap_or_else(|err| err.into_inner());
        cache.retain(|cached| {
            !(cached.client_id == client_id
                && cached.sql.as_ref() == sql
                && cached.param_types.as_ref() == param_types)
        });
    }

    /// Drops the cached statement for `sql`, so the next use prepares it
    /// again.
    pub(crate) fn evict(&self, client_id: u64, sql: &str, param_types: &[Type]) {
        let mut cache = self.0.lock().unwrap_or_else(|err| err.into_inner());
        cache.retain(|cached| {
            !(cached.client_id == client_id
                && cached.sql.as_ref() == sql
                && cached.param_types.as_ref() == param_types)
        });
    }

    /// Drops every statement cached for a connection: after its schema
    /// changed, or when the connection itself may be replaced.
    pub(crate) fn clear_client(&self, client_id: u64) {
        let mut cache = self.0.lock().unwrap_or_else(|err| err.into_inner());
        cache.retain(|cached| cached.client_id != client_id);
    }

    pub(crate) fn statement(
        &self,
        client: &mut Client,
        sql: &str,
        param_types: &[Type],
    ) -> Result<Statement, postgres::Error> {
        let Some(client_id) = registered_client_id(client) else {
            return client.prepare_typed(sql, param_types);
        };

        if let Some(statement) = self.lookup(client_id, sql, param_types) {
            return Ok(statement);
        }
        let statement = client.prepare_typed(sql, param_types)?;
        Ok(self.store(client_id, sql, param_types, statement))
    }

    /// Resolves a statement for a query running inside a transaction.
    ///
    /// Shares the connection's cache rather than keeping a per-transaction one:
    /// a `Statement` is bound to the *session*, not the transaction, and
    /// PostgreSQL does not roll back the prepared-statement catalog. A statement
    /// first prepared inside a transaction therefore stays valid after that
    /// transaction commits, rolls back, or rolls back to a savepoint, so later
    /// transactions and the connection runner reuse the same Parse.
    pub(crate) fn transaction_statement(
        &self,
        client_id: u64,
        tx: &mut postgres::Transaction<'_>,
        sql: &str,
        param_types: &[Type],
    ) -> Result<Statement, postgres::Error> {
        if let Some(statement) = self.lookup(client_id, sql, param_types) {
            return Ok(statement);
        }
        let statement = tx.prepare_typed(sql, param_types)?;
        Ok(self.store(client_id, sql, param_types, statement))
    }

    fn lookup(&self, client_id: u64, sql: &str, param_types: &[Type]) -> Option<Statement> {
        let mut cache = self.0.lock().unwrap_or_else(|err| err.into_inner());
        let pos = cache.iter().position(|cached| {
            cached.client_id == client_id
                && cached.sql.as_ref() == sql
                && cached.param_types.as_ref() == param_types
        })?;
        let cached = cache.remove(pos);
        let statement = cached.statement.clone();
        cache.insert(0, cached);
        Some(statement)
    }

    /// Inserts `statement` unless a concurrent caller already cached this key,
    /// in which case the already-cached statement wins and is returned.
    fn store(
        &self,
        client_id: u64,
        sql: &str,
        param_types: &[Type],
        statement: Statement,
    ) -> Statement {
        let mut cache = self.0.lock().unwrap_or_else(|err| err.into_inner());
        if let Some(pos) = cache.iter().position(|cached| {
            cached.client_id == client_id
                && cached.sql.as_ref() == sql
                && cached.param_types.as_ref() == param_types
        }) {
            let cached = cache.remove(pos);
            let statement = cached.statement.clone();
            cache.insert(0, cached);
            return statement;
        }
        cache.insert(
            0,
            CachedStatement {
                client_id,
                sql: sql.into(),
                param_types: param_types.into(),
                statement: statement.clone(),
            },
        );
        cache.truncate(STATEMENT_CACHE_CAP);
        statement
    }
}

impl<Marker, DecodedRow> From<OwnedPreparedStatement<Marker, DecodedRow>>
    for PreparedStatement<'_, Marker, DecodedRow>
{
    fn from(value: OwnedPreparedStatement<Marker, DecodedRow>) -> Self {
        let postgres_params = value.inner.params.iter().map(|v| {
            Param::new(
                v.placeholder,
                v.value.clone().map(|v| Cow::Owned(PostgresValue::from(v))),
            )
        });
        let inner = CorePreparedStatement {
            text_segments: value.inner.text_segments,
            params: postgres_params.collect::<Box<[_]>>(),
            sql: value.inner.sql,
        };
        PreparedStatement {
            inner,
            statement_cache: value.statement_cache,
            marker: PhantomData,
        }
    }
}

impl<'a, Marker, DecodedRow> PreparedStatement<'a, Marker, DecodedRow> {
    pub(crate) fn new(inner: CorePreparedStatement<'a, PostgresValue<'a>>) -> Self {
        Self {
            inner,
            statement_cache: StatementCache::default(),
            marker: PhantomData,
        }
    }

    /// Returns the rendered SQL, with `$n` placeholders.
    pub fn sql(&self) -> &str {
        self.inner.sql()
    }

    pub(crate) fn driver_statement(
        &self,
        client: &mut Client,
        sql: &str,
        param_types: &[Type],
    ) -> Result<Statement, postgres::Error> {
        self.statement_cache.statement(client, sql, param_types)
    }

    pub(crate) fn evict_statement(&self, client: &Client, sql: &str, param_types: &[Type]) {
        self.statement_cache.evict_for(client, sql, param_types);
    }

    /// Returns how many `$n` parameters the SQL has, counting embedded values
    /// as well as placeholders.
    pub fn param_count(&self) -> usize {
        self.inner.params.len()
    }

    /// Copies the embedded values so the statement no longer borrows them.
    pub fn into_owned(self) -> OwnedPreparedStatement<Marker, DecodedRow> {
        let owned_params = self
            .inner
            .params
            .into_vec()
            .into_iter()
            .map(|p| OwnedParam {
                placeholder: p.placeholder,
                value: p.value.map(|v| OwnedPostgresValue::from(v.into_owned())),
            });

        let inner = CoreOwnedPreparedStatement {
            text_segments: self.inner.text_segments,
            params: owned_params.collect::<Box<[_]>>(),
            sql: self.inner.sql,
        };

        OwnedPreparedStatement {
            inner,
            statement_cache: self.statement_cache,
            marker: PhantomData,
        }
    }
}

/// A [`PreparedStatement`] that owns its embedded values, so it can be stored
/// or moved freely.
#[derive(Debug, Clone)]
pub struct OwnedPreparedStatement<Marker = (), DecodedRow = ()> {
    pub(crate) inner: CoreOwnedPreparedStatement<OwnedPostgresValue>,
    pub(crate) statement_cache: StatementCache,
    pub(crate) marker: PhantomData<(Marker, DecodedRow)>,
}

impl<'a, Marker, DecodedRow> From<PreparedStatement<'a, Marker, DecodedRow>>
    for OwnedPreparedStatement<Marker, DecodedRow>
{
    fn from(value: PreparedStatement<'a, Marker, DecodedRow>) -> Self {
        value.into_owned()
    }
}

impl<Marker, DecodedRow> OwnedPreparedStatement<Marker, DecodedRow> {
    /// Returns the rendered SQL, with `$n` placeholders.
    pub fn sql(&self) -> &str {
        self.inner.sql()
    }

    /// Returns how many `$n` parameters the SQL has, counting embedded values
    /// as well as placeholders.
    pub fn param_count(&self) -> usize {
        self.inner.params.len()
    }

    pub(crate) fn driver_statement(
        &self,
        client: &mut Client,
        sql: &str,
        param_types: &[Type],
    ) -> Result<Statement, postgres::Error> {
        self.statement_cache.statement(client, sql, param_types)
    }

    pub(crate) fn evict_statement(&self, client: &Client, sql: &str, param_types: &[Type]) {
        self.statement_cache.evict_for(client, sql, param_types);
    }
}

postgres_prepared_sync_impl!(Client, Row, ToSql);

impl<Marker, DecodedRow> std::fmt::Display for PreparedStatement<'_, Marker, DecodedRow> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.inner)
    }
}

impl<Marker, DecodedRow> std::fmt::Display for OwnedPreparedStatement<Marker, DecodedRow> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.inner)
    }
}

impl<'a, Marker, DecodedRow> ToSQL<'a, PostgresValue<'a>>
    for PreparedStatement<'a, Marker, DecodedRow>
{
    fn to_sql(&self) -> drizzle_core::sql::SQL<'a, PostgresValue<'a>> {
        self.inner.to_sql()
    }
}

impl<'a, Marker, DecodedRow> ToSQL<'a, OwnedPostgresValue>
    for OwnedPreparedStatement<Marker, DecodedRow>
{
    fn to_sql(&self) -> drizzle_core::sql::SQL<'a, OwnedPostgresValue> {
        self.inner.to_sql()
    }
}

impl<'a, Marker, DecodedRow> ToSQL<'a, PostgresValue<'a>>
    for OwnedPreparedStatement<Marker, DecodedRow>
{
    fn to_sql(&self) -> drizzle_core::sql::SQL<'a, PostgresValue<'a>> {
        self.inner.to_sql().map_params(PostgresValue::from)
    }
}