poimen-memory/crates/mem-cli/src/dual_write_indexer.rs

//! M8.2 — Dual-write indexing pipeline
//!
//! Coordinates atomic writes to both pgvector (embedding search) and OpenSearch (lexical search).
//! Same chunk_id in both stores. If OpenSearch fails, marks `opensearch_pending=true` for eventual
//! consistency retry loop.

use anyhow::{anyhow, Result};
use sha2::{Digest, Sha256};
use sqlx::PgPool;
use uuid::Uuid;
use pgvector::Vector;
use std::sync::Arc;
use crate::opensearch_client::OpenSearchClient;
use crate::queue_adapter::QueueAdapter;

#[derive(Clone)]
pub struct DualWriteIndexer {
    pool: PgPool,
    opensearch: Option<Arc<OpenSearchClient>>,
    /// Queue adapter for concurrent dual-write processing
    /// Can be: kmsvc (production), in-memory (testing), or SQS (future)
    queue: Arc<dyn QueueAdapter>,
}

/// Input chunk for dual-write
#[derive(Debug, Clone)]
pub struct ChunkInput {
    pub content: String,
    pub source: String,
    pub project: String,
    pub level: String,  // "L0", "L1", "L2", "R"
    pub breadcrumb: Vec<String>,
}

/// Result of dual-write operation
#[derive(Debug, Clone)]
pub struct DualWriteResult {
    pub chunk_id: Uuid,
    pub chunk_hash: String,
    pub pgvector_success: bool,
    pub opensearch_success: bool,
    pub opensearch_pending: bool,  // true if OpenSearch failed
    pub error: Option<String>,
}

impl DualWriteIndexer {
    /// Create dual-write indexer with queue adapter
    pub fn new(
        pool: PgPool,
        opensearch: Option<Arc<OpenSearchClient>>,
        queue: Arc<dyn QueueAdapter>,
    ) -> Self {
        Self {
            pool,
            opensearch,
            queue,
        }
    }

    /// Queue chunk for dual-write processing
    ///
    /// Sequence:
    /// 1. Check dedup (chunk_hash exists AND indexed_in_pgvector AND indexed_in_opensearch)
    /// 2. Queue message to external queue service (kmsvc/SQS/etc)
    /// 3. Concurrent workers receive from queue and perform dual-write
    ///
    /// Returns message_id for tracking progress
    pub async fn queue_chunk(
        &self,
        chunk: &ChunkInput,
        embedding: &[f32],
    ) -> Result<String> {
        let chunk_id = Uuid::new_v4();
        let chunk_hash = self.compute_hash(&chunk.content);

        // Check deduplication
        if self.is_already_indexed(&chunk_hash, &chunk.project).await? {
            tracing::debug!("Chunk already indexed (dedup): {}", chunk_hash);
            return Ok(Uuid::nil().to_string());
        }

        // Build message attributes
        let mut attributes = std::collections::HashMap::new();
        attributes.insert("source".to_string(), chunk.source.clone());
        attributes.insert("level".to_string(), chunk.level.clone());
        attributes.insert("breadcrumb".to_string(), serde_json::to_string(&chunk.breadcrumb)?);
        attributes.insert("embedding_size".to_string(), embedding.len().to_string());

        // Build message body
        let body = serde_json::json!({
            "chunk_id": chunk_id,
            "content": chunk.content,
            "source": chunk.source,
            "level": chunk.level,
            "breadcrumb": chunk.breadcrumb,
            "embedding": embedding,
        }).to_string();

        // Queue message
        let message_id = self.queue.send_chunk(
            chunk_id,
            body,
            chunk.project.clone(),
            attributes,
        ).await?;

        tracing::info!("Chunk queued for dual-write: message_id={}, chunk_hash={}", message_id, chunk_hash);

        Ok(message_id)
    }

