#[path = "../support/utils.rs"]
mod _utils;

use _utils::{boxed_error, cleanup_run, create_client, new_run_id, print_summary, require};
use mubit_sdk::{GetContextOptions, RecallOptions, RememberOptions, TransportMode};
use serde_json::json;
use std::error::Error;
use std::time::Instant;

#[tokio::main(flavor = "current_thread")]
async fn main() -> Result<(), Box<dyn Error>> {
    let name = "public_11_temporal_entity_mental_model";
    let started = Instant::now();
    let client = create_client().await?;
    let run_id = new_run_id("public_temporal_entity");
    client.set_run_id(Some(run_id.clone()));
    client.set_transport(TransportMode::Http);

    let mut passed = true;
    let mut detail =
        "validated temporal queries, entity aliasing, and mental model priority".to_string();
    let mut metrics = json!({});

    let scenario = async {
        // 1. Store facts with occurrence_time
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_secs() as i64;

        let mut fact1 =
            RememberOptions::new("Alice Chen joined the platform team in March 2024.");
        fact1.run_id = Some(run_id.clone());
        fact1.intent = Some("fact".to_string());
        fact1.occurrence_time = Some(now - 86400 * 365); // ~1 year ago
        fact1.metadata = Some(json!({"speaker": "Alice Chen"}));
        client.remember(fact1).await?;

        let mut fact2 =
            RememberOptions::new("Alice was promoted to Senior Engineer in January 2025.");
        fact2.run_id = Some(run_id.clone());
        fact2.intent = Some("fact".to_string());
        fact2.occurrence_time = Some(now - 86400 * 60); // ~2 months ago
        fact2.metadata = Some(json!({"speaker": "Alice Chen"}));
        client.remember(fact2).await?;

        let mut fact3 =
            RememberOptions::new("Dr. Chen published a paper on distributed consensus.");
        fact3.run_id = Some(run_id.clone());
        fact3.intent = Some("fact".to_string());
        fact3.metadata = Some(json!({"speaker": "Alice Chen"}));
        client.remember(fact3).await?;

        // 2. Store a mental model (curated summary)
        let mut mental = RememberOptions::new(
            "Alice Chen (Dr. Chen) is a Senior Engineer on the platform team. \
             She specializes in distributed consensus and was promoted in January 2025.",
        );
        mental.run_id = Some(run_id.clone());
        mental.intent = Some("mental_model".to_string());
        mental.importance = Some("critical".to_string());
        mental.metadata = Some(json!({"entity": "alice chen", "consolidated": true}));
        client.remember(mental).await?;

        // 3. Recall with temporal bounds (last 90 days)
        let mut temporal_recall = RecallOptions::new("What happened recently with Alice?");
        temporal_recall.run_id = Some(run_id.clone());
        temporal_recall.min_timestamp = Some(now - 86400 * 90);
        temporal_recall.max_timestamp = Some(now);
        let temporal_result = client.recall(temporal_recall).await?;
        let temporal_evidence = temporal_result
            .get("evidence")
            .and_then(|v| v.as_array())
            .ok_or_else(|| {
                boxed_error("temporal recall missing evidence".to_string())
            })?;
        require(
            !temporal_evidence.is_empty(),
            format!("temporal recall should return evidence: {temporal_result}"),
        )?;

        // 4. Recall with budget tier
        let mut budget_recall = RecallOptions::new("Tell me about Alice Chen");
        budget_recall.run_id = Some(run_id.clone());
        budget_recall.budget = Some("high".to_string());
        let budget_result = client.recall(budget_recall).await?;
        let budget_evidence = budget_result
            .get("evidence")
            .and_then(|v| v.as_array())
            .ok_or_else(|| {
                boxed_error("budget recall missing evidence".to_string())
            })?;
        require(
            !budget_evidence.is_empty(),
            format!("budget recall should return evidence: {budget_result}"),
        )?;

        // 5. get_context should show Mental Models section first
        let mut ctx_opts = GetContextOptions::default();
        ctx_opts.run_id = Some(run_id.clone());
        ctx_opts.query = Some("Who is Alice Chen?".to_string());
        ctx_opts.entry_types = vec!["mental_model".to_string(), "fact".to_string()];
        ctx_opts.limit = Some(10);
        let ctx = client.get_context(ctx_opts).await?;
        let context_block = ctx
            .get("context_block")
            .and_then(|v| v.as_str())
            .unwrap_or("");
        // Mental Models section should appear before Facts
        let has_mental = context_block.contains("Mental Model");
        require(
            has_mental || context_block.contains("mental_model"),
            format!(
                "context should include mental model section: {}",
                &context_block[..context_block.len().min(300)]
            ),
        )?;

        metrics = json!({
            "run_id": run_id,
            "temporal_evidence_count": temporal_evidence.len(),
            "budget_evidence_count": budget_evidence.len(),
            "context_has_mental_model": has_mental,
            "context_block_length": context_block.len(),
        });
        Ok::<(), Box<dyn Error>>(())
    }
    .await;

    if let Err(err) = scenario {
        passed = false;
        detail = err.to_string();
    }

    let cleanup_ok = cleanup_run(&client, &run_id).await;
    if !cleanup_ok {
        passed = false;
        detail = format!("{detail} | cleanup failures");
    }

    print_summary(
        name,
        passed,
        &detail,
        &metrics,
        started.elapsed().as_secs_f64(),
        cleanup_ok,
    );
    if passed {
        Ok(())
    } else {
        Err(boxed_error(detail))
    }
}
