post-cortex-daemon 0.4.0

// Copyright (c) 2025, 2026 Julius ML
// MIT License

//! Tool 7: assemble_context — graph-aware retrieval: query-term matching over
//! recent updates, entity-graph traversal and impact analysis.

use crate::daemon::assembly_scope;
use crate::daemon::coerce::{CoercionError, coerce_and_validate};
use crate::daemon::validate::validate_session_id;
use post_cortex_mcp::MCPToolResult;
use post_cortex_memory::ConversationMemorySystem;
use post_cortex_memory::context_assembly;
use rmcp::{
    handler::server::wrapper::Parameters,
    model::{CallToolResult, ErrorData as McpError},
};
use std::sync::Arc;
use uuid::Uuid;

use super::AssembleContextRequest;
use super::mcp_result_to_call_result;

pub(super) async fn handle(
    memory_system: &Arc<ConversationMemorySystem>,
    params: Parameters<serde_json::Value>,
) -> Result<CallToolResult, McpError> {
    let req: AssembleContextRequest = coerce_and_validate(params.0).map_err(|e| {
        if e.message.contains("query") {
            e.clone()
                .with_parameter_path("query".to_string())
                .with_expected_type("non-empty query/topic string")
                .with_hint("Provide the question or topic to assemble context for")
                .to_mcp_error()
        } else {
            e.to_mcp_error()
        }
    })?;

    if req.query.trim().is_empty() {
        return Err(McpError::invalid_params(
            "query cannot be empty".to_string(),
            Some(serde_json::Value::String("query".to_string())),
        ));
    }

    if req.session_id.is_none() && req.workspace_id.is_none() {
        return Err(McpError::invalid_params(
            "either session_id or workspace_id is required".to_string(),
            None,
        ));
    }

    let token_budget = req
        .token_budget
        .filter(|b| *b > 0)
        .map(|b| b as usize)
        .unwrap_or(4000);

    let (updates, graph) = if let Some(ws_raw) = req.workspace_id.as_ref().filter(|s| !s.is_empty())
    {
        let ws_uuid = Uuid::parse_str(ws_raw).map_err(|_| {
            CoercionError::new(
                "Invalid UUID",
                std::io::Error::new(std::io::ErrorKind::InvalidInput, "Invalid UUID format"),
                Some(serde_json::Value::String(ws_raw.clone())),
            )
            .with_parameter_path("workspace_id".to_string())
            .with_expected_type("Valid UUID string")
            .to_mcp_error()
        })?;

        let workspace = memory_system
            .workspace_manager
            .get_workspace(&ws_uuid)
            .ok_or_else(|| {
                McpError::invalid_params(
                    format!("Workspace {} not found", ws_uuid),
                    Some(serde_json::Value::String("workspace_id".to_string())),
                )
            })?;

        let mut all_updates = Vec::new();
        let mut graphs = Vec::new();

        for (ws_session_id, _role) in workspace.get_all_sessions() {
            if let Ok(session_arc) = memory_system.get_session(ws_session_id).await {
                let session = session_arc.load();
                assembly_scope::session_updates(&session, &mut all_updates);
                graphs.push((ws_session_id, Arc::clone(&session.entity_graph)));
            }
        }

        let merged = assembly_scope::WORKSPACE_GRAPHS.merged(ws_uuid, graphs);
        (all_updates, merged)
    } else {
        let sid_raw = req.session_id.as_ref().unwrap();
        let session_id = validate_session_id(sid_raw).map_err(|e| e.to_mcp_error())?;

        let session_arc = memory_system
            .get_session(session_id)
            .await
            .map_err(|e| McpError::internal_error(format!("Session not found: {}", e), None))?;
        let session = session_arc.load();
        let mut updates = Vec::new();
        assembly_scope::session_updates(&session, &mut updates);
        (updates, Arc::clone(&session.entity_graph))
    };

    let assembled = context_assembly::assemble_context(&req.query, &graph, &updates, token_budget);
    let formatted_text = context_assembly::format_for_llm(&assembled);

    let payload = serde_json::json!({
        "query": req.query,
        "token_budget": token_budget,
        "total_tokens": assembled.total_tokens,
        "items": assembled.items,
        "entity_context": assembled.entity_context,
        "impact": assembled.impact,
        "formatted_text": formatted_text,
    });

    // Apply optional JSONPath select filter
    let payload = if req.select.is_some() || req.select_map.is_some() {
        use post_cortex_memory::jsonpath;

        let select_string = req.select.as_deref();
        let select_map = req.select_map.as_ref();

        match jsonpath::apply_select(&payload, select_string, select_map) {
            Ok(filtered) => filtered,
            Err(e) => {
                return Err(McpError::invalid_params(
                    format!("Invalid JSONPath expression: {}", e.message),
                    Some(serde_json::json!({
                        "offset": e.byte_offset,
                        "hint": "Allowed: $, .field, [n], [-n], [a:b], [*], [0,3,7]. \
                                 Forbidden: .., [?(...)], functions, [a:b:c]"
                    })),
                ));
            }
        }
    } else {
        payload
    };

    Ok(mcp_result_to_call_result(MCPToolResult::success(
        format!(
            "Assembled {} items ({} tokens) for query: {}",
            assembled.items.len(),
            assembled.total_tokens,
            req.query
        ),
        Some(payload),
    )))
}