    /// Worker: Process queued chunk for dual-write
    ///
    /// Called by concurrent workers receiving from queue.
    /// Sequence:
    /// 1. Receive message from queue
    /// 2. Write to pgvector with embedding
    /// 3. Write to OpenSearch (fail-soft)
    /// 4. Delete from queue on success, or extend visibility on retry
    pub async fn process_queued_chunk(
        &self,
        message: &crate::queue_adapter::QueueMessage,
        embedding: &[f32],
    ) -> Result<DualWriteResult> {
        let body: serde_json::Value = serde_json::from_str(&message.body)?;
        let chunk_id = body["chunk_id"].as_str().ok_or_else(|| anyhow!("Missing chunk_id"))?
            .parse::<Uuid>()?;
        let content = body["content"].as_str().ok_or_else(|| anyhow!("Missing content"))?.to_string();
        let source = body["source"].as_str().ok_or_else(|| anyhow!("Missing source"))?.to_string();
        let project = message.project.clone();
        let level = body["level"].as_str().ok_or_else(|| anyhow!("Missing level"))?.to_string();
        let breadcrumb: Vec<String> = serde_json::from_value(body["breadcrumb"].clone())?;

        let chunk_hash = self.compute_hash(&content);

        // Write to pgvector
        let pgvector_success = self
            .write_pgvector(
                &chunk_id,
                &chunk_hash,
                &content,
                &source,
                &project,
                &level,
                &breadcrumb,
                embedding,
            )
            .await;

        if !pgvector_success.is_ok() {
            tracing::error!("pgvector write failed: {}", pgvector_success.as_ref().err().unwrap());
            // Extend visibility timeout for retry
            self.queue.change_visibility(&message.message_id, &message.receipt_handle, 300).await.ok();
            return Ok(DualWriteResult {
                chunk_id,
                chunk_hash,
                pgvector_success: false,
                opensearch_success: false,
                opensearch_pending: false,
                error: Some(format!("{:?}", pgvector_success.err())),
            });
        }

        // Write to OpenSearch (fail-soft)
        let opensearch_success = if let Some(os_client) = &self.opensearch {
            self.write_opensearch(
                os_client,
                &chunk_id,
                &content,
                &source,
                &project,
                &level,
                &breadcrumb,
            )
            .await
        } else {
            Ok(())
        };

        let opensearch_pending = opensearch_success.is_err();

        if opensearch_pending {
            tracing::warn!(
                "OpenSearch write failed, marking for retry: {}",
                opensearch_success.as_ref().err().unwrap()
            );
            self.queue.change_visibility(&message.message_id, &message.receipt_handle, 300).await.ok();
        } else {
            // Success: delete from queue
            self.queue.delete_chunk(&message.message_id, &message.receipt_handle).await.ok();
        }

        Ok(DualWriteResult {
            chunk_id,
            chunk_hash,
            pgvector_success: pgvector_success.is_ok(),
            opensearch_success: opensearch_success.is_ok(),
            opensearch_pending,
            error: if opensearch_pending {
                Some(format!("{:?}", opensearch_success.err()))
            } else {
                None
            },
        })
    }

    /// Legacy: Direct dual-write (for backward compatibility)
    ///
    /// If queue adapter is not available, use this for synchronous processing.
    pub async fn dual_write(
        &self,
        chunk: &ChunkInput,
        embedding: &[f32],
    ) -> Result<DualWriteResult> {
        let chunk_id = Uuid::new_v4();
        let chunk_hash = self.compute_hash(&chunk.content);

        // Step 1: Check deduplication
        if self.is_already_indexed(&chunk_hash, &chunk.project).await? {
            tracing::debug!("Chunk already indexed (dedup): {}", chunk_hash);
            return Ok(DualWriteResult {
                chunk_id: Uuid::nil(),  // Placeholder
                chunk_hash,
                pgvector_success: true,
                opensearch_success: true,
                opensearch_pending: false,
                error: Some("already_indexed".to_string()),
            });
        }

        // Step 2: Write to pgvector
        let pgvector_success = self.write_pgvector(
            &chunk_id,
            &chunk_hash,
            &chunk.content,
            &chunk.source,
            &chunk.project,
            &chunk.level,
            &chunk.breadcrumb,
            embedding,
        )
        .await;

        if !pgvector_success.is_ok() {
            tracing::error!("pgvector write failed: {}", pgvector_success.as_ref().err().unwrap());
            return Ok(DualWriteResult {
                chunk_id,
                chunk_hash,
                pgvector_success: false,
                opensearch_success: false,
                opensearch_pending: false,
                error: Some(format!("{:?}", pgvector_success.err())),
            });
        }

        // Step 3: Write to OpenSearch (fail-soft)
        let opensearch_success = if let Some(os_client) = &self.opensearch {
            self.write_opensearch(
                os_client,
                &chunk_id,
                &chunk.content,
                &chunk.source,
                &chunk.project,
                &chunk.level,
                &chunk.breadcrumb,
            )
            .await
        } else {
            // OpenSearch not configured, skip
            Ok(())
        };

        let opensearch_pending = opensearch_success.is_err();

        if opensearch_pending {
            tracing::warn!(
                "OpenSearch write failed for chunk {}, marked for retry: {}",
                chunk_id,
                opensearch_success.as_ref().err().unwrap()
            );
            // Mark as pending in pgvector
            self.mark_opensearch_pending(&chunk_id).await.ok();
        }

        // Step 4: Update indexed flags
        let pgvector_ok = pgvector_success.is_ok();
        let opensearch_ok = opensearch_success.is_ok();
        
        if pgvector_ok {
            self.update_pgvector_indexed(&chunk_id).await.ok();
        }

        if opensearch_ok {
            self.update_opensearch_indexed(&chunk_id).await.ok();
        }

        Ok(DualWriteResult {
            chunk_id,
            chunk_hash,
            pgvector_success: pgvector_ok,
            opensearch_success: opensearch_ok,
            opensearch_pending,
            error: if opensearch_pending {
                Some(format!("{:?}", opensearch_success.err()))
            } else {
                None
            },
        })
    }

