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

use _utils::{boxed_error, cleanup_run, create_client, new_run_id, print_summary, require};
use mubit_sdk::{
    ArchiveOptions, CheckpointOptions, DereferenceOptions, DiagnoseOptions, FeedbackOptions,
    ForgetOptions, GetContextOptions, HandoffOptions, ListAgentsOptions, MemoryHealthOptions, RecallOptions,
    RecordOutcomeOptions, ReflectOptions, RegisterAgentOptions, RememberOptions,
    SurfaceStrategiesOptions, TransportMode,
};
use serde_json::{json, Value};
use std::error::Error;
use std::time::Instant;

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

    let mut passed = true;
    let mut detail = "validated all helper flows".to_string();
    let mut metrics = json!({});

    let scenario = async {
        let mut remember = RememberOptions::new("Policy limit is 5000 and claimant city is Leeds.");
        remember.run_id = Some(run_id.clone());
        remember.agent_id = Some("planner".to_string());
        remember.intent = Some("fact".to_string());
        remember.metadata = Some(json!({"suite": "all-helpers", "family": "claims"}));
        remember.importance = Some("high".to_string());
        let remember_response = client.remember(remember).await?;

        let mut lesson = RememberOptions::new("If city is missing from the latest step, check stored claim facts before escalating.");
        lesson.run_id = Some(run_id.clone());
        lesson.agent_id = Some("planner".to_string());
        lesson.intent = Some("lesson".to_string());
        lesson.lesson_type = Some("success".to_string());
        lesson.lesson_scope = Some("session".to_string());
        lesson.lesson_importance = Some("high".to_string());
        let lesson_response = client.remember(lesson).await?;

        let mut recall = RecallOptions::new("What is the policy limit?");
        recall.run_id = Some(run_id.clone());
        recall.agent_id = Some("planner".to_string());
        recall.limit = 5;
        let recall_response = client.recall(recall).await?;

        let mut context = GetContextOptions::default();
        context.run_id = Some(run_id.clone());
        context.query = Some("Prepare a short adjuster summary".to_string());
        context.agent_id = Some("planner".to_string());
        context.limit = Some(6);
        context.max_token_budget = Some(700);
        let context_response = client.get_context(context).await?;

        let mut health = MemoryHealthOptions::default();
        health.run_id = Some(run_id.clone());
        health.limit = 50;
        let health_response = client.memory_health(health).await?;

        let mut diagnose = DiagnoseOptions::new("policy limit lookup returned empty");
        diagnose.run_id = Some(run_id.clone());
        diagnose.error_type = Some("retrieval".to_string());
        diagnose.limit = 5;
        let diagnose_response = client.diagnose(diagnose).await?;

        let mut reflect = ReflectOptions::default();
        reflect.run_id = Some(run_id.clone());
        let reflect_response = client.reflect(reflect).await?;

        let mut register = RegisterAgentOptions::new("planner");
        register.run_id = Some(run_id.clone());
        register.role = "planner".to_string();
        register.read_scopes = vec!["fact".into(), "lesson".into(), "rule".into(), "archive_block".into(), "handoff".into(), "feedback".into()];
        register.write_scopes = vec!["fact".into(), "trace".into(), "lesson".into(), "observation".into(), "archive_block".into()];
        register.shared_memory_lanes = vec!["knowledge".into(), "history".into()];
        let register_response = client.register_agent(register).await?;

        let mut list_agents = ListAgentsOptions::default();
        list_agents.run_id = Some(run_id.clone());
        let agents_response = client.list_agents(list_agents).await?;

        let mut archive = ArchiveOptions::new(
            "Exact adjuster note: claimant confirmed city is Leeds during triage.",
            "adjuster_note",
        );
        archive.run_id = Some(run_id.clone());
        archive.agent_id = Some("planner".to_string());
        archive.origin_agent_id = Some("planner".to_string());
        archive.source_attempt_id = Some("attempt-1".to_string());
        archive.source_tool = Some("phone_call".to_string());
        archive.labels = vec!["claims".to_string(), "triage".to_string()];
        archive.family = Some("claims-adjustment".to_string());
        archive.metadata = Some(json!({"case": "all-helpers"}));
        let archive_response = client.archive(archive).await?;
        let reference_id = archive_response
            .get("reference_id")
            .and_then(Value::as_str)
            .ok_or_else(|| boxed_error(format!("archive reference missing: {archive_response}")))?
            .to_string();

        let mut dereference = DereferenceOptions::new(reference_id.clone());
        dereference.run_id = Some(run_id.clone());
        dereference.agent_id = Some("planner".to_string());
        let dereference_response = client.dereference(dereference).await?;

        let mut checkpoint = CheckpointOptions::new("Planner captured claim facts before window compaction.");
        checkpoint.run_id = Some(run_id.clone());
        checkpoint.label = Some("pre-compaction".to_string());
        checkpoint.metadata = Some(json!({"stage": "analysis"}));
        checkpoint.agent_id = Some("planner".to_string());
        let checkpoint_response = client.checkpoint(checkpoint).await?;

        let lessons_response = client.control.lessons(json!({"run_id": run_id, "limit": 20})).await?;
        let lesson_id = lessons_response
            .get("lessons")
            .and_then(Value::as_array)
            .and_then(|items| items.iter().find_map(|item| item.get("id").and_then(Value::as_str)))
            .ok_or_else(|| boxed_error(format!("expected lesson id for outcome flow: {lessons_response}")))?
            .to_string();

        let mut outcome = RecordOutcomeOptions::new(lesson_id, "success");
        outcome.run_id = Some(run_id.clone());
        outcome.signal = 0.75;
        outcome.rationale = "Planner reused stored claim fact and resolved the missing-city branch.".to_string();
        outcome.agent_id = Some("planner".to_string());
        let outcome_response = client.record_outcome(outcome).await?;

        let mut strategies = SurfaceStrategiesOptions::default();
        strategies.run_id = Some(run_id.clone());
        strategies.lesson_types = vec!["success".to_string(), "rule".to_string()];
        strategies.max_strategies = 5;
        let strategies_response = client.surface_strategies(strategies).await?;

        let mut handoff = HandoffOptions::new(
            "claim-1",
            "planner",
            "reviewer",
            "Review whether the stored claimant city is sufficient for routing.",
        );
        handoff.run_id = Some(run_id.clone());
        handoff.requested_action = "review".to_string();
        handoff.metadata = Some(json!({"team": "claims"}));
        let handoff_response = client.handoff(handoff).await?;
        let handoff_id = handoff_response
            .get("handoff_id")
            .and_then(Value::as_str)
            .ok_or_else(|| boxed_error(format!("handoff id missing: {handoff_response}")))?
            .to_string();

        let mut feedback = FeedbackOptions::new(handoff_id.clone(), "approve");
        feedback.run_id = Some(run_id.clone());
        feedback.comments = "Stored fact is sufficient; proceed without escalation.".to_string();
        feedback.from_agent_id = Some("reviewer".to_string());
        feedback.metadata = Some(json!({"team": "claims"}));
        let feedback_response = client.feedback(feedback).await?;

        let scratch_run = format!("{run_id}_forget");
        let mut scratch = RememberOptions::new("Scratch run for forget helper coverage.");
        scratch.run_id = Some(scratch_run.clone());
        scratch.agent_id = Some("planner".to_string());
        scratch.intent = Some("fact".to_string());
        scratch.metadata = Some(json!({"suite": "all-helpers", "temporary": true}));
        let scratch_response = client.remember(scratch).await?;

        let forget_response = client
            .forget(ForgetOptions::for_run(scratch_run.clone()))
            .await?;

        let mut forgotten = RecallOptions::new("What temporary scratch fact exists?");
        forgotten.run_id = Some(scratch_run.clone());
        forgotten.limit = 3;
        let forgotten_response = client.recall(forgotten).await?;

        require(remember_response.get("job_id").is_some() || remember_response.get("accepted").and_then(Value::as_bool).unwrap_or(false), format!("remember failed: {remember_response}"))?;
        require(lesson_response.get("job_id").is_some() || lesson_response.get("accepted").and_then(Value::as_bool).unwrap_or(false), format!("lesson remember failed: {lesson_response}"))?;
        require(recall_response.get("final_answer").is_some() || recall_response.get("evidence").is_some(), format!("recall failed: {recall_response}"))?;
        require(context_response.get("sources").and_then(Value::as_array).is_some(), format!("context malformed: {context_response}"))?;
        require(health_response.get("entry_counts").and_then(Value::as_object).is_some(), format!("memory_health malformed: {health_response}"))?;
        require(diagnose_response.get("failure_lessons").and_then(Value::as_array).is_some(), format!("diagnose malformed: {diagnose_response}"))?;
        require(reflect_response.get("lessons").and_then(Value::as_array).is_some(), format!("reflect malformed: {reflect_response}"))?;
        require(register_response.get("success").and_then(Value::as_bool).unwrap_or(false), format!("register_agent failed: {register_response}"))?;
        require(agents_response.get("agents").and_then(Value::as_array).map(|items| !items.is_empty()).unwrap_or(false), format!("list_agents failed: {agents_response}"))?;
        require(archive_response.get("success").and_then(Value::as_bool).unwrap_or(false), format!("archive failed: {archive_response}"))?;
        require(dereference_response.get("found").and_then(Value::as_bool).unwrap_or(false), format!("dereference failed: {dereference_response}"))?;
        require(checkpoint_response.get("success").and_then(Value::as_bool).unwrap_or(false), format!("checkpoint failed: {checkpoint_response}"))?;
        require(outcome_response.get("success").and_then(Value::as_bool).unwrap_or(false), format!("record_outcome failed: {outcome_response}"))?;
        require(strategies_response.get("strategies").and_then(Value::as_array).is_some(), format!("surface_strategies malformed: {strategies_response}"))?;
        require(handoff_response.get("success").and_then(Value::as_bool).unwrap_or(false), format!("handoff failed: {handoff_response}"))?;
        require(feedback_response.get("success").and_then(Value::as_bool).unwrap_or(false), format!("feedback failed: {feedback_response}"))?;
        require(scratch_response.get("job_id").is_some() || scratch_response.get("accepted").and_then(Value::as_bool).unwrap_or(false), format!("scratch remember failed: {scratch_response}"))?;
        require(forget_response.get("success").and_then(Value::as_bool).unwrap_or(true), format!("forget failed: {forget_response}"))?;
        let scratch_fact_visible = forgotten_response
            .get("evidence")
            .and_then(Value::as_array)
            .map(|items| {
                items.iter().any(|item| {
                    item.get("run_id").and_then(Value::as_str) == Some(scratch_run.as_str())
                        || item
                            .get("content")
                            .and_then(Value::as_str)
                            .map(|content| {
                                content
                                    .to_ascii_lowercase()
                                    .contains("scratch run for forget helper coverage")
                            })
                            .unwrap_or(false)
                })
            })
            .unwrap_or(false);
        require(!scratch_fact_visible, format!("scratch fact should be forgotten even if same-user lessons still surface: {forgotten_response}"))?;

        let strategy_count = strategies_response
            .get("strategies")
            .and_then(Value::as_array)
            .map(|items| items.len())
            .unwrap_or(0);
        let agent_count = agents_response
            .get("agents")
            .and_then(Value::as_array)
            .map(|items| items.len())
            .unwrap_or(0);
        let lesson_count = lessons_response
            .get("lessons")
            .and_then(Value::as_array)
            .map(|items| items.len())
            .unwrap_or(0);

        metrics = json!({
            "run_id": run_id,
            "agent_count": agent_count,
            "strategy_count": strategy_count,
            "reference_id": reference_id,
            "lesson_count": lesson_count,
            "forgotten_run": scratch_run,
        });

        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)) }
}
