ainl-memory 0.1.13

//! Mission subgraph query helpers (graph-native mission substrate).

use crate::edge_labels::{HAS_ASSERTION, HAS_FEATURE, HANDED_OFF_BY, PROGRESS_FOR};
use crate::node::AinlMemoryNode;
use crate::snapshot::SnapshotEdge;
use crate::store::GraphStore;
use ainl_contracts::{Assertion, Feature, Handoff, Mission, MissionState};
use serde::{Deserialize, Serialize};
use std::collections::HashSet;
use uuid::Uuid;

/// Lightweight mission row for orchestrator prompts and API list views.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct MissionSummary {
    pub mission_id: String,
    pub state: MissionState,
    pub objective_preview: String,
    pub milestone_count: usize,
}

/// Mission node plus connected features, assertions, handoffs, and edge rows.
#[derive(Debug, Clone)]
pub struct MissionSubgraph {
    pub mission: Mission,
    pub mission_node_id: Uuid,
    pub features: Vec<Feature>,
    pub assertions: Vec<Assertion>,
    pub handoffs: Vec<Handoff>,
    pub edges: Vec<SnapshotEdge>,
}

fn node_matches_agent(node: &AinlMemoryNode, agent_id: &str) -> bool {
    node.agent_id.is_empty() || node.agent_id == agent_id
}

fn objective_preview(objective_md: &str, max_chars: usize) -> String {
    let flat: String = objective_md
        .lines()
        .map(str::trim)
        .filter(|l| !l.is_empty())
        .collect::<Vec<_>>()
        .join(" ");
    if flat.chars().count() <= max_chars {
        flat
    } else {
        flat.chars().take(max_chars).collect::<String>() + "…"
    }
}

fn find_mission_node(
    store: &dyn GraphStore,
    agent_id: &str,
    mission_id: &str,
) -> Result<AinlMemoryNode, String> {
    for node in store.find_by_type("mission")? {
        if !node_matches_agent(&node, agent_id) {
            continue;
        }
        if node
            .mission()
            .is_some_and(|m| m.mission_id.as_str() == mission_id)
        {
            return Ok(node);
        }
    }
    Err(format!("mission not found: {mission_id}"))
}

/// Active missions for this agent (not `Completed` / `Cancelled`).
pub fn find_active_missions(
    store: &dyn GraphStore,
    agent_id: &str,
) -> Result<Vec<MissionSummary>, String> {
    let mut out = Vec::new();
    for node in store.find_by_type("mission")? {
        if !node_matches_agent(&node, agent_id) {
            continue;
        }
        let Some(mission) = node.mission() else {
            continue;
        };
        if matches!(
            mission.state,
            MissionState::Completed | MissionState::Cancelled
        ) {
            continue;
        }
        out.push(MissionSummary {
            mission_id: mission.mission_id.as_str().to_string(),
            state: mission.state,
            objective_preview: objective_preview(&mission.objective_md, 160),
            milestone_count: mission.milestone_ids.len(),
        });
    }
    out.sort_by(|a, b| a.mission_id.cmp(&b.mission_id));
    Ok(out)
}

/// Load mission record + connected DAG nodes in one sweep.
pub fn mission_subgraph(
    store: &dyn GraphStore,
    agent_id: &str,
    mission_id: &str,
) -> Result<MissionSubgraph, String> {
    let mission_node = find_mission_node(store, agent_id, mission_id)?;
    let mission = mission_node
        .mission()
        .cloned()
        .ok_or_else(|| format!("invalid mission node payload for {mission_id}"))?;

    let feature_nodes = store.walk_edges(mission_node.id, HAS_FEATURE)?;
    let assertion_nodes = store.walk_edges(mission_node.id, HAS_ASSERTION)?;

    let features = feature_nodes
        .iter()
        .filter_map(|n| n.feature().cloned())
        .collect::<Vec<_>>();
    let assertions = assertion_nodes
        .iter()
        .filter_map(|n| n.assertion().cloned())
        .collect::<Vec<_>>();

    let mut handoffs = Vec::new();
    let mut handoff_ids = HashSet::new();
    for feat_node in &feature_nodes {
        for handoff_node in store.walk_edges(feat_node.id, HANDED_OFF_BY)? {
            if handoff_ids.insert(handoff_node.id) {
                if let Some(h) = handoff_node.handoff() {
                    handoffs.push(h.clone());
                }
            }
        }
    }

    Ok(MissionSubgraph {
        mission,
        mission_node_id: mission_node.id,
        features,
        assertions,
        handoffs,
        edges: Vec::new(),
    })
}

/// Count `Episode --PROGRESS_FOR--> Feature` edges for features in this mission since `since_ts`.
pub fn count_progress_events_since(
    store: &dyn GraphStore,
    agent_id: &str,
    mission_id: &str,
    since_ts: i64,
) -> Result<usize, String> {
    let subgraph = mission_subgraph(store, agent_id, mission_id)?;
    if subgraph.features.is_empty() {
        return Ok(0);
    }
    let feature_ids: HashSet<String> = subgraph
        .features
        .iter()
        .map(|f| f.feature_id.as_str().to_string())
        .collect();

    let mut count = 0usize;
    for node in store.query_episodes_since(since_ts, 10_000)? {
        for target in store.walk_edges(node.id, PROGRESS_FOR)? {
            if target
                .feature()
                .is_some_and(|f| feature_ids.contains(f.feature_id.as_str()))
            {
                count += 1;
            }
        }
    }
    Ok(count)
}

/// Features filtered by milestone label within a mission subgraph.
pub fn find_features_by_milestone(
    store: &dyn GraphStore,
    agent_id: &str,
    mission_id: &str,
    milestone: &str,
) -> Result<Vec<Feature>, String> {
    Ok(mission_subgraph(store, agent_id, mission_id)?
        .features
        .into_iter()
        .filter(|f| f.milestone.as_deref() == Some(milestone))
        .collect())
}