    /// Compute SHA256 hash of content for deduplication
    fn compute_hash(&self, content: &str) -> String {
        let mut hasher = Sha256::new();
        hasher.update(content.as_bytes());
        format!("{:x}", hasher.finalize())
    }

    /// Check if chunk is already fully indexed
    async fn is_already_indexed(&self, chunk_hash: &str, project: &str) -> Result<bool> {
        let row = sqlx::query_scalar::<_, bool>(
            "SELECT (indexed_in_pgvector AND indexed_in_opensearch) 
             FROM chunks 
             WHERE chunk_hash = $1 AND project = $2 
             LIMIT 1"
        )
        .bind(chunk_hash)
        .bind(project)
        .fetch_optional(&self.pool)
        .await?;

        Ok(row.unwrap_or(false))
    }

    /// Write chunk to pgvector
    async fn write_pgvector(
        &self,
        chunk_id: &Uuid,
        chunk_hash: &str,
        content: &str,
        source: &str,
        project: &str,
        level: &str,
        breadcrumb: &[String],
        embedding: &[f32],
    ) -> Result<()> {
        let embedding_vec = Vector::from(embedding.to_vec());

        sqlx::query(
            "INSERT INTO chunks (id, chunk_hash, content, source, project, level, breadcrumb, embedding, indexed_in_pgvector, pgvector_indexed_at)
             VALUES ($1, $2, $3, $4, $5, $6, $7, $8, true, now())
             ON CONFLICT (id) DO UPDATE SET
               embedding = EXCLUDED.embedding,
               indexed_in_pgvector = true,
               pgvector_indexed_at = now()"
        )
        .bind(chunk_id)
        .bind(chunk_hash)
        .bind(content)
        .bind(source)
        .bind(project)
        .bind(level)
        .bind(breadcrumb)
        .bind(embedding_vec)
        .execute(&self.pool)
        .await?;

        Ok(())
    }

    /// Write chunk to OpenSearch
    async fn write_opensearch(
        &self,
        os_client: &Arc<OpenSearchClient>,
        chunk_id: &Uuid,
        content: &str,
        source: &str,
        project: &str,
        level: &str,
        breadcrumb: &[String],
    ) -> Result<()> {
        // Note: JWT token handling would come from AppState in http_server
        // For now, we'll pass empty token—production code should inject from context
        os_client
            .index_document(
                &chunk_id.to_string(),
                content,
                source,
                level,
                breadcrumb.to_vec(),
                "",  // TODO: inject JWT from AppState
            )
            .await?;

        Ok(())
    }

    /// Mark chunk as pending OpenSearch retry
    async fn mark_opensearch_pending(&self, chunk_id: &Uuid) -> Result<()> {
        sqlx::query(
            "UPDATE chunks 
             SET opensearch_pending = true, opensearch_retry_count = opensearch_retry_count + 1, opensearch_last_retry_at = now()
             WHERE id = $1"
        )
        .bind(chunk_id)
        .execute(&self.pool)
        .await?;

        Ok(())
    }

