rust-ef 2.0.0

//! `DbContext` struct and non-save methods.
//!
//! `DbContext` is the concrete context type. Entity sets use a type-map:
//! `ctx.set::<Blog>()` lazy-creates `DbSet<Blog>`. `SetOps<T>` dispatchers
//! enable `save_changes()` to iterate all entity types.
//!
//! ## Ownership and Mutation
//!
//! `DbContext` methods (`set::<T>()`, `save_changes()`, `detect_changes()`)
//! require `&mut self` — this is idiomatic Rust, not a limitation. The DI
//! integration (`add_dbcontext`) registers the context as **Scoped** and
//! supports owned resolution via `get_owned::<DbContext>()`.
//!
//! `DbContext` is **not** thread-safe — a single instance must not be shared
//! across threads (aligned with EFCore).

use crate::db_set::DbSet;
use crate::entity::{
    EntityState, IEntitySnapshot, IEntityType, IFromRow, IGetKeyValues, INavigationSetter,
};
use crate::error::{EFError, EFResult};
use crate::interceptor::InterceptorPipeline;
use crate::metadata::EntityTypeMeta;
use crate::migration::MigrationEngine;
use crate::model_builder::ModelBuilder;
use crate::provider::{DbValue, IDatabaseProvider, IsolationLevel};
use crate::session::{Session, SessionConnection, SessionProvider};
use crate::tracking::ChangeTracker;
use crate::transaction::{DbTransaction, ITransaction};
use std::any::{Any, TypeId};
use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;

use super::{DbContextOptions, ErasedSetOps, SetOps};

/// The session / unit-of-work entry point.
///
/// Entity sets are lazy-created via `ctx.set::<T>()`. Metadata is loaded from
/// the process-level `MetadataCache` on `DbContextOptions` (built once per
/// `context_key`, then `Arc`-shared). `save_changes()` commits all pending
/// changes in a transaction.
pub struct DbContext {
    pub(crate) sets: HashMap<TypeId, Box<dyn Any + Send + Sync>>,
    pub(crate) savers: HashMap<TypeId, Box<dyn ErasedSetOps>>,
    pub(crate) entity_metas: HashMap<TypeId, EntityTypeMeta>,
    pub(crate) model_builder: ModelBuilder,
    pub(crate) change_tracker: ChangeTracker,
    pub(crate) provider: Arc<dyn IDatabaseProvider>,
    /// Keyed context key (`None` for the default model). Used only for
    /// actionable diagnostics when a keyed replica forgets
    /// `use_default_entity_model()`.
    pub(crate) context_key: Option<String>,
    /// Per-context connection session. While an explicit transaction is open,
    /// every statement issued through `provider` borrows the transaction's
    /// connection instead of acquiring a new pooled connection.
    pub(crate) session: Arc<Session>,
    pub(crate) interceptor_pipeline: InterceptorPipeline,
    pub(crate) lazy_loading_enabled: bool,
}

impl DbContext {
    /// Creates the context from options (uses the provider factory stored in options).
    ///
    /// Entity metadata is loaded from the process-level `MetadataCache` on
    /// `DbContextOptions` — `inventory::iter` + `IEntityTypeConfiguration::configure()`
    /// run once per `context_key` (first call), then the result is `Arc`-shared
    /// across all `DbContext` instances. Per-instance `ModelBuilder` mutations
    /// (`has_query_filter`, etc.) after construction only affect this instance.
    pub fn from_options(options: &DbContextOptions) -> EFResult<Self> {
        let write_provider = options.create_provider()?;
        let session = Arc::new(Session::new(write_provider));
        let provider: Arc<dyn IDatabaseProvider> =
            Arc::new(SessionProvider::new(Arc::clone(&session)));
        let built = options
            .metadata_cache
            .get_or_build(options.context_key.as_deref());
        let ctx = Self {
            sets: HashMap::new(),
            savers: HashMap::new(),
            entity_metas: built.entity_metas.clone(),
            model_builder: ModelBuilder::from_built(&built),
            change_tracker: ChangeTracker::new(),
            provider,
            context_key: options.context_key.clone(),
            session,
            interceptor_pipeline: InterceptorPipeline::new(options.interceptors.clone()),
            lazy_loading_enabled: options.lazy_loading_enabled,
        };
        Ok(ctx)
    }

