lix 0.17.1

use std::collections::BTreeSet;
use std::ptr::NonNull;
use std::sync::Arc;

use async_trait::async_trait;
use serde_json::Value as JsonValue;
use tokio::sync::Mutex;

use crate::LixError;
use crate::binary_cas::{BlobBytesBatch, BlobDataReader, BlobId};
use crate::branch::{BranchHead, BranchRefReader};
use crate::changelog::CommitId;
use crate::commit_graph::CommitGraphReader;
use crate::filesystem::{
    FilesystemPathIndex, FilesystemPathIndexReader, FilesystemPathIndexRequest,
    UncachedFilesystemPathIndexReader,
};
use crate::functions::FunctionProviderHandle;
use crate::hot_state::{
    HotStateExactBatchRequest, HotStateReader, HotStateScanRequest, MaterializedHotStateBatch,
    MaterializedHotStateExactBatch,
};
use crate::plugin::runtime::PluginRuntimeHost;
use crate::plugin::runtime::UnsupportedWasmRuntime;
use crate::storage_adapter::StorageAdapterRead;
use crate::transaction_types::{
    CertifiedParameterInsertBatch, CertifiedParameterReplacementBatch, RawWriteBatch,
    TransactionWrite, TransactionWriteMode, TransactionWriteOutcome, TypedMutationJournalBatch,
};

use super::{PublicCatalog, SessionFileViews};

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum DiffCommand {
    Revert,
    Apply,
    CreateCheckpoint,
}

/// Relation-row identity selected by a public diff command.
#[derive(Clone, Debug, Eq, PartialEq)]
pub(crate) struct DiffCommandSelection {
    pub(crate) relation: String,
    pub(crate) row_pk: crate::row_pk::RowPk,
    pub(crate) source_commits: Option<(String, String)>,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub(crate) struct DiffCommandOutcome {
    pub(crate) rows_affected: u64,
    pub(crate) commit_id: Option<String>,
    pub(crate) parent_commit_id: Option<String>,
}

pub(crate) type SqlChangelogQuerySource<S> = ChangelogQuerySource<S>;
#[derive(Clone)]
pub(crate) struct ChangelogQuerySource<S> {
    pub(crate) store: S,
}

/// Read-only context used while executing one SQL statement.
///
/// Session and transaction orchestration stay above `sql2`. They provide the
/// execution-scoped committed read context for each call.
///
/// This trait is for read SQL session construction. Write SQL should use
/// `SqlWriteExecutionContext` so transaction-scoped reads and staging stay in
/// the transaction capability instead of flowing through committed read
/// sources.
#[async_trait]
pub(crate) trait SqlExecutionContext: Sync {
    type ReadStore: StorageAdapterRead + Clone + Send + Sync + 'static;

    fn active_branch_id(&self) -> &str;
    fn datafusion_session(&self) -> datafusion::prelude::SessionContext {
        super::session::new_sql_session_context()
    }
    fn datafusion_read_session(&self) -> super::planning_cache::PooledReadSession {
        super::planning_cache::PooledReadSession::standalone(self.datafusion_session())
    }
    async fn sql_planning_environment(
        &self,
    ) -> Result<
        Option<(
            Arc<super::SqlPlanningCache<crate::catalog::CatalogFingerprint>>,
            crate::catalog::CatalogFingerprint,
        )>,
        LixError,
    > {
        Ok(None)
    }
    fn active_account_id(&self) -> &str {
        crate::ANONYMOUS_ACCOUNT_ID
    }
    fn read_interest_registry(&self) -> Option<Arc<crate::hot_state::ReadInterestRegistry>> {
        None
    }
    fn hot_state(&self) -> Arc<dyn HotStateReader>;
    /// Supplies the committed tracked-head row snapshot capability when the
    /// read context can prove it is scoped to one immutable storage snapshot.
    /// Generic and transaction contexts intentionally retain the default
    /// materialized-row path.
    fn row_snapshot_reader(&self) -> Option<Arc<dyn super::RowSnapshotReader>> {
        None
    }
    fn filesystem_path_index(&self) -> Arc<dyn FilesystemPathIndexReader> {
        Arc::new(UncachedFilesystemPathIndexReader::new(self.hot_state()))
    }
    fn functions(&self) -> FunctionProviderHandle;
    fn changelog_query_source(&self) -> SqlChangelogQuerySource<Self::ReadStore>;
    fn commit_graph(&self) -> Box<dyn CommitGraphReader>;
    fn branch_ref(&self) -> Arc<dyn BranchRefReader>;
    fn blob_reader(&self) -> Arc<dyn BlobDataReader>;
    /// Loads runtime-defined SQL row metadata when provider selection could
    /// not be satisfied entirely by compile-time system surfaces.
    async fn load_visible_schemas(&self) -> Result<Vec<JsonValue>, LixError>;

