// Copyright (c) 2023-2026 ParadeDB, Inc.
//
// This file is part of ParadeDB - Postgres for Search and Analytics
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
//! Provides a reference-counted wrapper around an open Postgres [`pg_sys::Relation`].
use crate::api::version::Version;
use crate::postgres::build::is_bm25_index;
use crate::postgres::options::BM25IndexOptions;
use crate::postgres::storage::metadata::MetaPage;
use crate::schema::SearchIndexSchema;
use pgrx::pg_sys::WalLevel::WAL_LEVEL_REPLICA;
use pgrx::{name_data_to_str, pg_sys, PgList, PgTupleDesc};
use std::cell::RefCell;
use std::error::Error;
use std::fmt::{Debug, Display, Formatter};
use std::ops::Deref;
use std::ptr::NonNull;
use std::rc::Rc;
use tantivy::TantivyError;

type NeedClose = bool;

#[repr(transparent)]
struct IsCreateIndex(Rc<RefCell<bool>>);
impl Default for IsCreateIndex {
    fn default() -> Self {
        Self(Rc::new(RefCell::new(false)))
    }
}
impl IsCreateIndex {
    fn set(&self, value: bool) {
        self.0.replace(value);
    }

    fn get(&self) -> bool {
        *self.0.borrow()
    }
}

#[repr(transparent)]
struct ForkNumber(Rc<RefCell<pg_sys::ForkNumber::Type>>);
impl Default for ForkNumber {
    fn default() -> Self {
        Self(Rc::new(RefCell::new(pg_sys::ForkNumber::MAIN_FORKNUM)))
    }
}
impl ForkNumber {
    fn set(&self, value: pg_sys::ForkNumber::Type) {
        self.0.replace(value);
    }

    fn get(&self) -> pg_sys::ForkNumber::Type {
        *self.0.borrow()
    }
}

/// A lazily-evaluated flag for if a relation needs WAL or not.
#[repr(transparent)]
struct NeedWal(Rc<RefCell<Option<bool>>>);
impl Default for NeedWal {
    /// Default WAL status is "unknown"
    fn default() -> Self {
        Self(Rc::new(RefCell::new(None)))
    }
}

impl NeedWal {
    /// Set if WAL is needed or not, replacing any previous value
    fn set(&self, value: bool) {
        self.0.replace(Some(value));
    }

    /// Indicates if WAL is needed or not.  If the decision is currently unknown, this
    /// will check to see if Postgres generally thinks the Relation itself needs WAL
    fn get(&self, rel: pg_sys::Relation) -> bool {
        *self
            .0
            .borrow_mut()
            .get_or_insert_with(|| relation_needs_wal(rel))
    }
}

#[derive(Debug, Clone)]
pub enum SchemaError {
    RelationNotBM25Index,
    Other(TantivyError),
}

impl Display for SchemaError {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::RelationNotBM25Index => write!(f, "relation is not a BM25 index"),
            Self::Other(e) => write!(f, "{e}"),
        }
    }
}

impl Error for SchemaError {}

/// Represents an opened Postgres relation to be used by pg_search.
///
/// [`PgSearchRelation`] is reference counted and will close the underlying
/// [`pg_sys::Relation`] when the last reference is dropped, accounting for
/// the state of the current transaction.
///
/// Instances of [`PgSearchRelation`] can be closed as necessary.
#[allow(clippy::type_complexity)]
#[derive(Clone)]
#[repr(transparent)]
pub struct PgSearchRelation(
    Option<
        Rc<(
            NonNull<pg_sys::RelationData>,
            NeedClose,
            Option<pg_sys::LOCKMODE>,
            RefCell<Option<Result<SearchIndexSchema, SchemaError>>>,
            BM25IndexOptions,
            IsCreateIndex,
            ForkNumber,
            NeedWal,
        )>,
    >,
);

impl Debug for PgSearchRelation {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PgSearchRelation")
            .field("relation", &self.oid())
            .field("lockmode", &self.lockmode())
            .finish()
    }
}

crate::impl_safe_drop!(PgSearchRelation, |self| {
    let Some(rc) = self.0.take() else {
        return;
    };
    let Some((relation, need_close, lockmode, ..)) = Rc::into_inner(rc) else {
        return;
    };
    unsafe {
        if need_close && pg_sys::IsTransactionState() {
            match lockmode {
                Some(lockmode) => pg_sys::relation_close(relation.as_ptr(), lockmode),
                None => pg_sys::RelationClose(relation.as_ptr()),
            }
        }
    }
});