    pub fn set<T>(&mut self) -> &mut DbSet<T>
    where
        T: IEntityType
            + IEntitySnapshot
            + IGetKeyValues
            + IFromRow
            + INavigationSetter
            + Send
            + Sync
            + 'static,
    {
        let type_id = TypeId::of::<T>();
        self.savers
            .entry(type_id)
            .or_insert_with(|| Box::new(SetOps::<T>::new()));
        self.entity_metas
            .entry(type_id)
            .or_insert_with(T::entity_meta);
        if !self.model_builder.has_entity(type_id) {
            self.model_builder.register_entity_meta(T::entity_meta());
        }
        self.sets.entry(type_id).or_insert_with(|| {
            let table_name = self
                .model_builder
                .build()
                .into_iter()
                .find(|m| m.type_id == type_id)
                .map(|m| m.table_name.to_string())
                .unwrap_or_else(|| T::entity_meta().table_name.to_string());
            let mut db_set = DbSet::<T>::with_provider(table_name, Arc::clone(&self.provider));
            if let Some(filter) = self.model_builder.get_query_filter(&type_id) {
                db_set.set_query_filter(filter.clone());
            }
            db_set.set_filter_map(self.model_builder.filters_by_table());
            db_set.set_lazy_loading_enabled(self.lazy_loading_enabled);
            Box::new(db_set)
        });
        self.sets
            .get_mut(&type_id)
            .and_then(|b| b.downcast_mut::<DbSet<T>>())
            .expect("DbSet type mismatch")
    }

    /// Returns the model builder for Fluent API configuration.
    pub fn model(&mut self) -> &mut ModelBuilder {
        &mut self.model_builder
    }

    /// Returns a read-only reference to the model builder.
    pub fn model_builder(&self) -> &ModelBuilder {
        &self.model_builder
    }

    /// Returns `true` if an entity of type `T` has been discovered and
    /// registered in the entity metadata map.
    pub fn entity_metas_contains<T: IEntityType>(&self) -> bool {
        self.entity_metas.contains_key(&TypeId::of::<T>())
    }

    /// No-op since metadata caching was introduced.
    ///
    /// Historically, this method iterated `inventory::iter` to register entity
    /// metadata and apply `#[entity(T)]` configurations. That work now happens
    /// in `MetadataCache::build()`, invoked lazily by `from_options()` via
    /// `MetadataCache::get_or_build()`. The result is `Arc`-shared across all
    /// `DbContext` instances with the same `context_key`.
    ///
    /// Retained as a public method for backward compatibility — existing code
    /// calling `ctx.discover_entities()?` after `from_options()` continues to
    /// compile and run (as a no-op). The metadata is already populated.
    pub fn discover_entities(&mut self) -> EFResult<()> {
        Ok(())
    }

    /// Detects changes on all tracked DbSets by comparing property snapshots.
    pub fn detect_changes(&mut self) {
        let type_ids: Vec<TypeId> = self.sets.keys().copied().collect();
        for type_id in type_ids {
            let Some(set) = self.sets.get_mut(&type_id) else {
                continue;
            };
            let Some(saver) = self.savers.get(&type_id) else {
                continue;
            };
            saver.detect_changes(set.as_mut(), &mut self.change_tracker);
        }
    }

    /// Creates all tables for registered entity types.
    ///
    /// Sources metas from `model_builder.build()`, which applies all Fluent
    /// API configurations and `#[entity(T)]` overrides. Entities are
    /// discovered automatically via `#[derive(EntityType)]`; call
    /// `discover_entities()` first, or use `set::<T>()` to register manually.
    pub async fn ensure_created(&self) -> EFResult<()> {
        let mut metas: Vec<EntityTypeMeta> = self.model_builder.build();
        if metas.is_empty() {
            metas = self.entity_metas.values().cloned().collect();
        }
        if metas.is_empty() {
            return Err(self.empty_model_error("ensure_created"));
        }
        let dialect = self.provider.migration_dialect();
        MigrationEngine::new(dialect)
            .ensure_created(&*self.provider, &metas)
            .await?;

        for meta in &metas {
            let rows = self.model_builder.seed_rows_for(&meta.type_id);
            if !rows.is_empty() {
                MigrationEngine::new(dialect)
                    .apply_seed_data(&*self.provider, meta, rows)
                    .await?;
            }
        }
        Ok(())
    }

