//! MuBit Learn Module End-to-End Smoke Test (Rust).
//!
//! Tests the learn module against a live local MuBit instance.
//! Verifies:
//! 1. LearnConfig + LearnSession creation
//! 2. Lesson seeding via SDK Client
//! 3. enrich_messages() retrieves and injects context
//! 4. inject_context_openai() produces correct output
//! 5. RunManager call counting
//! 6. Session end triggers reflect

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

use _utils::{new_run_id, require};
use mubit_sdk::learn::injection::{extract_query, inject_context_openai};
use mubit_sdk::learn::{LearnConfig, LearnSession};
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_13_learn_module_smoke";
    let started = Instant::now();

    let api_key = std::env::var("MUBIT_API_KEY")
        .map_err(|_| "MUBIT_API_KEY must be set")?;
    let endpoint = std::env::var("MUBIT_ENDPOINT")
        .unwrap_or_else(|_| "http://127.0.0.1:3000".to_string());

    let run_id = new_run_id("learn_smoke");

    let mut passed = true;
    let mut detail = String::new();
    let mut metrics = json!({});

    let scenario = async {
        // -----------------------------------------------------------------
        // Step 1: Seed a lesson via raw HTTP (bootstrap key works with HTTP)
        // -----------------------------------------------------------------
        println!("[INFO] Seeding lesson in run_id={run_id}...");
        let http_client = reqwest::Client::new();
        let seed_resp = http_client
            .post(format!("{endpoint}/v2/control/ingest"))
            .bearer_auth(&api_key)
            .json(&json!({
                "run_id": run_id,
                "agent_id": "learn-smoke",
                "items": [{
                    "item_id": format!("learn-smoke-{}", run_id),
                    "content_type": "text",
                    "text": "Always validate the JWT token expiry before retrying a failed API call",
                    "intent": "lesson",
                    "metadata_json": "{\"lesson_type\":\"success\",\"lesson_scope\":\"session\",\"lesson_importance\":\"high\"}"
                }]
            }))
            .send()
            .await?;

        require(
            seed_resp.status().is_success(),
            format!("ingest should succeed, got status {}", seed_resp.status()),
        )?;
        println!("[OK]   Lesson seeded");

        // Wait for indexing
        tokio::time::sleep(std::time::Duration::from_secs(5)).await;

        // -----------------------------------------------------------------
        // Step 2: Create LearnSession
        // -----------------------------------------------------------------
        println!("[INFO] Creating LearnSession...");
        let config = LearnConfig::default()
            .api_key(&api_key)
            .endpoint(&endpoint)
            .session_id(&run_id)
            .inject_lessons(true)
            .max_token_budget(2048)
            .auto_reflect(false); // don't reflect until we explicitly call end()

        let session = LearnSession::new(config).await;
        require(
            session.session_id() == run_id,
            format!("session_id should be {run_id}, got {}", session.session_id()),
        )?;
        println!("[OK]   LearnSession created, session_id={}", session.session_id());

        // -----------------------------------------------------------------
        // Step 3: Test inject_context_openai (pure function)
        // -----------------------------------------------------------------
        println!("[INFO] Testing inject_context_openai...");
        let messages = vec![
            json!({"role": "system", "content": "You are a helpful assistant."}),
            json!({"role": "user", "content": "What should I do before retrying?"}),
        ];

        let context_block = "### Lessons\n- Always validate JWT before retry";
        let enriched = inject_context_openai(&messages, context_block, "system");

        let sys_content = enriched[0]
            .get("content")
            .and_then(|v| v.as_str())
            .unwrap_or("");
        require(
            sys_content.contains("<memory_context>"),
            "enriched system message should contain <memory_context>",
        )?;
        require(
            sys_content.contains("You are a helpful assistant."),
            "enriched system message should preserve original content",
        )?;
        require(
            sys_content.contains("Always validate JWT"),
            "enriched system message should contain the lesson",
        )?;
        println!("[OK]   inject_context_openai works correctly");

        // -----------------------------------------------------------------
        // Step 4: Test extract_query
        // -----------------------------------------------------------------
        let query = extract_query(&messages, 200);
        require(
            query == "What should I do before retrying?",
            format!("extract_query should return user message, got: {query}"),
        )?;
        println!("[OK]   extract_query works correctly");

        // -----------------------------------------------------------------
        // Step 5: Test enrich_messages (live server call)
        // -----------------------------------------------------------------
        println!("[INFO] Testing enrich_messages (live context retrieval)...");
        let enriched_live = session.enrich_messages(&messages).await;

        let live_sys = enriched_live[0]
            .get("content")
            .and_then(|v| v.as_str())
            .unwrap_or("");

        if live_sys.contains("<memory_context>") {
            println!("[OK]   enrich_messages injected live context from server");
            require(
                live_sys.contains("You are a helpful assistant."),
                "enrich_messages should preserve original system message",
            )?;
            metrics["context_injected"] = json!(true);
        } else {
            println!("[WARN] enrich_messages returned no context (lesson may not be indexed yet)");
            println!("       System content: {}...", &live_sys[..live_sys.len().min(200)]);
            metrics["context_injected"] = json!(false);
        }

        // -----------------------------------------------------------------
        // Step 6: Test record() and call counting
        // -----------------------------------------------------------------
        println!("[INFO] Testing record + call counting...");
        session.record("test response", "gpt-4o-mini", 100.0).await;
        session.record("another response", "gpt-4o-mini", 200.0).await;

        require(
            session.call_count() == 2,
            format!("call_count should be 2, got {}", session.call_count()),
        )?;
        println!("[OK]   call_count is {}", session.call_count());

        // -----------------------------------------------------------------
        // Step 7: Test session.end() (triggers reflect)
        // -----------------------------------------------------------------
        println!("[INFO] Testing session.end()...");
        session.end().await;
        println!("[OK]   session.end() completed");

        metrics["session_id"] = json!(run_id);
        metrics["call_count"] = json!(2);

        Ok::<(), Box<dyn Error>>(())
    };

    match scenario.await {
        Ok(()) => {
            detail = "validated learn module: config, injection, enrichment, run lifecycle".into();
        }
        Err(err) => {
            passed = false;
            detail = format!("scenario failed: {err}");
        }
    }

    let elapsed = started.elapsed().as_secs_f64();

    let status = if passed { "PASS" } else { "FAIL" };
    println!();
    println!("[{status}] {name} {elapsed:.2}s  {detail}");
    println!("metrics: {metrics}");
    println!("cleanup: PASS");

    if passed { Ok(()) } else { Err(detail.into()) }
}
