yantrikdb 0.7.1

use rusqlite::params;

use crate::error::Result;
use crate::serde_helpers::serialize_f32;
use crate::types::*;

use super::{now, embedding_hash, YantrikDB};

impl YantrikDB {
    /// Store a new memory and return its RID.
    #[tracing::instrument(skip(self, metadata, embedding), fields(memory_type, namespace))]
    pub fn record(
        &self,
        text: &str,
        memory_type: &str,
        importance: f64,
        valence: f64,
        half_life: f64,
        metadata: &serde_json::Value,
        embedding: &[f32],
        namespace: &str,
        certainty: f64,
        domain: &str,
        source: &str,
        emotional_state: Option<&str>,
    ) -> Result<String> {
        let rid = crate::id::new_id();
        let ts = now();
        let emb_blob = serialize_f32(embedding);
        let meta_str = serde_json::to_string(metadata)?;

        // Encrypt fields if encryption is enabled
        let stored_text = self.encrypt_text(text)?;
        let stored_meta = self.encrypt_text(&meta_str)?;
        let stored_emb = self.encrypt_embedding(&emb_blob)?;

        // Read active session for this namespace into a local before acquiring conn
        let session_id = self.active_sessions.read().get(namespace).cloned();

        // Acquire conn, do all SQL, then drop before other locks
        {
            let conn = self.conn();
            conn.execute(
                "INSERT INTO memories \
                 (rid, type, text, embedding, created_at, updated_at, importance, \
                  half_life, last_access, valence, metadata, namespace, \
                  certainty, domain, source, emotional_state) \
                 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16)",
                params![rid, memory_type, stored_text, stored_emb, ts, ts, importance, half_life, ts, valence, stored_meta, namespace,
                        certainty, domain, source, emotional_state],
            )?;

            // Auto-link to active session for this namespace
            if let Some(session_id) = &session_id {
                conn.execute(
                    "UPDATE memories SET session_id = ?1 WHERE rid = ?2",
                    params![session_id, rid],
                )?;
                conn.execute(
                    "UPDATE sessions SET memory_count = memory_count + 1 WHERE session_id = ?1",
                    params![session_id],
                )?;
            }
        }
        // conn dropped here

        // Insert into vector index (lock ordering: conn already dropped)
        let seq = self.vec_seq.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
        self.vec_index.append(rid.clone(), embedding.to_vec(), seq)?;
        self.bump_visible_seq(namespace, seq);

        // Insert into scoring cache (conn and vec_index dropped)
        self.cache_insert(rid.clone(), ScoringRow {
            created_at: ts,
            importance,
            half_life,
            last_access: ts,
            access_count: 0,
            valence,
            consolidation_status: "active".to_string(),
            memory_type: memory_type.to_string(),
            namespace: namespace.to_string(),
            certainty,
            domain: domain.to_string(),
            source: source.to_string(),
            emotional_state: emotional_state.map(|s| s.to_string()),
        });

        // Log the user-facing "record" op FIRST so external consumers
        // (replication extract_ops_since, oplog inspectors) see records
        // in their natural causal order: the record came before any
        // post-record materialization queued in its wake.
        let emb_hash = embedding_hash(embedding);
        self.log_op(
            "record",
            Some(&rid),
            &serde_json::json!({
                "rid": rid,
                "type": memory_type,
                "text": text,
                "importance": importance,
                "valence": valence,
                "half_life": half_life,
                "metadata": metadata,
                "created_at": ts,
                "updated_at": ts,
                "namespace": namespace,
                "certainty": certainty,
                "domain": domain,
                "source": source,
                "emotional_state": emotional_state,
            }),
            Some(&emb_hash),
        )?;

        // **Phase 4.3 Commit B (saga task 3, 2026-05-08).** The
        // unbounded entity / memory_entities / claims loops that used
        // to live here are now enqueued for the materializer thread to
        // run off the request path. See docs/phase_4_3_design.md for
        // the contract change (synchronous read-after-write of
        // entity-graph queries shifts from immediate to ms-scale; the
        // delta-recall path is unaffected since DeltaIndex.append
        // happened above on the foreground thread).
        {
            let post_payload = serde_json::json!({
                "rid": rid,
                "text": stored_text,
                "namespace": namespace,
                "ts_secs": ts,
                "domain": domain,
                "source": source,
            });
            self.log_op_pending(
                crate::engine::op_types::OP_MATERIALIZE_RECORD_POST,
                Some(&rid),
                &post_payload,
                None,
                None,
            )?;
        }

        Ok(rid)
    }