/// Assertions filtered by milestone label within a mission subgraph.
pub fn find_assertions_for_milestone(
    store: &dyn GraphStore,
    agent_id: &str,
    mission_id: &str,
    milestone: &str,
) -> Result<Vec<Assertion>, String> {
    Ok(mission_subgraph(store, agent_id, mission_id)?
        .assertions
        .into_iter()
        .filter(|a| a.milestone.as_deref() == Some(milestone))
        .collect())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::node::{AinlMemoryNode, AinlNodeType, MemoryCategory};
    use crate::GraphMemory;
    use ainl_contracts::{FeatureId, MissionCapabilityFlags, MissionId};
    use chrono::Utc;

    fn write_mission(memory: &GraphMemory, agent_id: &str, mid: &str) -> Uuid {
        let mission = Mission {
            mission_id: MissionId(mid.into()),
            objective_md: "Build the mission substrate".into(),
            state: MissionState::Running,
            milestone_ids: vec!["m1".into()],
            mission_root: None,
            created_at: Utc::now(),
            last_orchestrator_turn_at: None,
            capability_flags: MissionCapabilityFlags::default(),
        };
        let mut node = AinlMemoryNode {
            id: Uuid::new_v4(),
            memory_category: MemoryCategory::Mission,
            importance_score: 0.5,
            agent_id: agent_id.to_string(),
            project_id: None,
            node_type: AinlNodeType::Mission { mission },
            edges: Vec::new(),
            plugin_data: None,
        };
        let id = node.id;
        memory.write_node(&node).expect("write mission");
        id
    }

    fn write_feature(memory: &GraphMemory, agent_id: &str, fid: &str, milestone: Option<&str>) -> AinlMemoryNode {
        let feature = Feature {
            feature_id: FeatureId(fid.into()),
            description: format!("feature {fid}"),
            status: ainl_contracts::FeatureStatus::Pending,
            milestone: milestone.map(str::to_string),
            skill_name: None,
            touches_files: vec!["src/lib.rs".into()],
            preconditions: vec![],
            expected_behavior: vec![],
            verification_steps: vec![],
            fulfills: vec![],
            snapshot: None,
        };
        let node = AinlMemoryNode {
            id: Uuid::new_v4(),
            memory_category: MemoryCategory::Feature,
            importance_score: 0.5,
            agent_id: agent_id.to_string(),
            project_id: None,
            node_type: AinlNodeType::Feature { feature },
            edges: Vec::new(),
            plugin_data: None,
        };
        memory.write_node(&node).expect("write feature");
        node
    }

    #[test]
    fn mission_subgraph_and_active_missions_roundtrip() {
        let dir = tempfile::tempdir().expect("tempdir");
        let db = dir.path().join("mission_query.db");
        let memory = GraphMemory::new(&db).expect("graph memory");
        let agent = "agent-mission-q";

        let mission_id = write_mission(&memory, agent, "mission-1");
        let feat = write_feature(&memory, agent, "f1", Some("m1"));
        let mut mission_node = memory
            .store()
            .read_node(mission_id)
            .expect("read")
            .expect("node");
        mission_node.add_edge(feat.id, HAS_FEATURE);
        memory.write_node(&mission_node).expect("edge mission->feature");

        let store = memory.store();
        let active = find_active_missions(store, agent).expect("active");
        assert_eq!(active.len(), 1);
        assert_eq!(active[0].mission_id, "mission-1");

        let subgraph = mission_subgraph(store, agent, "mission-1").expect("subgraph");
        assert_eq!(subgraph.features.len(), 1);
        assert_eq!(subgraph.features[0].feature_id.as_str(), "f1");

        let by_m = find_features_by_milestone(store, agent, "mission-1", "m1").expect("by milestone");
        assert_eq!(by_m.len(), 1);
    }

    #[test]
    fn mission_migration_idempotent_and_fts_indexed() {
        let dir = tempfile::tempdir().expect("tempdir");
        let db = dir.path().join("mission_migrate.db");
        let memory = GraphMemory::new(&db).expect("open v1");
        let mission_node_id = write_mission(&memory, "agent-a", "m-fts");
        let _feat = write_feature(&memory, "agent-a", "f-fts", Some("m1"));
        drop(memory);

        let memory2 = GraphMemory::new(&db).expect("reopen migrates idempotently");
        let hits = memory2
            .search_all_nodes_fts("agent-a", "substrate", None, 10)
            .expect("fts");
        assert!(
            hits.iter().any(|n| n.id == mission_node_id),
            "mission nodes should be FTS-indexed"
        );
        let _ = GraphMemory::new(&db).expect("third open is idempotent");
    }

    #[test]
    fn mission_subgraph_query_under_budget_on_fixture() {
        use std::time::Instant;
        let dir = tempfile::tempdir().expect("tempdir");
        let db = dir.path().join("mission_perf.db");
        let memory = GraphMemory::new(&db).expect("open");
        let agent = "perf-agent";
        let _mid = write_mission(&memory, agent, "perf-m1");
        for i in 0..50 {
            write_feature(&memory, agent, &format!("f{i}"), Some("m1"));
        }
        let store = memory.store();
        let start = Instant::now();
        for _ in 0..20 {
            let _ = mission_subgraph(store, agent, "perf-m1").expect("subgraph");
        }
        let elapsed_ms = start.elapsed().as_millis();
        assert!(
            elapsed_ms < 500,
            "20 mission_subgraph queries took {elapsed_ms}ms (budget 500ms on ~50-feature fixture)"
        );
    }
}