    /// Mark chunk as pgvector indexed
    async fn update_pgvector_indexed(&self, chunk_id: &Uuid) -> Result<()> {
        sqlx::query(
            "UPDATE chunks SET indexed_in_pgvector = true, pgvector_indexed_at = now() WHERE id = $1"
        )
        .bind(chunk_id)
        .execute(&self.pool)
        .await?;

        Ok(())
    }

    /// Mark chunk as OpenSearch indexed
    async fn update_opensearch_indexed(&self, chunk_id: &Uuid) -> Result<()> {
        sqlx::query(
            "UPDATE chunks SET indexed_in_opensearch = true, opensearch_pending = false, opensearch_indexed_at = now() WHERE id = $1"
        )
        .bind(chunk_id)
        .execute(&self.pool)
        .await?;

        Ok(())
    }

    /// Retry failed OpenSearch writes (background task)
    ///
    /// Polls for chunks where opensearch_pending=true and retries up to 3 times.
    /// Runs every 5 minutes.
    pub async fn retry_pending_chunks(&self, project: &str, max_retries: i32) -> Result<usize> {
        if self.opensearch.is_none() {
            return Ok(0);  // Skip if OpenSearch not configured
        }

        let pending = sqlx::query_as::<_, (Uuid, String, String, String, Vec<String>)>(
            "SELECT id, content, source, level, breadcrumb 
             FROM chunks 
             WHERE project = $1 AND opensearch_pending = true AND opensearch_retry_count < $2
             ORDER BY opensearch_last_retry_at ASC 
             LIMIT 100"
        )
        .bind(project)
        .bind(max_retries)
        .fetch_all(&self.pool)
        .await?;

        let mut succeeded = 0;

        for (chunk_id, content, source, level, breadcrumb) in pending {
            if let Err(e) = self
                .write_opensearch(
                    self.opensearch.as_ref().unwrap(),
                    &chunk_id,
                    &content,
                    &source,
                    project,
                    &level,
                    &breadcrumb,
                )
                .await
            {
                tracing::warn!("Retry failed for chunk {}: {}", chunk_id, e);
                // Increment retry count
                sqlx::query(
                    "UPDATE chunks SET opensearch_retry_count = opensearch_retry_count + 1, opensearch_last_retry_at = now() WHERE id = $1"
                )
                .bind(&chunk_id)
                .execute(&self.pool)
                .await
                .ok();
            } else {
                tracing::info!("Retry succeeded for chunk {}", chunk_id);
                self.update_opensearch_indexed(&chunk_id).await.ok();
                succeeded += 1;
            }
        }

        Ok(succeeded)
    }

    /// Get retry statistics
    pub async fn retry_stats(&self, project: &str) -> Result<(usize, usize)> {
        let pending: (i64,) = sqlx::query_as(
            "SELECT COUNT(*) FROM chunks WHERE project = $1 AND opensearch_pending = true"
        )
        .bind(project)
        .fetch_one(&self.pool)
        .await?;

        let failed: (i64,) = sqlx::query_as(
            "SELECT COUNT(*) FROM chunks WHERE project = $1 AND opensearch_retry_count >= 3"
        )
        .bind(project)
        .fetch_one(&self.pool)
        .await?;

        Ok((pending.0 as usize, failed.0 as usize))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_compute_hash() {
        let indexer = DualWriteIndexer::new(
            sqlx::pool::PoolOptions::new().max_connections(1).connect_lazy("postgresql://localhost").unwrap(),
            None,
        );

        let hash1 = indexer.compute_hash("same content");
        let hash2 = indexer.compute_hash("same content");
        assert_eq!(hash1, hash2, "Same content must produce same hash");

        let hash3 = indexer.compute_hash("different");
        assert_ne!(hash1, hash3, "Different content must produce different hash");
    }

    #[test]
    fn test_hash_deterministic() {
        let indexer = DualWriteIndexer::new(
            sqlx::pool::PoolOptions::new().max_connections(1).connect_lazy("postgresql://localhost").unwrap(),
            None,
        );

        let content = "ERROR: permission denied\nStack trace...";
        let hash1 = indexer.compute_hash(content);
        let hash2 = indexer.compute_hash(content);

        assert_eq!(hash1, hash2);
        assert_eq!(hash1.len(), 64); // SHA256 hex is 64 chars
    }
}