    /// Record multiple memories in a single transaction.
    /// Uses SAVEPOINT for atomicity while keeping `&self` (no `&mut self`).
    #[tracing::instrument(skip(self, inputs), fields(batch_size = inputs.len()))]
    pub fn record_batch(&self, inputs: &[RecordInput]) -> Result<Vec<String>> {
        if inputs.is_empty() {
            return Ok(vec![]);
        }

        // Clone active sessions map before acquiring conn
        let sessions = self.active_sessions.read().clone();

        // Precompute entity candidates per memory before touching conn/graph_index.
        // Two sources:
        //   (a) heuristic extraction from text (capitalized proper-nouns)
        //   (b) match against already-known entities in graph_index
        let known_entities = self.graph_index.read().all_entity_names();
        let per_memory_linkage: Vec<(Vec<String>, std::collections::HashSet<String>)> = inputs
            .iter()
            .map(|input| {
                let text_tokens = crate::graph::tokenize(&input.text);
                let heuristic = crate::graph::extract_heuristic_entities(&input.text);
                let mut candidates: std::collections::HashSet<String> =
                    heuristic.iter().cloned().collect();
                for known in &known_entities {
                    if crate::graph::entity_matches_text(known, &text_tokens) {
                        candidates.insert(known.clone());
                    }
                }
                (heuristic, candidates)
            })
            .collect();

        let mut rids = Vec::with_capacity(inputs.len());

        // Lock conn once for the entire batch SQL work
        {
            let conn = self.conn();
            conn.execute_batch("SAVEPOINT batch_record")?;

            for input in inputs {
                let rid = crate::id::new_id();
                let ts = now();
                let emb_blob = serialize_f32(&input.embedding);
                let meta_str = serde_json::to_string(&input.metadata)?;

                // Encrypt fields if encryption is enabled
                let stored_text = self.encrypt_text(&input.text)?;
                let stored_meta = self.encrypt_text(&meta_str)?;
                let stored_emb = self.encrypt_embedding(&emb_blob)?;

                let result = conn.execute(
                    "INSERT INTO memories \
                     (rid, type, text, embedding, created_at, updated_at, importance, \
                      half_life, last_access, valence, metadata, namespace, \
                      certainty, domain, source, emotional_state) \
                     VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16)",
                    params![rid, input.memory_type, stored_text, stored_emb, ts, ts,
                            input.importance, input.half_life, ts, input.valence, stored_meta,
                            input.namespace, input.certainty, input.domain, input.source,
                            input.emotional_state],
                );

                if let Err(e) = result {
                    conn.execute_batch("ROLLBACK TO batch_record")?;
                    return Err(e.into());
                }

                rids.push(rid);
            }

            // Auto-link batch to active sessions
            for (rid, input) in rids.iter().zip(inputs.iter()) {
                if let Some(session_id) = sessions.get(&input.namespace) {
                    conn.execute(
                        "UPDATE memories SET session_id = ?1 WHERE rid = ?2",
                        params![session_id, rid],
                    )?;
                    conn.execute(
                        "UPDATE sessions SET memory_count = memory_count + 1 WHERE session_id = ?1",
                        params![session_id],
                    )?;
                }
            }

            // Persist entity linkage (SQL side). graph_index in-memory update
            // happens after conn is dropped to avoid holding two write locks.
            let batch_ts = now();
            for (rid, (heuristic, candidates)) in rids.iter().zip(per_memory_linkage.iter()) {
                for entity in heuristic {
                    let entity_type = crate::graph::classify_entity_type(entity);
                    conn.execute(
                        "INSERT INTO entities (name, entity_type, first_seen, last_seen, mention_count) \
                         VALUES (?1, ?2, ?3, ?3, 1) \
                         ON CONFLICT(name) DO UPDATE SET \
                            last_seen = ?3, \
                            mention_count = mention_count + 1, \
                            entity_type = CASE \
                                WHEN entity_type = 'unknown' AND ?2 != 'unknown' THEN ?2 \
                                ELSE entity_type END",
                        params![entity, entity_type, batch_ts],
                    )?;
                }
                for entity in candidates {
                    conn.execute(
                        "INSERT OR IGNORE INTO memory_entities (memory_rid, entity_name) VALUES (?1, ?2)",
                        params![rid, entity],
                    )?;
                }
            }

            conn.execute_batch("RELEASE batch_record")?;
        }
        // conn dropped; now update graph_index in-memory.
        {
            let mut gi = self.graph_index.write();
            for (rid, (_, candidates)) in rids.iter().zip(per_memory_linkage.iter()) {
                for entity in candidates {
                    let entity_type = crate::graph::classify_entity_type(entity);
                    gi.add_entity(entity, entity_type);
                    gi.link_memory(rid, entity);
                }
            }
        }

        // RFC 006 Phase 0: emit one audit event per memory in the batch.
        for (rid, (input, (heuristic_entities, candidates))) in
            rids.iter().zip(inputs.iter().zip(per_memory_linkage.iter()))
        {
            let heuristic_vec: Vec<String> = heuristic_entities.iter().cloned().collect();
            let features = crate::graph::analyze_text_features(&input.text, &heuristic_vec);
            tracing::info!(
                target: "yantrikdb::audit::extraction",
                namespace = %input.namespace,
                memory_rid = %rid,
                domain = %input.domain,
                source = %input.source,
                extractor_version = "heuristic_v1",
                batch = true,
                char_length = features.char_length,
                sentence_count = features.sentence_count,
                entity_count = features.entity_count,
                entities_matched_in_graph = candidates.len().saturating_sub(heuristic_entities.len()),
                negation_cue_count = features.negation_cue_count,
                temporal_cue_count = features.temporal_cue_count,
                modality_cue_count = features.modality_cue_count,
                has_compound_markers = features.has_compound_markers,
                likely_assertion = features.likely_assertion,
                "extraction audit"
            );
        }

        // Append to vec_index (DeltaIndex) after SQL commit
        for (rid, input) in rids.iter().zip(inputs.iter()) {
            let seq = self.vec_seq.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
            self.vec_index.append(rid.clone(), input.embedding.clone(), seq)?;
            self.bump_visible_seq(&input.namespace, seq);
        }
        // vec_index dropped, now scoring_cache
        {
            let mut cache = self.scoring_cache.write();
            for (rid, input) in rids.iter().zip(inputs.iter()) {
                let ts = now();
                cache.insert(rid.clone(), ScoringRow {
                    created_at: ts,
                    importance: input.importance,
                    half_life: input.half_life,
                    last_access: ts,
                    access_count: 0,
                    valence: input.valence,
                    consolidation_status: "active".to_string(),
                    memory_type: input.memory_type.clone(),
                    namespace: input.namespace.clone(),
                    certainty: input.certainty,
                    domain: input.domain.clone(),
                    source: input.source.clone(),
                    emotional_state: input.emotional_state.clone(),
                });
            }
        }

        // Log a single batch op (log_op locks conn internally)
        self.log_op(
            "record_batch",
            None,
            &serde_json::json!({
                "count": rids.len(),
                "rids": rids,
            }),
            None,
        )?;

        Ok(rids)
    }