    /// Drops all tables for registered entity types.
    pub async fn ensure_deleted(&self) -> EFResult<()> {
        let mut metas: Vec<EntityTypeMeta> = self.model_builder.build();
        if metas.is_empty() {
            metas = self.entity_metas.values().cloned().collect();
        }
        if metas.is_empty() {
            return Err(self.empty_model_error("ensure_deleted"));
        }
        let dialect = self.provider.migration_dialect();
        MigrationEngine::new(dialect)
            .ensure_deleted(&*self.provider, &metas)
            .await
    }

    /// Actionable error for a context that discovered no entities.
    fn empty_model_error(&self, operation: &str) -> EFError {
        let mut message = format!(
            "No entity types registered. Call ctx.set::<T>() or ctx.discover_entities() before {operation}()."
        );
        if let Some(key) = &self.context_key {
            message.push_str(&format!(
                " This is keyed context '{key}', which only discovers entities tagged #[context(\"{key}\")]. \
                 For a read/write replica sharing the default schema, call \
                 DbContextOptionsBuilder::use_default_entity_model()."
            ));
        }
        EFError::configuration(message)
    }
}

// ---------------------------------------------------------------------------
// DbContext inherent methods (formerly on IDbContext / IDbContextExt traits)
// ---------------------------------------------------------------------------

impl DbContext {
    /// Returns the database provider.
    pub fn provider(&self) -> &dyn IDatabaseProvider {
        &*self.provider
    }

    /// Returns a read-only reference to the change tracker.
    pub fn change_tracker(&self) -> &ChangeTracker {
        &self.change_tracker
    }

    /// Returns a mutable reference to the change tracker.
    pub fn change_tracker_mut(&mut self) -> &mut ChangeTracker {
        &mut self.change_tracker
    }

    /// Detaches every tracked entity and clears the change tracker.
    ///
    /// Call this after a rollback when the handler is not going to retry, so
    /// the rolled-back entities cannot leak into a later `save_changes`.
    /// Entity instances themselves are not modified — only the context's
    /// tracking state is dropped.
    pub fn clear_tracking(&mut self) {
        let type_ids: Vec<TypeId> = self.sets.keys().copied().collect();
        for type_id in type_ids {
            let Some(saver) = self.savers.get(&type_id) else {
                continue;
            };
            let Some(set) = self.sets.get_mut(&type_id) else {
                continue;
            };
            saver.clear_entries(set.as_mut());
        }
        self.change_tracker.clear();
    }

    /// Executes a raw SQL query and materializes the result rows into entities.
    ///
    /// This is the escape hatch for complex queries (multi-table JOINs, CTEs,
    /// window functions) that are hard to express via LINQ. The caller is
    /// responsible for SQL correctness and parameterization.
    ///
    /// # Example
    /// ```rust,ignore
    /// let blogs: Vec<Blog> = ctx
    ///     .sql_query("SELECT * FROM blogs WHERE id = ?", &[DbValue::I32(1)])
    ///     .await?;
    /// ```
    pub async fn sql_query<T: IFromRow + IEntityType>(
        &self,
        sql: &str,
        params: &[DbValue],
    ) -> EFResult<Vec<T>> {
        let mut conn = self.provider.get_connection().await?;
        let rows = conn.query(sql, params).await?;
        crate::entity::materialize_entities(&rows)
    }

    /// Begins a context-scoped transaction and returns its handle.
    ///
    /// This is the EF Core-style API: the transaction is registered on the
    /// context, so queries, navigation loading, bulk DML, raw SQL and
    /// `save_changes()` all use the transaction connection. Nothing is
    /// committed or rolled back until the returned handle is consumed.
    ///
    /// ```rust,ignore
    /// let txn = ctx.begin_transaction().await?;
    /// ctx.add::<Account>(account);
    /// ctx.save_changes().await?; // joins, does not commit
    /// txn.commit().await?;       // or txn.rollback().await?
    /// ```
    ///
    /// Dropping the handle without committing rolls back (best effort).
    /// Nested transactions are not supported.
    pub async fn begin_transaction(&mut self) -> EFResult<Box<dyn ITransaction>> {
        self.begin_transaction_with_opt(None).await
    }

    /// Like [`DbContext::begin_transaction`], with an isolation level applied
    /// at transaction start.
    pub async fn begin_transaction_with(
        &mut self,
        isolation: IsolationLevel,
    ) -> EFResult<Box<dyn ITransaction>> {
        self.begin_transaction_with_opt(Some(isolation)).await
    }

