remem-ai 0.5.96

use anyhow::Result;
use rusqlite::Connection;
use std::time::Instant;

use crate::db;

pub(in crate::cli) fn run_backfill_entities() -> Result<()> {
    let conn = db::open_db()?;
    let count: i64 = conn
        .query_row(
            "SELECT COUNT(*) FROM memories WHERE status = 'active'",
            [],
            |row| row.get(0),
        )
        .unwrap_or(0);
    println!("Backfilling entities from {} active memories...", count);

    let mut stmt =
        conn.prepare("SELECT id, title, content FROM memories WHERE status = 'active'")?;
    let rows = stmt.query_map([], |row| {
        Ok((
            row.get::<_, i64>(0)?,
            row.get::<_, String>(1)?,
            row.get::<_, String>(2)?,
        ))
    })?;

    print_backfill_progress(&conn, rows, count)
}

fn print_backfill_progress<F>(
    conn: &Connection,
    rows: rusqlite::MappedRows<'_, F>,
    count: i64,
) -> Result<()>
where
    F: FnMut(&rusqlite::Row<'_>) -> rusqlite::Result<(i64, String, String)>,
{
    let mut total_entities = 0usize;
    let mut memories_processed = 0usize;

    for row in rows {
        let (id, title, content) = row?;
        let entities = crate::retrieval::entity::extract_entities(&title, &content);
        if !entities.is_empty() {
            crate::retrieval::entity::link_entities(conn, id, &entities)?;
            total_entities += entities.len();
        }
        memories_processed += 1;
        if memories_processed.is_multiple_of(100) {
            println!("  processed {}/{}", memories_processed, count);
        }
    }

    let unique: i64 = conn
        .query_row("SELECT COUNT(*) FROM entities", [], |row| row.get(0))
        .unwrap_or(0);
    println!(
        "Done. {} entities extracted, {} unique entities, {} memories processed.",
        total_entities, unique, memories_processed
    );
    Ok(())
}

pub(in crate::cli) fn run_backfill_embeddings(limit: i64, batch_size: i64) -> Result<()> {
    let conn = db::open_db()?;
    let limit = limit.max(1);
    let batch_size = batch_size.max(1);
    let pending_before = count_missing_embeddings(&conn)?;
    println!(
        "Backfilling/reindexing up to {limit} missing or stale memory embeddings, batch_size={batch_size}, pending_before={pending_before}..."
    );

    let started = Instant::now();
    let mut backfilled = 0usize;
    let mut remaining_limit = limit;
    let mut printed_profile = false;
    while remaining_limit > 0 {
        let batch_limit = remaining_limit.min(batch_size);
        let report =
            crate::retrieval::vector::reindex_memory_embeddings_with_report(&conn, batch_limit)?;
        if !printed_profile && !report.model.is_empty() {
            println!(
                "Embedding profile: model={} dimensions={}",
                report.model, report.dimensions
            );
            printed_profile = true;
        }
        if report.processed == 0 {
            break;
        }

        backfilled += report.processed;
        remaining_limit -= report.processed as i64;
        let remaining = count_missing_embeddings(&conn)?;
        let elapsed_ms = report
            .timings
            .iter()
            .find(|timing| timing.phase == "total")
            .map(|timing| timing.elapsed_ms)
            .unwrap_or(0);
        let rows_per_sec = if elapsed_ms == 0 {
            report.processed as f64
        } else {
            report.processed as f64 * 1000.0 / elapsed_ms as f64
        };
        println!(
            "  batch processed={} selected={} remaining={} rows_per_sec={rows_per_sec:.1} {}",
            report.processed,
            report.selected,
            remaining,
            crate::perf::format_phase_timings(&report.timings)
        );
        if report.processed < batch_limit as usize {
            break;
        }
    }

    let remaining = count_missing_embeddings(&conn)?;
    println!(
        "Done. {backfilled} embeddings backfilled/reindexed, {remaining} remaining, elapsed_ms={}.",
        started.elapsed().as_millis()
    );
    Ok(())
}

fn count_missing_embeddings(conn: &Connection) -> Result<i64> {
    crate::retrieval::vector::pending_memory_embedding_reindex_count(conn)
}