mubit-sdk 0.13.1

//! MuBit Learn Run Manager.
//!
//! Manages session lifecycle: call counting, background reflection.

use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Mutex;
use std::time::Duration;

use crate::learn::LearnConfig;

pub(crate) struct RunManagerInner {
    pub session_id: String,
    pub config: LearnConfig,
    call_count: AtomicU64,
    ended: AtomicBool,
    /// Entry IDs recalled for the most recent call, so feedback() can credit them.
    last_recalled_ids: Mutex<Vec<String>>,
    http_client: reqwest::Client,
}

impl RunManagerInner {
    pub fn new(config: LearnConfig) -> Self {
        let session_id = config.session_id.clone().unwrap_or_else(|| {
            uuid::Uuid::new_v4()
                .to_string()
                .replace('-', "")
                .chars()
                .take(16)
                .collect()
        });

        Self {
            session_id,
            config,
            call_count: AtomicU64::new(0),
            ended: AtomicBool::new(false),
            last_recalled_ids: Mutex::new(Vec::new()),
            // Generous global fallback; each request sets its own (shorter)
            // per-request timeout below (context = configurable hot path).
            http_client: reqwest::Client::builder()
                .timeout(Duration::from_secs(30))
                .build()
                .unwrap_or_default(),
        }
    }

    /// Stash the entry IDs recalled for the latest call.
    pub fn set_recalled_ids(&self, ids: Vec<String>) {
        if let Ok(mut guard) = self.last_recalled_ids.lock() {
            *guard = ids.into_iter().filter(|s| !s.is_empty()).collect();
        }
    }

    /// The entry IDs recalled for the most recent call.
    pub fn last_recalled_ids(&self) -> Vec<String> {
        self.last_recalled_ids
            .lock()
            .map(|g| g.clone())
            .unwrap_or_default()
    }

    pub fn increment(&self) {
        let count = self.call_count.fetch_add(1, Ordering::Relaxed) + 1;
        if let Some(n) = self.config.reflect_after_n_calls {
            if n > 0 && count % n == 0 {
                self.background_reflect();
            }
        }
    }

    pub fn call_count(&self) -> u64 {
        self.call_count.load(Ordering::Relaxed)
    }

    pub async fn end(&self) {
        if self.ended.swap(true, Ordering::SeqCst) {
            return;
        }
        if self.config.auto_reflect {
            let _ = self.do_reflect().await;
        }
    }

    pub fn background_reflect(&self) {
        let endpoint = self.config.endpoint.clone();
        let api_key = self.config.api_key.clone();
        let session_id = self.session_id.clone();
        let client = self.http_client.clone();

        tokio::spawn(async move {
            let _ = Self::reflect_http(&client, &endpoint, &api_key, &session_id).await;
        });
    }

    async fn do_reflect(&self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        Self::reflect_http(
            &self.http_client,
            &self.config.endpoint,
            &self.config.api_key,
            &self.session_id,
        )
        .await
    }

    async fn reflect_http(
        client: &reqwest::Client,
        endpoint: &str,
        api_key: &str,
        session_id: &str,
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        if api_key.is_empty() {
            return Ok(());
        }
        let url = format!("{}/v2/control/reflect", endpoint.trim_end_matches('/'));
        client
            .post(&url)
            .bearer_auth(api_key)
            .timeout(Duration::from_secs(10))
            .json(&serde_json::json!({"run_id": session_id}))
            .send()
            .await?;
        Ok(())
    }