impl Deref for PgSearchRelation {
    type Target = pg_sys::RelationData;

    fn deref(&self) -> &Self::Target {
        // SAFETY: the backing pointer is always correct for use by Rust as we couldn't have
        // gotten here otherwise
        unsafe { self.as_ptr().as_ref().unwrap_unchecked() }
    }
}

impl PgSearchRelation {
    /// Take ownership of a [`pg_sys::Relation`] pointer previously created by Postgres
    ///
    /// This relation will not be closed when we're dropped.
    pub unsafe fn from_pg(relation: pg_sys::Relation) -> Self {
        Self(Some(Rc::new((
            NonNull::new(relation)
                .expect("PgSearchRelation::from_pg: provided relation cannot be NULL"),
            false,
            None,
            Default::default(),
            BM25IndexOptions::from_relation(relation),
            IsCreateIndex::default(),
            ForkNumber::default(),
            NeedWal::default(),
        ))))
    }

    /// Open a relation with the specified [`pg_sys::Oid`].
    ///
    /// This relation will be closed when we're the last of our reference-counted clones to be dropped.
    pub fn open(oid: pg_sys::Oid) -> Self {
        unsafe {
            // SAFETY: RelationIdGetRelation() should return a valid RelationData pointer
            // unless no pg_class row could be found (suggesting the relation was just dropped)
            let relation = pg_sys::RelationIdGetRelation(oid);
            if relation.is_null() {
                panic!("relation not found, suggesting it was just dropped");
            }

            Self(Some(Rc::new((
                NonNull::new_unchecked(relation),
                true,
                None,
                Default::default(),
                BM25IndexOptions::from_relation(relation),
                IsCreateIndex::default(),
                ForkNumber::default(),
                NeedWal::default(),
            ))))
        }
    }

    /// Open a relation with the specified [`pg_sys::Oid`]
    ///
    /// Like [`Self::with_lock`], but fallible
    pub fn try_open(oid: pg_sys::Oid, lockmode: pg_sys::LOCKMODE) -> Option<Self> {
        // SAFETY: See `open`
        unsafe {
            let relation = pg_sys::try_relation_open(oid, lockmode);
            if relation.is_null() {
                None
            } else {
                Some(Self(Some(Rc::new((
                    NonNull::new_unchecked(relation),
                    true,
                    None,
                    Default::default(),
                    BM25IndexOptions::from_relation(relation),
                    IsCreateIndex::default(),
                    ForkNumber::default(),
                    NeedWal::default(),
                )))))
            }
        }
    }

    /// Open a relation with the specified [`pg_sys::Oid`] under the specified [`pg_sys::LOCKMODE`].
    ///
    /// This relation will be closed when we're the last of our reference-counted clones to be dropped.
    pub fn with_lock(oid: pg_sys::Oid, lockmode: pg_sys::LOCKMODE) -> Self {
        unsafe {
            // SAFETY: relation_open() always returns a valid RelationData pointer
            let relation = pg_sys::relation_open(oid, lockmode);
            Self(Some(Rc::new((
                NonNull::new_unchecked(relation),
                true,
                Some(lockmode),
                Default::default(),
                BM25IndexOptions::from_relation(relation),
                IsCreateIndex::default(),
                ForkNumber::default(),
                NeedWal::default(),
            ))))
        }
    }

    pub fn set_is_create_index(&mut self) {
        self.0.as_ref().unwrap().5.set(true);
    }

    pub fn is_create_index(&self) -> bool {
        self.0.as_ref().unwrap().5.get()
    }

    pub fn set_fork_number(&mut self, fork_number: pg_sys::ForkNumber::Type) {
        self.0.as_ref().unwrap().6.set(fork_number);
    }

    pub fn fork_number(&self) -> pg_sys::ForkNumber::Type {
        self.0.as_ref().unwrap().6.get()
    }

    /// Returns false if in the middle of a `REINDEX CONCURRENTLY`
    /// and the index is not yet ready to serve queries
    pub fn is_valid(&self) -> bool {
        unsafe { (*(*self.as_ptr()).rd_index).indisvalid }
    }

    /// Returns false if the index should not receive writes yet, such as during
    /// the early stages of `CREATE INDEX CONCURRENTLY` or `REINDEX CONCURRENTLY`.
    pub fn is_ready(&self) -> bool {
        unsafe { (*(*self.as_ptr()).rd_index).indisready }
    }

