adk-session 0.8.0

use adk_session::{
    CreateRequest, DeleteRequest, Event, GetRequest, KEY_PREFIX_TEMP, ListRequest, SessionService,
};
use chrono::{Duration, Utc};
use serde_json::json;
use std::collections::HashMap;

#[allow(dead_code)]
pub async fn assert_session_contract(
    service: &dyn SessionService,
    app_name: &str,
    other_app_name: &str,
) {
    assert_session_contract_with_users(service, app_name, other_app_name, "user1", "user2").await;
}

pub async fn assert_session_contract_with_users(
    service: &dyn SessionService,
    app_name: &str,
    other_app_name: &str,
    user_1: &str,
    user_2: &str,
) {
    let mut initial_state = HashMap::new();
    initial_state.insert("app:locale".to_string(), json!("en-US"));
    initial_state.insert("user:name".to_string(), json!("alice"));
    initial_state.insert("session_key".to_string(), json!("seed"));
    initial_state.insert(format!("{}create", KEY_PREFIX_TEMP), json!("drop-me"));

    let created = service
        .create(CreateRequest {
            app_name: app_name.to_string(),
            user_id: user_1.to_string(),
            session_id: None,
            state: initial_state,
        })
        .await
        .expect("create session should succeed");

    let session_id = created.id().to_string();
    assert!(!session_id.is_empty());
    assert_eq!(created.app_name(), app_name);
    assert_eq!(created.user_id(), user_1);

    let fetched = service
        .get(GetRequest {
            app_name: app_name.to_string(),
            user_id: user_1.to_string(),
            session_id: session_id.clone(),
            num_recent_events: None,
            after: None,
        })
        .await
        .expect("get session should succeed");

    assert_eq!(fetched.id(), session_id);
    assert_eq!(fetched.state().get("app:locale"), Some(json!("en-US")));
    assert_eq!(fetched.state().get("user:name"), Some(json!("alice")));
    assert_eq!(fetched.state().get("session_key"), Some(json!("seed")));
    assert_eq!(fetched.state().get("temp:create"), None);

    let t1 = Utc::now();
    let t2 = t1 + Duration::seconds(1);

    let mut event_1 = Event::new("inv-1");
    event_1.author = "agent".to_string();
    event_1.timestamp = t1;
    event_1.actions.state_delta.insert("result".to_string(), json!("ok-1"));
    event_1.actions.state_delta.insert(format!("{}event", KEY_PREFIX_TEMP), json!("skip"));

    service.append_event(&session_id, event_1).await.expect("append first event should succeed");

    let mut event_2 = Event::new("inv-2");
    event_2.author = "agent".to_string();
    event_2.timestamp = t2;
    event_2.actions.state_delta.insert("result".to_string(), json!("ok-2"));

    service.append_event(&session_id, event_2).await.expect("append second event should succeed");

    let with_events = service
        .get(GetRequest {
            app_name: app_name.to_string(),
            user_id: user_1.to_string(),
            session_id: session_id.clone(),
            num_recent_events: None,
            after: None,
        })
        .await
        .expect("get after append should succeed");

    assert_eq!(with_events.events().len(), 2);
    assert_eq!(with_events.state().get("result"), Some(json!("ok-2")));
    assert_eq!(with_events.state().get("temp:event"), None);

    let recent = service
        .get(GetRequest {
            app_name: app_name.to_string(),
            user_id: user_1.to_string(),
            session_id: session_id.clone(),
            num_recent_events: Some(1),
            after: None,
        })
        .await
        .expect("get recent events should succeed");

    assert_eq!(recent.events().len(), 1);
    assert_eq!(recent.events().at(0).expect("event 0").timestamp, t2);

    let after = service
        .get(GetRequest {
            app_name: app_name.to_string(),
            user_id: user_1.to_string(),
            session_id: session_id.clone(),
            num_recent_events: None,
            after: Some(t2),
        })
        .await
        .expect("get with after filter should succeed");

    assert_eq!(after.events().len(), 1);
    assert_eq!(after.events().at(0).expect("event 0").timestamp, t2);

    let sessions_user1 = service
        .list(ListRequest {
            app_name: app_name.to_string(),
            user_id: user_1.to_string(),
            limit: None,
            offset: None,
        })
        .await
        .expect("list for user1 should succeed");

    assert!(sessions_user1.iter().any(|session| session.id() == session_id));

    let user2 = service
        .create(CreateRequest {
            app_name: app_name.to_string(),
            user_id: user_2.to_string(),
            session_id: None,
            state: HashMap::new(),
        })
        .await
        .expect("create session for user2 should succeed");
    let user2_session_id = user2.id().to_string();

    let wrong_user_get = service
        .get(GetRequest {
            app_name: app_name.to_string(),
            user_id: user_2.to_string(),
            session_id: session_id.clone(),
            num_recent_events: None,
            after: None,
        })
        .await;
    assert!(wrong_user_get.is_err());

    let sessions_user2 = service
        .list(ListRequest {
            app_name: app_name.to_string(),
            user_id: user_2.to_string(),
            limit: None,
            offset: None,
        })
        .await
        .expect("list for user2 should succeed");

    assert!(sessions_user2.iter().any(|session| session.id() == user2_session_id));
    assert!(!sessions_user2.iter().any(|session| session.id() == session_id));

    let other_app = service
        .create(CreateRequest {
            app_name: other_app_name.to_string(),
            user_id: user_1.to_string(),
            session_id: None,
            state: HashMap::new(),
        })
        .await
        .expect("create session for second app should succeed");
    let other_app_session_id = other_app.id().to_string();

    let sessions_primary_app = service
        .list(ListRequest {
            app_name: app_name.to_string(),
            user_id: user_1.to_string(),
            limit: None,
            offset: None,
        })
        .await
        .expect("list primary app should succeed");

    assert!(sessions_primary_app.iter().any(|session| session.id() == session_id));
    assert!(!sessions_primary_app.iter().any(|session| session.id() == other_app_session_id));

    let sessions_other_app = service
        .list(ListRequest {
            app_name: other_app_name.to_string(),
            user_id: user_1.to_string(),
            limit: None,
            offset: None,
        })
        .await
        .expect("list secondary app should succeed");

    assert!(sessions_other_app.iter().any(|session| session.id() == other_app_session_id));
    assert!(!sessions_other_app.iter().any(|session| session.id() == session_id));

    service
        .delete(DeleteRequest {
            app_name: app_name.to_string(),
            user_id: user_1.to_string(),
            session_id: session_id.clone(),
        })
        .await
        .expect("delete should succeed");

    service
        .delete(DeleteRequest {
            app_name: app_name.to_string(),
            user_id: user_2.to_string(),
            session_id: user2_session_id,
        })
        .await
        .expect("delete secondary user session should succeed");

    service
        .delete(DeleteRequest {
            app_name: other_app_name.to_string(),
            user_id: user_1.to_string(),
            session_id: other_app_session_id,
        })
        .await
        .expect("delete secondary app session should succeed");

    let deleted_get = service
        .get(GetRequest {
            app_name: app_name.to_string(),
            user_id: user_1.to_string(),
            session_id,
            num_recent_events: None,
            after: None,
        })
        .await;
    assert!(deleted_get.is_err());
}