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

use _utils::{boxed_error, cleanup_run, create_client, new_run_id, print_summary, require};
use mubit_sdk::{GetContextOptions, ListAgentsOptions, RegisterAgentOptions, 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_05_agent_registration_and_scoped_access";
    let started = Instant::now();
    let client = create_client().await?;
    let run_id = new_run_id("public_agent_scopes");
    client.set_run_id(Some(run_id.clone()));
    client.set_transport(TransportMode::Http);

    let mut passed = true;
    let mut detail = "validated agent registration and scoped access".to_string();
    let mut metrics = json!({});

    let scenario = async {
        let mut planner = RegisterAgentOptions::new("planner");
        planner.run_id = Some(run_id.clone());
        planner.role = "planner".to_string();
        planner.read_scopes = vec!["fact".into(), "lesson".into()];
        planner.write_scopes = vec!["fact".into(), "lesson".into()];
        planner.shared_memory_lanes = vec!["knowledge".into(), "history".into()];
        client.register_agent(planner).await?;

        let mut reviewer = RegisterAgentOptions::new("reviewer");
        reviewer.run_id = Some(run_id.clone());
        reviewer.role = "reviewer".to_string();
        reviewer.read_scopes = vec!["rule".into(), "trace".into()];
        reviewer.write_scopes = vec!["rule".into(), "trace".into()];
        reviewer.shared_memory_lanes = vec!["knowledge".into(), "history".into()];
        client.register_agent(reviewer).await?;

        let mut planner_fact = RememberOptions::new("Claim requires police report only when payout exceeds 3000 GBP.");
        planner_fact.run_id = Some(run_id.clone());
        planner_fact.agent_id = Some("planner".into());
        planner_fact.intent = Some("fact".into());
        planner_fact.metadata = Some(json!({"source": "planner"}));
        let planner_write = client.remember(planner_fact).await?;

        let mut reviewer_rule = RememberOptions::new("Escalate to compliance if a police report is missing after payout approval.");
        reviewer_rule.run_id = Some(run_id.clone());
        reviewer_rule.agent_id = Some("reviewer".into());
        reviewer_rule.intent = Some("rule".into());
        reviewer_rule.metadata = Some(json!({"source": "reviewer"}));
        let reviewer_write = client.remember(reviewer_rule).await?;

        let mut planner_context_opts = GetContextOptions::default();
        planner_context_opts.run_id = Some(run_id.clone());
        planner_context_opts.query = Some("What should planner know before triage?".into());
        planner_context_opts.agent_id = Some("planner".into());
        planner_context_opts.entry_types = vec!["fact".into(), "lesson".into(), "rule".into()];
        planner_context_opts.mode = Some("sections".into());
        planner_context_opts.sections = vec!["facts".into(), "rules".into()];
        let planner_context = client.get_context(planner_context_opts).await?;

        let mut reviewer_context_opts = GetContextOptions::default();
        reviewer_context_opts.run_id = Some(run_id.clone());
        reviewer_context_opts.query = Some("What should reviewer know before approval?".into());
        reviewer_context_opts.agent_id = Some("reviewer".into());
        reviewer_context_opts.entry_types = vec!["fact".into(), "lesson".into(), "rule".into(), "trace".into()];
        reviewer_context_opts.mode = Some("sections".into());
        reviewer_context_opts.sections = vec!["facts".into(), "rules".into()];
        let reviewer_context = client.get_context(reviewer_context_opts).await?;

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

        require(planner_write.get("accepted").and_then(|v| v.as_bool()).unwrap_or(false) || planner_write.get("job_id").is_some(), format!("planner remember failed: {planner_write}"))?;
        require(reviewer_write.get("accepted").and_then(|v| v.as_bool()).unwrap_or(false) || reviewer_write.get("job_id").is_some(), format!("reviewer remember failed: {reviewer_write}"))?;
        let agent_count = agents.get("agents").and_then(|v| v.as_array()).map(|items| items.len()).unwrap_or(0);
        require(agent_count == 2, format!("agent registration failed: {agents}"))?;

        let planner_types: Vec<&str> = planner_context.get("sources").and_then(|v| v.as_array()).into_iter().flatten().filter_map(|source| source.get("entry_type").and_then(|v| v.as_str())).collect();
        let reviewer_types: Vec<&str> = reviewer_context.get("sources").and_then(|v| v.as_array()).into_iter().flatten().filter_map(|source| source.get("entry_type").and_then(|v| v.as_str())).collect();
        require(planner_types.iter().any(|t| *t == "fact"), format!("planner should see facts: {planner_context}"))?;
        require(!planner_types.iter().any(|t| *t == "rule"), format!("planner should not see reviewer rule scope: {planner_context}"))?;
        require(reviewer_types.iter().any(|t| *t == "rule"), format!("reviewer should see rules: {reviewer_context}"))?;

        metrics = json!({
            "run_id": run_id,
            "agent_count": agent_count,
            "planner_source_types": planner_types,
            "reviewer_source_types": reviewer_types,
        });
        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)) }
}