    /// Returns false if the index is in a concurrent drop phase and should not
    /// be touched by callers.
    pub fn is_live(&self) -> bool {
        unsafe { (*(*self.as_ptr()).rd_index).indislive }
    }

    /// Returns true when Postgres considers this index safe to inspect as a
    /// complete, usable index.
    pub fn is_usable(&self) -> bool {
        self.is_valid() && self.is_ready() && self.is_live()
    }

    /// Allows the user to decide for themselves if this [`PgSearchRelation`] instance (and all its
    /// clones) need to do WAL or not.
    pub fn set_need_wal(&mut self, need_wal: bool) {
        self.0.as_ref().unwrap().7.set(need_wal);
    }

    /// Returns true if Postgres thinks this relation needs WAL *or* instead returns the WAL-ness
    /// based on a prior call to `set_need_wal`.
    pub fn need_wal(&self) -> bool {
        self.0.as_ref().unwrap().7.get(self.as_ptr())
    }

    pub fn lockmode(&self) -> Option<pg_sys::LOCKMODE> {
        // SAFETY: self.0 is always Some
        unsafe { self.0.as_ref().unwrap_unchecked().2 }
    }

    pub fn oid(&self) -> pg_sys::Oid {
        // SAFETY: self.as_ptr() is always a valid pointer
        unsafe { (*self.as_ptr()).rd_id }
    }

    pub fn rel_oid(&self) -> Option<pg_sys::Oid> {
        if self.rd_index.is_null() {
            None
        } else {
            unsafe { Some((*self.rd_index).indrelid) }
        }
    }

    pub fn name(&self) -> &str {
        unsafe { name_data_to_str(&(*self.rd_rel).relname) }
    }

    pub fn namespace(&self) -> &str {
        unsafe {
            core::ffi::CStr::from_ptr(pg_sys::get_namespace_name((*self.rd_rel).relnamespace))
        }
        .to_str()
        .expect("unable to convert namespace name to UTF8")
    }

    pub fn tuple_desc(&self) -> PgTupleDesc<'_> {
        unsafe { PgTupleDesc::from_pg_unchecked(self.rd_att) }
    }

    pub fn reltuples(&self) -> Option<f32> {
        let reltuples = unsafe { (*self.rd_rel).reltuples };

        if reltuples == 0f32 {
            None
        } else {
            Some(reltuples)
        }
    }

    pub fn as_ptr(&self) -> pg_sys::Relation {
        // SAFETY: self.0 is always Some
        unsafe { self.0.as_ref().unwrap_unchecked().0.as_ptr() }
    }

    /// Bytes available to a non-privileged process on the filesystem backing this relation's
    /// tablespace. Returns `None` if the free space can't be determined.
    #[cfg(unix)]
    pub fn available_disk_bytes(&self) -> Option<u64> {
        let path = unsafe {
            let reltablespace = (*(*self.as_ptr()).rd_rel).reltablespace;
            let spc = if reltablespace == pg_sys::InvalidOid {
                pg_sys::MyDatabaseTableSpace
            } else {
                reltablespace
            };
            let db = if spc == pg_sys::GLOBALTABLESPACE_OID {
                pg_sys::InvalidOid
            } else {
                pg_sys::MyDatabaseId
            };

            // GetDatabasePath returns a palloc'd path relative to the data directory, which is the
            // backend's working directory; statvfs resolves it (and any tablespace symlink)
            // correctly.
            let path = pg_sys::GetDatabasePath(db, spc);
            if path.is_null() {
                return None;
            }
            let owned = std::ffi::CStr::from_ptr(path).to_owned();
            pg_sys::pfree(path.cast());
            owned
        };

        let stat = rustix::fs::statvfs(path.as_c_str()).ok()?;
        Some(stat.f_bavail.saturating_mul(stat.f_frsize))
    }

    #[cfg(not(unix))]
    pub fn available_disk_bytes(&self) -> Option<u64> {
        // No portable way to query filesystem free space here.
        None
    }

    pub fn heap_relation(&self) -> Option<PgSearchRelation> {
        if self.rd_index.is_null() {
            None
        } else {
            unsafe { Some(PgSearchRelation::open((*self.rd_index).indrelid)) }
        }
    }

    pub fn indices(&self, lockmode: pg_sys::LOCKMODE) -> impl Iterator<Item = PgSearchRelation> {
        // SAFETY: we know self.as_ptr() is a valid pointer as we created it
        let list =
            unsafe { PgList::<pg_sys::Oid>::from_pg(pg_sys::RelationGetIndexList(self.as_ptr())) };

        list.iter_oid()
            .filter(|oid| *oid != pg_sys::InvalidOid)
            .map(|oid| PgSearchRelation::with_lock(oid, lockmode))
            .collect::<Vec<_>>()
            .into_iter()
    }

    pub fn options(&self) -> &BM25IndexOptions {
        unsafe {
            // SAFETY: self.0 is always Some
            &self.0.as_ref().unwrap_unchecked().4
        }
    }

    pub fn schema(&self) -> Result<SearchIndexSchema, SchemaError> {
        let rc = self.0.as_ref().unwrap();
        let mut borrow = rc.3.borrow_mut();
        let schema = borrow.get_or_insert_with(|| {
            if !is_bm25_index(self) {
                return Err(SchemaError::RelationNotBM25Index);
            }

            SearchIndexSchema::open(self).map_err(SchemaError::Other)
        });

        match schema {
            Ok(schema) => Ok(schema.clone()),
            Err(e) => Err(e.clone()),
        }
    }

    /// This opens the MetaPage on every call, so use it carefully
    pub fn created_by_version(&self) -> Option<Version> {
        MetaPage::open(self).created_by_version()
    }

    /// Get the index info for this relation.
    pub fn index_info(&self) -> *mut pg_sys::IndexInfo {
        unsafe { pg_sys::BuildIndexInfo(self.as_ptr()) }
    }

    /// Extract index expressions from the index info.
    pub fn index_expressions(&self) -> PgList<pg_sys::Expr> {
        unsafe { PgList::<pg_sys::Expr>::from_pg((*self.index_info()).ii_Expressions) }
    }

    /// Check if a field supports aggregate pushdown.
    ///
    /// Returns `Ok(false)` for NUMERIC fields, `Ok(true)` for other fields,
    /// or an error if the schema cannot be loaded.
    pub fn field_supports_aggregate(&self, field: &str) -> Result<bool, SchemaError> {
        self.schema().map(|s| s.field_supports_aggregate(field))
    }
}