    /// Loads reusable public-surface metadata for this read snapshot.
    ///
    /// The default keeps lightweight test/read contexts simple. Session
    /// contexts override it with their revision-keyed catalog cache; providers
    /// themselves remain scoped to the current storage snapshot.
    async fn public_catalog(&self) -> Result<Arc<PublicCatalog>, LixError> {
        Ok(Arc::new(PublicCatalog::from_visible_schemas(
            &self.load_visible_schemas().await?,
        )?))
    }

    fn plugin_host(&self) -> PluginRuntimeHost {
        PluginRuntimeHost::new(Arc::new(UnsupportedWasmRuntime))
    }

    fn session_file_views(&self) -> Option<SessionFileViews> {
        None
    }
}

/// Write-capable SQL runtime boundary.
///
/// Providers that mutate engine state should target this shape instead of
/// reaching through session/storage escape hatches. The request and write
/// payloads stay in the existing engine forms so this boundary centralizes
/// authority without adding another translation layer.
#[async_trait]
pub(crate) trait SqlWriteExecutionContext: Send {
    fn is_partial_replica(&self) -> bool {
        false
    }
    fn read_interest_registry(&self) -> Option<Arc<crate::hot_state::ReadInterestRegistry>> {
        None
    }

    fn ensure_statement_allowed_after_restore(&self) -> Result<(), LixError> {
        Ok(())
    }

    fn active_branch_id(&self) -> &str;
    /// Revocation token for this context's borrow.
    ///
    /// The default hands back a token that is never retired, which is correct
    /// for contexts that are not the engine `Transaction`: they own no borrow
    /// that a `SqlWriteContext` can outlive.
    fn write_context_liveness(&self) -> WriteContextLiveness {
        WriteContextLiveness::new()
    }
    fn datafusion_session(&self) -> datafusion::prelude::SessionContext {
        super::session::new_sql_session_context()
    }
    fn active_account_id(&self) -> &str {
        crate::ANONYMOUS_ACCOUNT_ID
    }
    fn functions(&self) -> FunctionProviderHandle;
    fn current_timestamp(&mut self) -> crate::common::LixTimestamp {
        self.functions().call_timestamp()
    }
    fn list_visible_schemas(&self) -> Result<Vec<JsonValue>, LixError>;
    fn public_catalog(&self) -> Result<Arc<PublicCatalog>, LixError> {
        Ok(Arc::new(PublicCatalog::from_visible_schemas(
            &self.list_visible_schemas()?,
        )?))
    }
    fn schema_catalog_snapshot(&self) -> Option<Arc<crate::catalog::CatalogSnapshot>> {
        None
    }
    fn has_staged_schema_changes(&self) -> Result<bool, LixError> {
        Ok(false)
    }
    /// Current payload-validation authority, separate from opening SQL binding.
    async fn staged_schema_document(
        &mut self,
        _domain: &crate::domain::Domain,
        _schema_key: &str,
    ) -> Result<Option<Arc<JsonValue>>, LixError> {
        Ok(None)
    }
    fn staged_schema_plan(
        &self,
        _domain: &crate::domain::Domain,
        _schema_key: &str,
    ) -> Option<&crate::catalog::SchemaPlan> {
        None
    }
    /// Catalog visible to tracked writes in the active branch.
    ///
    /// SQL binding may also see untracked schema registrations. Certified
    /// tracked write lanes must pin their plan against this narrower catalog
    /// or defer to transaction normalization.
    fn tracked_schema_catalog_snapshot(&self) -> Option<Arc<crate::catalog::CatalogSnapshot>> {
        None
    }
    /// Whether the active plugin registry owns this schema's durable rows.
    /// Engine schemas also have compiled plans but retain generic JSON rows.
    fn plugin_owns_schema(&self, _schema_key: &str) -> bool {
        false
    }
    fn plugin_host(&self) -> PluginRuntimeHost {
        PluginRuntimeHost::new(Arc::new(UnsupportedWasmRuntime))
    }