    /// **Issue #9 — deterministic mutation primitive for cluster replication.**
    ///
    /// Sibling of `record()` that takes a caller-assigned rid + caller-supplied
    /// embedding + materialized extracted_entities + caller-supplied
    /// timestamp + embedding_model. Engine does NOT call its own embedder
    /// or NER. Used by yantrikdb-server's cluster-mode applier so
    /// replicated writes are byte-deterministic across leader + followers.
    ///
    /// # Contract
    ///
    /// - **Idempotent on rid**: a second call with the same rid + identical
    ///   other fields succeeds without error and produces identical engine
    ///   state (INSERT OR IGNORE on memories, INSERT OR IGNORE on entities,
    ///   INSERT OR IGNORE on memory_entities, DeltaIndex.append idempotent
    ///   on rid+seq).
    /// - **Caller supplies the embedding.** Engine validates dim and rejects
    ///   `Error::EmbeddingDimensionMismatch` on mismatch — diverged dim is
    ///   undetectable until a query notices, so we fail loudly.
    /// - **Caller supplies created_at_unix_micros.** Materialized into both
    ///   `created_at REAL` (for back-compat scoring) and the v25
    ///   `created_at_unix_micros INTEGER` column. No engine-side `now()`
    ///   call on this path — leader stamps once, followers replay verbatim.
    /// - **Caller supplies extracted_entities.** Engine writes entity_edges
    ///   accordingly. Empty slice = no edges; engine does NOT fall back to
    ///   its own NER. (Heuristic NER lives in `crate::knowledge::graph` and
    ///   is callable directly by the leader if needed — see issue #9 thread.)
    /// - **Caller supplies embedding_model.** Stored on the row as the
    ///   engine-deterministic-surface version pin. RFC 013 may swap the
    ///   field type later behind the same column name.
    /// - **Caller-supplied `seq`** (cluster mode): when `Some(n)`, the
    ///   engine uses `n` as the delta-entry seq and the visible_seq bump
    ///   value, and ratchets `vec_seq` up to at least `n`. Per design
    ///   lock 2026-05-07, the seq IS the openraft commit-log index in
    ///   cluster mode, giving byte-deterministic per-namespace
    ///   visible_seq across leader + followers. Single-node callers pass
    ///   `None` and the engine allocates the seq itself.
    ///
    /// # Returns
    ///
    /// `Ok(())` on success or idempotent re-apply. The rid is the input,
    /// not the output — caller already owns it.
    #[allow(clippy::too_many_arguments)]
    #[tracing::instrument(skip(self, metadata, embedding, extracted_entities), fields(rid, memory_type, namespace, embedding_model))]
    pub fn record_with_rid(
        &self,
        rid: &str,
        text: &str,
        memory_type: &str,
        importance: f64,
        valence: f64,
        half_life: f64,
        metadata: &serde_json::Value,
        embedding: &[f32],
        namespace: &str,
        certainty: f64,
        domain: &str,
        source: &str,
        emotional_state: Option<&str>,
        created_at_unix_micros: i64,
        extracted_entities: &[&str],
        embedding_model: &str,
        seq: Option<u64>,
    ) -> Result<()> {
        // Determinism gate: dim must match. Diverged dim = silent corruption.
        if embedding.len() != self.embedding_dim {
            return Err(crate::error::YantrikDbError::EmbeddingDimensionMismatch {
                expected: self.embedding_dim,
                got: embedding.len(),
            });
        }

        // Caller-supplied timestamp — NEVER call now() on this path.
        let ts_secs = (created_at_unix_micros as f64) / 1_000_000.0;
        let emb_blob = serialize_f32(embedding);
        let meta_str = serde_json::to_string(metadata)?;

        // Encryption is engine-side and deterministic given the same DEK +
        // same plaintext bytes (AES-GCM is non-deterministic across IVs but
        // the encrypt-once-on-leader model means each follower receives the
        // already-encrypted bytes via the WAL replication path — Phase 4
        // wires that. For now we encrypt locally; cluster-mode follower
        // apply will skip this step in a follow-up patch.)
        let stored_text = self.encrypt_text(text)?;
        let stored_meta = self.encrypt_text(&meta_str)?;
        let stored_emb = self.encrypt_embedding(&emb_blob)?;

        let session_id = self.active_sessions.read().get(namespace).cloned();

        // Single conn block: INSERT OR IGNORE on memories (idempotent on rid),
        // session links, entity persistence. SAVEPOINT for atomicity within
        // the call.
        let was_new_row: bool = {
            let conn = self.conn();
            conn.execute_batch("SAVEPOINT record_with_rid")?;

            let result: Result<bool> = (|| {
                let inserted = conn.execute(
                    "INSERT OR IGNORE INTO memories \
                     (rid, type, text, embedding, created_at, updated_at, importance, \
                      half_life, last_access, valence, metadata, namespace, \
                      certainty, domain, source, emotional_state, \
                      created_at_unix_micros, embedding_model) \
                     VALUES (?1, ?2, ?3, ?4, ?5, ?5, ?6, ?7, ?5, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16)",
                    params![
                        rid, memory_type, stored_text, stored_emb,
                        ts_secs,
                        importance, half_life, valence, stored_meta, namespace,
                        certainty, domain, source, emotional_state,
                        created_at_unix_micros, embedding_model,
                    ],
                )?;
                let was_new_row = inserted == 1;

                if was_new_row {
                    // Auto-link only on first insert. Replay should not
                    // re-bump session memory_count.
                    if let Some(session_id) = &session_id {
                        conn.execute(
                            "UPDATE memories SET session_id = ?1 WHERE rid = ?2",
                            params![session_id, rid],
                        )?;
                        conn.execute(
                            "UPDATE sessions SET memory_count = memory_count + 1 WHERE session_id = ?1",
                            params![session_id],
                        )?;
                    }
                }

                // **Phase 4.3 Commit C (saga task 19, 2026-05-08).** The
                // entity / memory_entities INSERT loop was previously here
                // inside the SAVEPOINT, holding `db.conn().lock()` for
                // O(extracted_entities.len()) statements. Now enqueued as
                // OP_MATERIALIZE_RECORD_WITH_RID_POST after the SAVEPOINT
                // releases. See docs/phase_4_3_design.md for the contract.

                Ok(was_new_row)
            })();

            match result {
                Ok(b) => { conn.execute_batch("RELEASE record_with_rid")?; b }
                Err(e) => {
                    let _ = conn.execute_batch("ROLLBACK TO record_with_rid");
                    let _ = conn.execute_batch("RELEASE record_with_rid");
                    return Err(e);
                }
            }
        };
        // conn dropped

        // DeltaIndex append. The seq is either caller-supplied (cluster
        // mode: openraft commit-log index for byte-deterministic replay)
        // or engine-allocated (single-node). On idempotent replay the rid
        // is the same and the seq is identical (cluster) or fresh
        // (single-node retry); the compactor's highest-seq-wins rule
        // converges state identically on both paths.
        let seq = self.assign_seq(seq);
        self.vec_index.append(rid.to_string(), embedding.to_vec(), seq)?;
        self.bump_visible_seq(namespace, seq);

        // Scoring cache (engine-internal; replay safe since insert is
        // overwrite-on-rid).
        if was_new_row {
            self.cache_insert(rid.to_string(), ScoringRow {
                created_at: ts_secs,
                importance,
                half_life,
                last_access: ts_secs,
                access_count: 0,
                valence,
                consolidation_status: "active".to_string(),
                memory_type: memory_type.to_string(),
                namespace: namespace.to_string(),
                certainty,
                domain: domain.to_string(),
                source: source.to_string(),
                emotional_state: emotional_state.map(|s| s.to_string()),
            });
        }

        // Op log entry — applied=1 since leader has materialized inline.
        // Followers will receive a separate replicated entry via the
        // cluster sync path; this path never logs applied=0.
        //
        // Logged BEFORE the post-record materialization enqueue so
        // extract_ops_since reports the user-data op in causal order
        // (record_with_rid arrived, then its entity-link materialization
        // was queued).
        let emb_hash = embedding_hash(embedding);
        if was_new_row {
            self.log_op(
                "record_with_rid",
                Some(rid),
                &serde_json::json!({
                    "rid": rid,
                    "type": memory_type,
                    "text": text,
                    "importance": importance,
                    "valence": valence,
                    "half_life": half_life,
                    "metadata": metadata,
                    "created_at_unix_micros": created_at_unix_micros,
                    "namespace": namespace,
                    "certainty": certainty,
                    "domain": domain,
                    "source": source,
                    "emotional_state": emotional_state,
                    "embedding_model": embedding_model,
                    "extracted_entities": extracted_entities,
                }),
                Some(&emb_hash),
            )?;
        }

        // **Phase 4.3 Commit C (saga task 19, 2026-05-08).** Enqueue the
        // entity / memory_entities / graph_index materialization for the
        // worker thread. Skip when there are no entities to apply — the
        // dispatch arm short-circuits the same way, but skipping avoids
        // a wasteful oplog row in the common no-entity case.
        //
        // Cluster determinism: the leader and each follower will both
        // enqueue + apply this op against their local state. Convergence
        // on entities + memory_entities is guaranteed by the same
        // INSERT OR IGNORE / ON CONFLICT idempotency the inline path
        // had. The convergence *time* differs by the materializer-lag
        // window (ms-scale), but the converged final state is identical.
        if !extracted_entities.is_empty() {
            let entities_json: Vec<&str> = extracted_entities.to_vec();
            let post_payload = serde_json::json!({
                "rid": rid,
                "namespace": namespace,
                "ts_secs": ts_secs,
                "extracted_entities": entities_json,
                "was_new_row": was_new_row,
            });
            self.log_op_pending(
                crate::engine::op_types::OP_MATERIALIZE_RECORD_WITH_RID_POST,
                Some(rid),
                &post_payload,
                None,
                None,
            )?;
        }

        Ok(())
    }
}