fn relation_needs_wal(relation: pg_sys::Relation) -> bool {
    // #define InvalidSubTransactionId		((SubTransactionId) 0)
    const INVALID_SUB_TRANSACTION_ID: pg_sys::SubTransactionId = 0;

    // /*
    //  * Is WAL-logging necessary for archival or log-shipping, or can we skip
    //  * WAL-logging if we fsync() the data before committing instead?
    //  */
    // #define XLogIsNeeded() (wal_level >= WAL_LEVEL_REPLICA)
    unsafe fn xlog_is_needed() -> bool {
        pg_sys::wal_level >= WAL_LEVEL_REPLICA as i32
    }

    // /*
    //  * RelationIsPermanent
    //  *		True if relation is permanent.
    //  */
    // #define RelationIsPermanent(relation) \
    // 	((relation)->rd_rel->relpersistence == RELPERSISTENCE_PERMANENT)

    unsafe fn relation_is_permanent(relation: pg_sys::Relation) -> bool {
        (*(*relation).rd_rel).relpersistence
            == pg_sys::RELPERSISTENCE_PERMANENT as core::ffi::c_char
    }

    // /*
    //  * RelationNeedsWAL
    //  *		True if relation needs WAL.
    //  *
    //  * Returns false if wal_level = minimal and this relation is created or
    //  * truncated in the current transaction.  See "Skipping WAL for New
    //  * RelFileLocator" in src/backend/access/transam/README.
    //  */
    // #define RelationNeedsWAL(relation)										\
    // 	(RelationIsPermanent(relation) && (XLogIsNeeded() ||				\
    // 	  (relation->rd_createSubid == InvalidSubTransactionId &&			\
    // 	   relation->rd_firstRelfilelocatorSubid == InvalidSubTransactionId)))

    #[cfg(feature = "pg15")]
    unsafe {
        relation_is_permanent(relation)
            && (xlog_is_needed()
                || ((*relation).rd_createSubid == INVALID_SUB_TRANSACTION_ID
                    && (*relation).rd_firstRelfilenodeSubid == INVALID_SUB_TRANSACTION_ID))
    }
    #[cfg(any(feature = "pg16", feature = "pg17", feature = "pg18"))]
    unsafe {
        relation_is_permanent(relation)
            && (xlog_is_needed()
                || ((*relation).rd_createSubid == INVALID_SUB_TRANSACTION_ID
                    && (*relation).rd_firstRelfilelocatorSubid == INVALID_SUB_TRANSACTION_ID))
    }
}