    fn session_file_views(&self) -> Option<SessionFileViews> {
        None
    }

    async fn require_referenced_manifests(&mut self, _hashes: &[BlobId]) -> Result<(), LixError> {
        Ok(())
    }

    async fn load_bytes_many(&mut self, hashes: &[BlobId]) -> Result<BlobBytesBatch, LixError>;

    async fn scan_hot_state_batch(
        &mut self,
        request: &HotStateScanRequest,
    ) -> Result<MaterializedHotStateBatch, LixError>;

    async fn load_exact_hot_state_batch(
        &mut self,
        request: &HotStateExactBatchRequest,
    ) -> Result<MaterializedHotStateExactBatch, LixError>;

    async fn filesystem_path_index(
        &mut self,
        request: &FilesystemPathIndexRequest,
    ) -> Result<Arc<FilesystemPathIndex>, LixError> {
        let rows = self
            .scan_hot_state_batch(&request.hot_state_request())
            .await?;
        Ok(Arc::new(FilesystemPathIndex::from_live_batch(&rows)?))
    }

    async fn load_branch_head(&mut self, branch_id: &str) -> Result<Option<CommitId>, LixError>;
    async fn load_branch_working_base(
        &mut self,
        _branch_id: &str,
    ) -> Result<Option<CommitId>, LixError> {
        Ok(None)
    }

    async fn load_collection_generation(
        &mut self,
        _branch_id: &str,
        _scope: crate::collection_generation::CollectionScopeRef<'_>,
    ) -> Result<Option<crate::collection_generation::CollectionGeneration>, LixError> {
        Ok(None)
    }

    async fn load_exact_collection_live_count(
        &mut self,
        _branch_id: &str,
        _scope: crate::collection_generation::CollectionScopeRef<'_>,
    ) -> Result<Option<u64>, LixError> {
        Ok(None)
    }

    fn has_staged_collection_rows(
        &self,
        _branch_id: &str,
        _scope: crate::collection_generation::CollectionScopeRef<'_>,
    ) -> Result<bool, LixError> {
        Ok(false)
    }

    async fn stage_write(
        &mut self,
        write: TransactionWrite,
    ) -> Result<TransactionWriteOutcome, LixError>;

    async fn stage_parameter_batch_insert(
        &mut self,
        rows: RawWriteBatch,
    ) -> Result<TransactionWriteOutcome, LixError> {
        self.stage_write(TransactionWrite::Rows {
            mode: TransactionWriteMode::Insert,
            rows,
        })
        .await
    }

    async fn stage_certified_parameter_batch_insert(
        &mut self,
        rows: CertifiedParameterInsertBatch,
    ) -> Result<TransactionWriteOutcome, LixError> {
        self.stage_parameter_batch_insert(rows.into_raw()?).await
    }

    async fn stage_parameter_batch_replace(
        &mut self,
        rows: RawWriteBatch,
    ) -> Result<TransactionWriteOutcome, LixError> {
        self.stage_write(TransactionWrite::Rows {
            mode: TransactionWriteMode::Replace,
            rows,
        })
        .await
    }

    async fn stage_certified_parameter_batch_replace(
        &mut self,
        rows: CertifiedParameterReplacementBatch,
    ) -> Result<TransactionWriteOutcome, LixError> {
        self.stage_parameter_batch_replace(rows.into_raw()?).await
    }

    async fn stage_typed_mutation_journal_replace(
        &mut self,
        rows: TypedMutationJournalBatch,
    ) -> Result<TransactionWriteOutcome, LixError>;

    async fn can_stage_typed_mutation_journal_replace(
        &mut self,
        schema_key: &str,
        live_count: u64,
        ordered_identity_digest: [u8; 32],
    ) -> Result<bool, LixError>;

    async fn execute_diff_command(
        &mut self,
        _command: DiffCommand,
        _selections: Vec<DiffCommandSelection>,
    ) -> Result<DiffCommandOutcome, LixError> {
        Err(LixError::new(
            LixError::CODE_UNSUPPORTED_SQL,
            "diff commands are not supported by this write context",
        ))
    }