    async fn begin_transaction_with_opt(
        &mut self,
        isolation: Option<IsolationLevel>,
    ) -> EFResult<Box<dyn ITransaction>> {
        let shared = self.session.begin_with(isolation).await?;
        let facade = SessionConnection::new(Arc::clone(&shared));
        Ok(Box::new(DbTransaction::new_shared(
            shared,
            facade,
            Arc::clone(&self.session),
        )))
    }

    /// Executes a closure within a context-scoped transaction.
    ///
    /// Registers the transaction on the context session for the duration of
    /// `f`, so every statement issued through the context — queries, navigation
    /// loading, bulk DML, raw SQL and `save_changes()` — reuses the same
    /// transaction connection. Commits on `Ok`, rolls back on `Err`.
    ///
    /// For a business rejection that must roll back and still return a value,
    /// use the manual [`DbContext::begin_transaction`] handle and call
    /// `rollback()` explicitly.
    ///
    /// The closure receives `&mut DbContext` and must return a pinned boxed
    /// future. This signature works around Rust's async borrow checker by
    /// letting the closure capture `ctx` by mutable reference while still
    /// producing a `Send` future.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// ctx.use_transaction(|ctx| Box::pin(async move {
    ///     ctx.add::<Blog>(blog);
    ///     ctx.save_changes().await?;
    ///     Ok(())
    /// })).await?;
    /// ```
    pub async fn use_transaction<F, R>(&mut self, f: F) -> EFResult<R>
    where
        for<'a> F: FnOnce(&'a mut Self) -> Pin<Box<dyn Future<Output = EFResult<R>> + Send + 'a>>,
        R: Send + 'static,
    {
        self.use_transaction_with_opt(None, f).await
    }

    /// Like [`DbContext::use_transaction`], with an explicit isolation level
    /// applied at transaction start.
    pub async fn use_transaction_with<F, R>(
        &mut self,
        isolation: IsolationLevel,
        f: F,
    ) -> EFResult<R>
    where
        for<'a> F: FnOnce(&'a mut Self) -> Pin<Box<dyn Future<Output = EFResult<R>> + Send + 'a>>,
        R: Send + 'static,
    {
        self.use_transaction_with_opt(Some(isolation), f).await
    }

    async fn use_transaction_with_opt<F, R>(
        &mut self,
        isolation: Option<IsolationLevel>,
        f: F,
    ) -> EFResult<R>
    where
        for<'a> F: FnOnce(&'a mut Self) -> Pin<Box<dyn Future<Output = EFResult<R>> + Send + 'a>>,
        R: Send + 'static,
    {
        if self.session.is_active().await {
            return Err(EFError::transaction(
                "a transaction is already active on this DbContext",
            ));
        }
        let shared = self.session.begin_with(isolation).await?;
        let facade = SessionConnection::new(Arc::clone(&shared));
        let txn: Box<dyn ITransaction> = Box::new(DbTransaction::new_shared(
            shared,
            facade,
            Arc::clone(&self.session),
        ));
        // `txn` is a local guard: if `f` returns Err, panics, or is cancelled,
        // its Drop discards the connection and rolls back.
        let result = f(self).await;
        match result {
            Ok(r) => {
                txn.commit().await?;
                Ok(r)
            }
            Err(e) => {
                let _ = txn.rollback().await;
                Err(e)
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Entity mutation methods — coordinate DbSet + ChangeTracker
// ---------------------------------------------------------------------------

impl DbContext {
    /// Ensures the entity type is registered (saver, metadata, model builder).
    /// Returns `(type_id, type_name)`.
    fn ensure_registered<T>(&mut self) -> (TypeId, String)
    where
        T: IEntityType
            + IEntitySnapshot
            + IGetKeyValues
            + IFromRow
            + INavigationSetter
            + Send
            + Sync
            + 'static,
    {
        let type_id = TypeId::of::<T>();
        self.savers
            .entry(type_id)
            .or_insert_with(|| Box::new(SetOps::<T>::new()));
        self.entity_metas
            .entry(type_id)
            .or_insert_with(T::entity_meta);
        if !self.model_builder.has_entity(type_id) {
            self.model_builder.register_entity_meta(T::entity_meta());
        }
        let type_name = self.entity_metas[&type_id].type_name.to_string();
        (type_id, type_name)
    }

    /// Adds a new entity in `Added` state. The entity will be INSERTed during
    /// `save_changes()`.
    pub fn add<T>(&mut self, entity: T)
    where
        T: IEntityType
            + IEntitySnapshot
            + IGetKeyValues
            + IFromRow
            + INavigationSetter
            + Send
            + Sync
            + 'static,
    {
        let (type_id, type_name) = self.ensure_registered::<T>();
        let entry_id =
            self.change_tracker
                .track(type_id, &type_name, EntityState::Added, None, false);
        self.set::<T>().push_entry(entity, entry_id);
    }

    /// Attaches an existing entity in `Unchanged` state with a snapshot for
    /// future change detection.
    pub fn attach<T>(&mut self, entity: T)
    where
        T: IEntityType
            + IEntitySnapshot
            + IGetKeyValues
            + IFromRow
            + INavigationSetter
            + Send
            + Sync
            + 'static,
    {
        let (type_id, type_name) = self.ensure_registered::<T>();
        let snapshot = entity.snapshot();
        let entry_id = self.change_tracker.track(
            type_id,
            &type_name,
            EntityState::Unchanged,
            Some(snapshot),
            false,
        );
        self.set::<T>().push_entry(entity, entry_id);
    }

    /// Marks an entity as `Modified` — all fields will be UPDATEd during
    /// `save_changes()`.
    pub fn update<T>(&mut self, entity: T)
    where
        T: IEntityType
            + IEntitySnapshot
            + IGetKeyValues
            + IFromRow
            + INavigationSetter
            + Send
            + Sync
            + 'static,
    {
        let (type_id, type_name) = self.ensure_registered::<T>();
        let entry_id =
            self.change_tracker
                .track(type_id, &type_name, EntityState::Modified, None, false);
        self.set::<T>().push_entry(entity, entry_id);
    }

    /// Marks an entity for upsert (INSERT ... ON CONFLICT DO UPDATE).
    pub fn upsert<T>(&mut self, entity: T)
    where
        T: IEntityType
            + IEntitySnapshot
            + IGetKeyValues
            + IFromRow
            + INavigationSetter
            + Send
            + Sync
            + 'static,
    {
        let (type_id, type_name) = self.ensure_registered::<T>();
        let entry_id =
            self.change_tracker
                .track(type_id, &type_name, EntityState::Added, None, true);
        self.set::<T>().push_entry(entity, entry_id);
    }

    /// Marks the entity at the given index as `Deleted`.
    pub fn remove_at<T>(&mut self, index: usize) -> EFResult<()>
    where
        T: IEntityType
            + IEntitySnapshot
            + IGetKeyValues
            + IFromRow
            + INavigationSetter
            + Send
            + Sync
            + 'static,
    {
        let entry_id = self
            .set::<T>()
            .entry_id_at(index)
            .ok_or_else(|| EFError::not_found("Entity not found at the given index".to_string()))?;
        self.change_tracker
            .set_state(entry_id, EntityState::Deleted);
        Ok(())
    }

    /// Marks all entities of type `T` as `Deleted`.
    pub fn remove_all<T>(&mut self)
    where
        T: IEntityType
            + IEntitySnapshot
            + IGetKeyValues
            + IFromRow
            + INavigationSetter
            + Send
            + Sync
            + 'static,
    {
        let entry_ids: Vec<u64> = self.set::<T>().entries.iter().map(|e| e.entry_id).collect();
        for entry_id in entry_ids {
            self.change_tracker
                .set_state(entry_id, EntityState::Deleted);
        }
    }

    /// Loads all rows from the database into the DbSet as `Unchanged`.
    /// Clears existing entries and tracker state for type `T` first.
    pub async fn load_all<T>(&mut self) -> EFResult<()>
    where
        T: IEntityType
            + IEntitySnapshot
            + IGetKeyValues
            + IFromRow
            + INavigationSetter
            + crate::entity::ILazyInit
            + Send
            + Sync
            + 'static,
    {
        let type_id = TypeId::of::<T>();
        let items = self.set::<T>().query().to_list().await?;
        self.set::<T>().clear_entries();
        self.change_tracker.clear_by_type(type_id);
        for item in items {
            self.attach::<T>(item);
        }
        Ok(())
    }
}