sol/src/memory/store.rs

use chrono::Utc;
use opensearch::OpenSearch;
use serde_json::json;
use uuid::Uuid;

use super::schema::MemoryDocument;

/// Search memories by content relevance, filtered to a specific user.
pub async fn query(
    client: &OpenSearch,
    index: &str,
    user_id: &str,
    query_text: &str,
    limit: usize,
) -> anyhow::Result<Vec<MemoryDocument>> {
    let body = json!({
        "size": limit,
        "query": {
            "bool": {
                "filter": [
                    { "term": { "user_id": user_id } }
                ],
                "must": [
                    { "match": { "content": query_text } }
                ]
            }
        },
        "sort": [{ "_score": "desc" }]
    });

    let response = client
        .search(opensearch::SearchParts::Index(&[index]))
        .body(body)
        .send()
        .await?;

    let data: serde_json::Value = response.json().await?;
    parse_hits(&data)
}

/// Get the most recent memories for a user, sorted by updated_at desc.
pub async fn get_recent(
    client: &OpenSearch,
    index: &str,
    user_id: &str,
    limit: usize,
) -> anyhow::Result<Vec<MemoryDocument>> {
    let body = json!({
        "size": limit,
        "query": {
            "bool": {
                "filter": [
                    { "term": { "user_id": user_id } }
                ]
            }
        },
        "sort": [{ "updated_at": "desc" }]
    });

    let response = client
        .search(opensearch::SearchParts::Index(&[index]))
        .body(body)
        .send()
        .await?;

    let data: serde_json::Value = response.json().await?;
    parse_hits(&data)
}

/// Store a new memory document for a user.
pub async fn set(
    client: &OpenSearch,
    index: &str,
    user_id: &str,
    content: &str,
    category: &str,
    source: &str,
) -> anyhow::Result<()> {
    let now = Utc::now().timestamp_millis();
    let id = Uuid::new_v4().to_string();

    let doc = MemoryDocument {
        id: id.clone(),
        user_id: user_id.to_string(),
        content: content.to_string(),
        category: category.to_string(),
        created_at: now,
        updated_at: now,
        source: source.to_string(),
    };

    let response = client
        .index(opensearch::IndexParts::IndexId(index, &id))
        .body(serde_json::to_value(&doc)?)
        .send()
        .await?;

    if !response.status_code().is_success() {
        let body = response.text().await?;
        anyhow::bail!("Failed to store memory: {body}");
    }

    Ok(())
}

pub(crate) fn parse_hits(data: &serde_json::Value) -> anyhow::Result<Vec<MemoryDocument>> {
    let hits = data["hits"]["hits"]
        .as_array()
        .cloned()
        .unwrap_or_default();

    let mut docs = Vec::with_capacity(hits.len());
    for hit in &hits {
        if let Ok(doc) = serde_json::from_value::<MemoryDocument>(hit["_source"].clone()) {
            docs.push(doc);
        }
    }
    Ok(docs)
}

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

    fn fake_os_response(sources: Vec<serde_json::Value>) -> serde_json::Value {
        let hits: Vec<serde_json::Value> = sources
            .into_iter()
            .map(|s| json!({ "_source": s }))
            .collect();
        json!({ "hits": { "hits": hits } })
    }

    #[test]
    fn test_parse_hits_multiple() {
        let data = fake_os_response(vec![
            json!({
                "id": "a", "user_id": "sienna@sunbeam.pt",
                "content": "prefers terse answers", "category": "preference",
                "created_at": 1710000000000_i64, "updated_at": 1710000000000_i64,
                "source": "auto"
            }),
            json!({
                "id": "b", "user_id": "sienna@sunbeam.pt",
                "content": "working on drive UI", "category": "fact",
                "created_at": 1710000000000_i64, "updated_at": 1710000000000_i64,
                "source": "script"
            }),
        ]);

        let docs = parse_hits(&data).unwrap();
        assert_eq!(docs.len(), 2);
        assert_eq!(docs[0].id, "a");
        assert_eq!(docs[0].content, "prefers terse answers");
        assert_eq!(docs[1].id, "b");
        assert_eq!(docs[1].category, "fact");
    }

    #[test]
    fn test_parse_hits_empty() {
        let data = json!({ "hits": { "hits": [] } });
        let docs = parse_hits(&data).unwrap();
        assert!(docs.is_empty());
    }

    #[test]
    fn test_parse_hits_missing_structure() {
        let data = json!({});
        let docs = parse_hits(&data).unwrap();
        assert!(docs.is_empty());
    }

    #[test]
    fn test_parse_hits_skips_malformed() {
        let data = fake_os_response(vec![
            json!({
                "id": "good", "user_id": "x@y",
                "content": "ok", "category": "fact",
                "created_at": 1, "updated_at": 1, "source": "auto"
            }),
            json!({ "bad": "no fields" }),
        ]);

        let docs = parse_hits(&data).unwrap();
        assert_eq!(docs.len(), 1);
        assert_eq!(docs[0].id, "good");
    }
}