    async fn restore_active_branch(&mut self, _commit_id: String) -> Result<(), LixError> {
        Err(LixError::new(
            LixError::CODE_UNSUPPORTED_SQL,
            "lix_restore is not supported by this write context",
        ))
    }

    fn staged_commit_id(&self, _branch_id: &str) -> Result<Option<String>, LixError> {
        Ok(None)
    }
}

#[derive(Clone)]
pub(crate) struct SqlWriteContext {
    ptr: Arc<SqlWriteContextPtr>,
    gate: Arc<Mutex<()>>,
    liveness: WriteContextLiveness,
    shared: Arc<SqlWriteContextShared>,
    explicit_insert_columns: Option<Arc<BTreeSet<String>>>,
    write_targets: Option<Arc<super::providers::WriteTargetRegistry>>,
}

struct SqlWriteContextPtr(NonNull<dyn SqlWriteExecutionContext>);

/// Revocation token for the transaction context behind `SqlWriteContextPtr`.
///
/// `SqlWriteContext::new` forges a `'static` pointer out of a `&mut` borrow, and
/// the borrow's real lifetime cannot be recovered by the type system (see the
/// note on `SqlWriteContextPtr` below). This token is the runtime substitute:
/// the borrowed `Transaction` flips it on teardown, and every gated
/// dereference checks it first, so a pointer that outlives its pointee produces
/// a deterministic `LixError` instead of undefined behaviour.
///
/// Measured before this existed: in one `cargo test -p lix -p lix_e2e` run,
/// **249 distinct `Transaction` objects were destroyed while a
/// `SqlWriteContext` still pointed at them** (116 with one live pointer, 49
/// with two, one with 26). None of them was dereferenced afterwards, so this is
/// hardening rather than a bug fix — but the safety of the whole construction
/// rested on drop ordering that nothing enforced or tested.
#[derive(Clone, Debug)]
pub(crate) struct WriteContextLiveness(Arc<std::sync::atomic::AtomicBool>);

impl WriteContextLiveness {
    pub(crate) fn new() -> Self {
        Self(Arc::new(std::sync::atomic::AtomicBool::new(true)))
    }

    /// Called when the borrowed context is torn down. Idempotent.
    pub(crate) fn retire(&self) {
        self.0.store(false, std::sync::atomic::Ordering::Release);
    }

    pub(crate) fn is_live(&self) -> bool {
        self.0.load(std::sync::atomic::Ordering::Acquire)
    }
}

impl Default for WriteContextLiveness {
    fn default() -> Self {
        Self::new()
    }
}

/// Values captured from the write execution context at construction time.
///
/// These were previously read back through `SqlWriteContextPtr` on every call,
/// producing a `&dyn` into the transaction context with no synchronization
/// while the gated methods could concurrently hold a reconstituted `&mut` to
/// the same object. Every one of them is a cheap getter returning owned,
/// `Arc`'d, or cloned data that is stable for this context's lifetime, so
/// capturing them here removes that shared borrow outright rather than
/// serializing it. Nothing below reads through the raw pointer.
struct SqlWriteContextShared {
    read_interest_registry: Option<Arc<crate::hot_state::ReadInterestRegistry>>,
    is_partial_replica: bool,
    functions: FunctionProviderHandle,
    /// Stored as the `Result` it was: the underlying catalog is memoized on a
    /// fingerprint that is fixed for the context's lifetime
    /// (`sql_schema_snapshot` is assigned once at construction and never
    /// reassigned), so a captured value cannot go stale.
    public_catalog: Result<Arc<PublicCatalog>, LixError>,
    active_branch_id: String,
    active_account_id: String,
    plugin_host: PluginRuntimeHost,
    /// A shared `Arc<Mutex<..>>` handle, not a snapshot: mutations made through
    /// a captured clone are observed by every other holder.
    session_file_views: Option<SessionFileViews>,
}

// DataFusion stores providers as owned Send + Sync trait objects. This context
// is only constructed for one write execution.
//
// SAFETY SCOPE: the pointer is reached only by the gate-serialized methods that
// reconstitute `&mut`. The shared accessors read `SqlWriteContextShared` and
// never touch it, so no `&dyn` into the transaction context can be alive while
// one of those `&mut` borrows is held.
//
// LIFETIME: this context CAN outlive the borrowed transaction context. The
// comment here used to claim it never did; that was false, and
// `WriteContextLiveness` above carries the measurement. Every gated
// dereference is therefore guarded at runtime.
//
// WHY THE LIFETIME IS FORGED, AND WHY IT CANNOT BE FIXED WITH A LIFETIME
// PARAMETER. A `SqlWriteContext<'ctx>` cannot reach DataFusion at all:
// `SessionContext::register_table` takes `Arc<dyn TableProvider>`, which is
// `Arc<dyn TableProvider + 'static>`, and `TableProvider::as_any` returns
// `&dyn Any` where `trait Any: 'static`. `ExecutionPlan` carries the same
// `as_any` requirement. Both bounds are unconditional in the DataFusion
// provider API, so the `'static` here is forced by that boundary rather than
// chosen. That is precisely why the invariant is enforced at runtime by
// `WriteContextLiveness` instead of by a lifetime: it is not expressible.
unsafe impl Send for SqlWriteContextPtr {}
unsafe impl Sync for SqlWriteContextPtr {}

impl SqlWriteContext {
    /// Refuses a dereference of a transaction context that has been torn down.
    ///
    /// Always compiled in: the whole point is to be present in the release
    /// configuration where a dangling dereference would actually bite.
    fn ensure_context_live(&self, site: &'static str) -> Result<(), LixError> {
        if self.liveness.is_live() {
            return Ok(());
        }
        Err(
            LixError::new(
                LixError::CODE_INTERNAL_ERROR,
                "SQL write context outlived the transaction it borrows; refusing to dereference a retired context",
            )
            .with_details(serde_json::json!({
                "invariant": "SqlWriteContext must not outlive the Transaction it borrows",
                "site": site,
            })),
        )
    }

