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

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

    let mut passed = true;
    let mut detail = "validated context budgeting and sectioned disclosure".to_string();
    let mut metrics = json!({});

    let scenario = async {
        let mut rule = RememberOptions::new("Always lead with policy facts before escalation notes.");
        rule.run_id = Some(run_id.clone());
        rule.agent_id = Some("planner".to_string());
        rule.intent = Some("rule".to_string());
        rule.metadata = Some(json!({"source": "public-example", "section": "rules"}));
        rule.importance = Some("critical".to_string());
        client.remember(rule).await?;

        let mut fact = RememberOptions::new("Claim CLM-4242 has a policy limit of 5000 GBP and a deductible of 250 GBP.");
        fact.run_id = Some(run_id.clone());
        fact.agent_id = Some("planner".to_string());
        fact.intent = Some("fact".to_string());
        fact.metadata = Some(json!({"source": "public-example", "section": "facts"}));
        fact.importance = Some("high".to_string());
        client.remember(fact).await?;

        let mut trace = RememberOptions::new("Planner noted that the claimant is waiting on a same-day callback.");
        trace.run_id = Some(run_id.clone());
        trace.agent_id = Some("planner".to_string());
        trace.intent = Some("trace".to_string());
        trace.metadata = Some(json!({"source": "public-example", "section": "traces"}));
        trace.importance = Some("medium".to_string());
        client.remember(trace).await?;

        let mut mental = RememberOptions::new("Claim CLM-4242 overview: 5000 GBP policy limit, 250 GBP deductible, claimant awaiting same-day callback. Priority: high.");
        mental.run_id = Some(run_id.clone());
        mental.agent_id = Some("planner".to_string());
        mental.intent = Some("mental_model".to_string());
        mental.metadata = Some(json!({"source": "public-example", "section": "mental_models"}));
        mental.importance = Some("critical".to_string());
        client.remember(mental).await?;

        let mut summary_opts = GetContextOptions::default();
        summary_opts.run_id = Some(run_id.clone());
        summary_opts.query = Some("Prepare a claims summary.".to_string());
        summary_opts.mode = Some("summary".to_string());
        summary_opts.entry_types = vec!["rule".to_string(), "fact".to_string(), "trace".to_string()];
        summary_opts.max_token_budget = Some(220);
        summary_opts.limit = Some(6);
        let summary = client.get_context(summary_opts).await?;

        let mut sections_opts = GetContextOptions::default();
        sections_opts.run_id = Some(run_id.clone());
        sections_opts.query = Some("Prepare a claims summary.".to_string());
        sections_opts.mode = Some("sections".to_string());
        sections_opts.sections = vec!["mental_models".to_string(), "rules".to_string(), "facts".to_string()];
        sections_opts.entry_types = vec!["rule".to_string(), "fact".to_string(), "trace".to_string()];
        sections_opts.max_token_budget = Some(220);
        sections_opts.limit = Some(6);
        let sections = client.get_context(sections_opts).await?;

        require(summary.get("context_block").and_then(|v| v.as_str()).unwrap_or("").is_empty(), format!("summary mode should suppress context_block: {summary}"))?;
        let summary_sections = summary.get("section_summaries").and_then(|v| v.as_array()).ok_or_else(|| boxed_error(format!("summary mode missing section summaries: {summary}")))?;
        let section_sources = sections.get("sources").and_then(|v| v.as_array()).ok_or_else(|| boxed_error(format!("sections mode missing sources: {sections}")))?;
        require(!summary_sections.is_empty(), format!("summary mode missing section summaries: {summary}"))?;
        require(!section_sources.is_empty(), format!("sections mode missing sources: {sections}"))?;
        require(sections.get("budget_used").and_then(|v| v.as_u64()).unwrap_or(0) >= 1, format!("budget_used missing: {sections}"))?;
        require(sections.get("budget_remaining").and_then(|v| v.as_u64()).unwrap_or(0) <= 220, format!("budget_remaining missing: {sections}"))?;
        require(section_sources.iter().any(|s| s.get("entry_type").and_then(|v| v.as_str()) == Some("rule")), format!("expected rule source in sections mode: {sections}"))?;

        metrics = json!({
            "run_id": run_id,
            "summary_section_count": summary_sections.len(),
            "section_source_count": section_sources.len(),
            "budget_used": sections.get("budget_used").cloned().unwrap_or(serde_json::Value::Null),
            "budget_remaining": sections.get("budget_remaining").cloned().unwrap_or(serde_json::Value::Null),
        });
        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)) }
}