    /// Fetch the context block plus the entry IDs of its source evidence, so a
    /// later outcome can be attributed back to them. One call serves both
    /// injection and feedback attribution. Uses the configurable hot-path
    /// timeout. Returns `(String::new(), vec![])` on a non-success status.
    pub async fn get_context_with_ids_http(
        &self,
        query: &str,
    ) -> Result<(String, Vec<String>), Box<dyn std::error::Error + Send + Sync>> {
        if self.config.api_key.is_empty() {
            return Ok((String::new(), Vec::new()));
        }

        let url = format!(
            "{}/v2/control/context",
            self.config.endpoint.trim_end_matches('/')
        );

        let mut payload = serde_json::json!({
            "run_id": self.session_id,
            "query": query,
            "format": "structured",
            "max_token_budget": self.config.max_token_budget,
        });

        if !self.config.entry_types.is_empty() {
            payload["entry_types"] = serde_json::json!(self.config.entry_types);
        }
        if !self.config.context_sections.is_empty() {
            payload["sections"] = serde_json::json!(self.config.context_sections);
        }

        let resp = self
            .http_client
            .post(&url)
            .bearer_auth(&self.config.api_key)
            .timeout(Duration::from_secs_f64(self.config.context_fetch_timeout.max(0.1)))
            .json(&payload)
            .send()
            .await?;

        if !resp.status().is_success() {
            return Ok((String::new(), Vec::new()));
        }

        let body: serde_json::Value = resp.json().await?;
        let block = body
            .get("context_block")
            .and_then(|v| v.as_str())
            .unwrap_or("")
            .to_string();
        let ids = body
            .get("sources")
            .and_then(|v| v.as_array())
            .map(|arr| {
                arr.iter()
                    .filter_map(|s| s.get("id").and_then(|i| i.as_str()).map(String::from))
                    .collect::<Vec<_>>()
            })
            .unwrap_or_default();
        Ok((block, ids))
    }

    /// Record an outcome that credits recalled entries. With an empty
    /// reference_id and populated `entry_ids`, the server applies per-entry
    /// reinforcement. Fire-and-forget.
    pub async fn record_outcome_http(
        &self,
        outcome: &str,
        signal: f32,
        entry_ids: &[String],
        verified_in_production: bool,
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        if self.config.api_key.is_empty() {
            return Ok(());
        }
        let url = format!(
            "{}/v2/control/outcome",
            self.config.endpoint.trim_end_matches('/')
        );
        let payload = outcome_payload(
            &self.session_id,
            &self.config.agent_id,
            outcome,
            signal,
            entry_ids,
            verified_in_production,
        );
        self.http_client
            .post(&url)
            .bearer_auth(&self.config.api_key)
            .timeout(Duration::from_secs(8))
            .json(&payload)
            .send()
            .await?;
        Ok(())
    }
}

/// Build the `/v2/control/outcome` request body. An empty `reference_id` with
/// populated `entry_ids` is the run-level multi-entry attribution form the
/// server applies per entry. Pure (no I/O) so it can be unit-tested.
pub(crate) fn outcome_payload(
    session_id: &str,
    agent_id: &str,
    outcome: &str,
    signal: f32,
    entry_ids: &[String],
    verified_in_production: bool,
) -> serde_json::Value {
    let mut payload = serde_json::json!({
        "run_id": session_id,
        "reference_id": "",
        "outcome": outcome,
        "signal": signal,
        "agent_id": agent_id,
    });
    if !entry_ids.is_empty() {
        payload["entry_ids"] = serde_json::json!(entry_ids);
    }
    if verified_in_production {
        payload["verified_in_production"] = serde_json::json!(true);
    }
    payload
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::learn::LearnConfig;

    #[test]
    fn recalled_ids_are_stashed_and_filtered() {
        let inner = RunManagerInner::new(LearnConfig::default().session_id("s"));
        assert!(inner.last_recalled_ids().is_empty());
        inner.set_recalled_ids(vec!["a".into(), "".into(), "b".into()]);
        assert_eq!(inner.last_recalled_ids(), vec!["a".to_string(), "b".to_string()]);
    }

    #[test]
    fn outcome_payload_shape() {
        let p = outcome_payload("run1", "ag", "success", 1.0, &["e1".into(), "e2".into()], true);
        assert_eq!(p["run_id"], "run1");
        assert_eq!(p["reference_id"], "");
        assert_eq!(p["outcome"], "success");
        assert_eq!(p["signal"], 1.0);
        assert_eq!(p["agent_id"], "ag");
        assert_eq!(p["entry_ids"], serde_json::json!(["e1", "e2"]));
        assert_eq!(p["verified_in_production"], true);
    }

    #[test]
    fn outcome_payload_omits_optional_fields() {
        let p = outcome_payload("run1", "ag", "failure", -0.5, &[], false);
        assert!(p.get("entry_ids").is_none());
        assert!(p.get("verified_in_production").is_none());
        assert_eq!(p["outcome"], "failure");
    }
}