    #[cfg(test)]
    pub(crate) fn liveness_for_test(&self) -> &WriteContextLiveness {
        &self.liveness
    }
}

impl SqlWriteContext {
    pub(crate) fn new(ctx: &mut dyn SqlWriteExecutionContext) -> Self {
        // Capture the shared surface while the `&mut` borrow is still held
        // legitimately, so no later call has to forge one.
        let shared = Arc::new(SqlWriteContextShared {
            read_interest_registry: ctx.read_interest_registry(),
            is_partial_replica: ctx.is_partial_replica(),
            functions: ctx.functions(),
            public_catalog: ctx.public_catalog(),
            active_branch_id: ctx.active_branch_id().to_string(),
            active_account_id: ctx.active_account_id().to_string(),
            plugin_host: ctx.plugin_host(),
            session_file_views: ctx.session_file_views(),
        });
        let liveness = ctx.write_context_liveness();
        let ptr = NonNull::from(ctx);
        let ptr = unsafe {
            std::mem::transmute::<
                NonNull<dyn SqlWriteExecutionContext + '_>,
                NonNull<dyn SqlWriteExecutionContext + 'static>,
            >(ptr)
        };
        Self {
            ptr: Arc::new(SqlWriteContextPtr(ptr)),
            gate: Arc::new(Mutex::new(())),
            liveness,
            shared,
            explicit_insert_columns: None,
            write_targets: Some(Arc::new(super::providers::WriteTargetRegistry::default())),
        }
    }

    pub(crate) fn with_explicit_insert_columns(
        mut self,
        columns: Option<BTreeSet<String>>,
    ) -> Self {
        self.explicit_insert_columns = columns.map(Arc::new);
        self
    }

    pub(crate) fn explicit_insert_columns(&self) -> Option<&BTreeSet<String>> {
        self.explicit_insert_columns.as_deref()
    }

    pub(crate) fn write_targets(
        &self,
    ) -> Result<Arc<super::providers::WriteTargetRegistry>, LixError> {
        self.write_targets.clone().ok_or_else(|| {
            LixError::unknown("physical SQL write target cannot own a write-target registry")
        })
    }

    pub(crate) fn into_physical_target(mut self) -> Self {
        self.write_targets = None;
        self
    }

    pub(crate) fn functions(&self) -> FunctionProviderHandle {
        self.shared.functions.clone()
    }

    pub(crate) fn blob_reader(&self) -> Arc<dyn BlobDataReader> {
        Arc::new(WriteContextBlobDataReader::new(self.clone()))
    }

    pub(crate) fn public_catalog(&self) -> Result<Arc<PublicCatalog>, LixError> {
        self.shared.public_catalog.clone()
    }

    pub(crate) fn active_branch_id(&self) -> String {
        self.shared.active_branch_id.clone()
    }

    pub(crate) fn active_account_id(&self) -> String {
        self.shared.active_account_id.clone()
    }

    pub(crate) fn plugin_host(&self) -> PluginRuntimeHost {
        self.shared.plugin_host.clone()
    }

    pub(crate) fn session_file_views(&self) -> Option<SessionFileViews> {
        self.shared.session_file_views.clone()
    }

    pub(crate) async fn staged_schema_document(
        &self,
        domain: &crate::domain::Domain,
        schema_key: &str,
    ) -> Result<Option<Arc<JsonValue>>, LixError> {
        let _guard = self.gate.lock().await;
        self.ensure_context_live("staged_schema_document")?;
        unsafe {
            self.ptr
                .0
                .as_ptr()
                .as_mut()
                .unwrap()
                .staged_schema_document(domain, schema_key)
                .await
        }
    }

    pub(crate) async fn scan_hot_state_batch(
        &self,
        request: &HotStateScanRequest,
    ) -> Result<MaterializedHotStateBatch, LixError> {
        let _guard = self.gate.lock().await;
        self.ensure_context_live("scan_hot_state_batch")?;
        unsafe {
            self.ptr
                .0
                .as_ptr()
                .as_mut()
                .unwrap()
                .scan_hot_state_batch(request)
                .await
        }
    }

    pub(crate) async fn load_exact_hot_state_batch(
        &self,
        request: &HotStateExactBatchRequest,
    ) -> Result<MaterializedHotStateExactBatch, LixError> {
        let _guard = self.gate.lock().await;
        self.ensure_context_live("load_exact_hot_state_batch")?;
        unsafe {
            self.ptr
                .0
                .as_ptr()
                .as_mut()
                .unwrap()
                .load_exact_hot_state_batch(request)
                .await
        }
    }

    pub(crate) async fn require_referenced_manifests(
        &self,
        hashes: &[BlobId],
    ) -> Result<(), LixError> {
        let _guard = self.gate.lock().await;
        self.ensure_context_live("require_referenced_manifests")?;
        unsafe {
            self.ptr
                .0
                .as_ptr()
                .as_mut()
                .unwrap()
                .require_referenced_manifests(hashes)
                .await
        }
    }

    pub(crate) async fn load_bytes_many(
        &self,
        hashes: &[BlobId],
    ) -> Result<BlobBytesBatch, LixError> {
        let _guard = self.gate.lock().await;
        self.ensure_context_live("load_bytes_many")?;
        unsafe {
            self.ptr
                .0
                .as_ptr()
                .as_mut()
                .unwrap()
                .load_bytes_many(hashes)
                .await
        }
    }

    pub(crate) async fn load_branch_head(
        &self,
        branch_id: &str,
    ) -> Result<Option<CommitId>, LixError> {
        let _guard = self.gate.lock().await;
        self.ensure_context_live("load_branch_head")?;
        unsafe {
            self.ptr
                .0
                .as_ptr()
                .as_mut()
                .unwrap()
                .load_branch_head(branch_id)
                .await
        }
    }

    pub(crate) async fn load_branch_working_base(
        &self,
        branch_id: &str,
    ) -> Result<Option<CommitId>, LixError> {
        let _guard = self.gate.lock().await;
        self.ensure_context_live("load_branch_working_base")?;
        unsafe {
            self.ptr
                .0
                .as_ptr()
                .as_mut()
                .unwrap()
                .load_branch_working_base(branch_id)
                .await
        }
    }

    pub(crate) async fn filesystem_path_index(
        &self,
        request: &FilesystemPathIndexRequest,
    ) -> Result<Arc<FilesystemPathIndex>, LixError> {
        let _guard = self.gate.lock().await;
        self.ensure_context_live("filesystem_path_index")?;
        unsafe {
            self.ptr
                .0
                .as_ptr()
                .as_mut()
                .unwrap()
                .filesystem_path_index(request)
                .await
        }
    }

    pub(crate) async fn stage_write(
        &self,
        write: TransactionWrite,
    ) -> Result<TransactionWriteOutcome, LixError> {
        let _guard = self.gate.lock().await;
        self.ensure_context_live("stage_write")?;
        unsafe {
            self.ptr
                .0
                .as_ptr()
                .as_mut()
                .unwrap()
                .stage_write(write)
                .await
        }
    }

    pub(crate) async fn execute_diff_command(
        &self,
        command: DiffCommand,
        selections: Vec<DiffCommandSelection>,
    ) -> Result<DiffCommandOutcome, LixError> {
        let _guard = self.gate.lock().await;
        self.ensure_context_live("execute_diff_command")?;
        unsafe {
            self.ptr
                .0
                .as_ptr()
                .as_mut()
                .unwrap()
                .execute_diff_command(command, selections)
                .await
        }
    }
}

pub(crate) struct WriteContextBlobDataReader {
    ctx: SqlWriteContext,
}

impl WriteContextBlobDataReader {
    pub(crate) fn new(ctx: SqlWriteContext) -> Self {
        Self { ctx }
    }
}

#[async_trait]
impl BlobDataReader for WriteContextBlobDataReader {
    fn requires_referenced_content_preparation(&self) -> bool {
        self.ctx.shared.is_partial_replica
    }

    async fn require_referenced_manifests(&self, hashes: &[BlobId]) -> Result<(), LixError> {
        self.ctx.require_referenced_manifests(hashes).await
    }
    async fn load_bytes_many(&self, hashes: &[BlobId]) -> Result<BlobBytesBatch, LixError> {
        self.ctx.load_bytes_many(hashes).await
    }
}

#[derive(Clone)]
pub(crate) enum WriteAccess {
    ReadOnly,
    Write { ctx: SqlWriteContext },
}

impl WriteAccess {
    pub(crate) fn read_only() -> Self {
        Self::ReadOnly
    }

    pub(crate) fn write(ctx: SqlWriteContext) -> Self {
        Self::Write { ctx }
    }

    pub(crate) fn into_write_context(self) -> Option<SqlWriteContext> {
        match self {
            Self::ReadOnly => None,
            Self::Write { ctx } => Some(ctx),
        }
    }
}

pub(crate) struct WriteContextHotStateReader {
    ctx: SqlWriteContext,
}

impl WriteContextHotStateReader {
    pub(crate) fn new(ctx: SqlWriteContext) -> Self {
        Self { ctx }
    }
}

#[async_trait]
impl HotStateReader for WriteContextHotStateReader {
    fn is_partial_replica(&self) -> bool {
        self.ctx.shared.is_partial_replica
    }
    fn read_interest_registry(&self) -> Option<Arc<crate::hot_state::ReadInterestRegistry>> {
        self.ctx.shared.read_interest_registry.clone()
    }

    async fn scan_batch(
        &self,
        request: &HotStateScanRequest,
    ) -> Result<MaterializedHotStateBatch, LixError> {
        self.ctx.scan_hot_state_batch(request).await
    }

    async fn load_exact_batch(
        &self,
        request: &HotStateExactBatchRequest,
    ) -> Result<MaterializedHotStateExactBatch, LixError> {
        self.ctx.load_exact_hot_state_batch(request).await
    }
}

#[async_trait]
impl FilesystemPathIndexReader for WriteContextHotStateReader {
    async fn path_index(
        &self,
        request: &FilesystemPathIndexRequest,
    ) -> Result<Arc<FilesystemPathIndex>, LixError> {
        self.ctx.filesystem_path_index(request).await
    }
}

pub(crate) struct WriteContextBranchRefReader {
    ctx: SqlWriteContext,
}

impl WriteContextBranchRefReader {
    pub(crate) fn new(ctx: SqlWriteContext) -> Self {
        Self { ctx }
    }
}

#[async_trait]
impl BranchRefReader for WriteContextBranchRefReader {
    async fn load_head(&self, branch_id: &str) -> Result<Option<BranchHead>, LixError> {
        let working_base_commit_id = self.ctx.load_branch_working_base(branch_id).await?;
        Ok(self
            .ctx
            .load_branch_head(branch_id)
            .await?
            .map(|commit_id| BranchHead {
                branch_id: branch_id.to_string(),
                commit_id,
                working_base_commit_id,
            }))
    }

    async fn scan_heads(&self) -> Result<Vec<BranchHead>, LixError> {
        Err(LixError::new(
            "LIX_ERROR_UNKNOWN",
            "scan_heads is not available through sql2 write context",
        ))
    }
}