zesdex/apps/infrastructure/src/tools/parallel_delegate.rs

//! Parallel delegation tool — splits a large task into sub-tasks that run
//! concurrently across multiple subagents, then synthesises results.
//!
//! Flow: receive task description + optional breakdown → LLM analyses the
//! task and splits it into parallel directives → spawn subagents for each →
//! collect results → return consolidated output.

use anyhow::Result;
use serde_json::{json, Value};
use tracing::{debug, info, instrument, warn};

use zesdex_domain::cms::{AppConfigRepository, SettingsRepository};
use zesdex_domain::core::Store;

use crate::persistence::{JsonAppConfigRepository, JsonSettingsRepository};
use crate::subagent::context::SubagentContext;
use crate::subagent::division::AccessTier;
use crate::subagent::spawn::spawn_subagent;
use crate::tools::{Tool, ToolCtx};

/// Delegate a large task to multiple subagents running in parallel.
///
/// Flow:
///   1. If `directives` are provided explicitly, use them directly.
///   2. Otherwise, use LLM to analyse the task and split it into directives.
///   3. Spawn a subagent for each directive concurrently.
///   4. Join all results and consolidate into a single response.
pub struct ParallelDelegate;

impl Tool for ParallelDelegate {
    fn name(&self) -> &'static str {
        "parallel_delegate"
    }

    fn description(&self) -> &'static str {
        "Delegate a large task to multiple subagents running in parallel for 3x faster completion"
    }

    fn parameters(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "task": {
                    "type": "string",
                    "description": "The task to be split and delegated to parallel agents"
                },
                "directives": {
                    "type": "array",
                    "items": {
                        "type": "object",
                        "properties": {
                            "directive": {"type": "string", "description": "Directive for one parallel agent"},
                            "access": {"type": "string", "enum": ["read", "write", "full"], "description": "Access tier for this agent (default: write)"}
                        },
                        "required": ["directive"]
                    },
                    "description": "Optional explicit list of parallel directives (if not provided, the LLM will auto-split the task)"
                },
                "synthesize": {
                    "type": "boolean",
                    "description": "Whether to synthesize results into a unified response (default: true)",
                    "default": true
                },
                "max_parallel": {
                    "type": "integer",
                    "description": "Maximum number of parallel agents (default: 3, max: 8)",
                    "default": 3
                }
            },
            "required": ["task"]
        })
    }

    #[instrument(skip(self, ctx, args))]
    fn run(&self, ctx: &ToolCtx, args: &Value) -> Result<String> {
        let task = crate::tools::arg_str(args, "task")?;
        let synthesize = args
            .get("synthesize")
            .and_then(|v| v.as_bool())
            .unwrap_or(true);
        let max_parallel = args
            .get("max_parallel")
            .and_then(|v| v.as_u64())
            .unwrap_or(3)
            .min(8) as usize;

        // Load LLM credentials
        let store = Store::new();
        let settings = JsonSettingsRepository::new()
            .load(&store.base_dir)
            .unwrap_or_default();
        let app_config = JsonAppConfigRepository::new()
            .load(&store.base_dir)
            .unwrap_or_default();

        let (provider, model) =
            crate::subagent::provider::resolve_subagent_provider(&settings, &app_config);

        let base_url = app_config
            .providers
            .get(&provider)
            .map(|p| p.api_base.clone())
            .unwrap_or_else(|| "https://opencode.ai/zen/v1".to_string());

        let api_key = crate::llm::provider::resolve_api_key(&settings, &app_config);

        // Determine directives
        let directives: Vec<(String, AccessTier)> = if let Some(dirs) =
            args.get("directives").and_then(|v| v.as_array())
        {
            // Explicit directives provided
            dirs.iter()
                .filter_map(|d| {
                    let directive = d.get("directive").and_then(|v| v.as_str())?;
                    let access_str = d
                        .get("access")
                        .and_then(|v| v.as_str())
                        .unwrap_or("write");
                    let access = match access_str {
                        "read" => AccessTier::Read,
                        "full" => AccessTier::Full,
                        _ => AccessTier::Write,
                    };
                    Some((directive.to_string(), access))
                })
                .collect()
        } else {
            // Auto-split using LLM
            let rt = tokio::runtime::Runtime::new()?;
            let directives = rt.block_on(auto_split_task(
                &task,
                max_parallel,
                &base_url,
                &api_key,
                &model,
            ))?;
            directives
        };

        if directives.is_empty() {
            anyhow::bail!("no directives could be derived for the task");
        }

        let directive_count = directives.len();
        info!(
            task = %task,
            directive_count,
            max_parallel,
            "parallel delegation: starting subagents"
        );

        // Spawn agents in parallel
        let mut handles = Vec::new();
        for (i, (directive, access)) in directives.iter().enumerate() {
            let subagent_ctx = SubagentContext::new(
                directive.clone(),
                ctx.clone(),
                format!("{access:?}"),
                base_url.clone(),
                api_key.clone(),
                model.clone(),
            );

            debug!(agent_index = i, access = ?access, "spawning parallel agent");
            let handle = spawn_subagent(
                subagent_ctx,
                directive.clone(),
                *access,
                ctx.clone(),
            );
            handles.push((i, handle));
        }

        // Join all results
        let mut results: Vec<(usize, String, String)> = Vec::new();
        for (i, handle) in handles {
            match handle.join() {
                Ok(Ok(output)) => {
                    info!(agent_index = i, "parallel agent completed");
                    results.push((i, directives[i].0.clone(), output));
                }
                Ok(Err(e)) => {
                    warn!(agent_index = i, error = %e, "parallel agent failed");
                    results.push((
                        i,
                        directives[i].0.clone(),
                        format!("[ERROR] {e}"),
                    ));
                }
                Err(e) => {
                    warn!(agent_index = i, error = ?e, "parallel agent panicked");
                    results.push((
                        i,
                        directives[i].0.clone(),
                        "[ERROR] Agent panicked".to_string(),
                    ));
                }
            }
        }

        // Consolidate results
        if synthesize && results.len() > 1 {
            let consolidated = consolidate_results(&results, &base_url, &api_key, &model)?;
            Ok(format!(
                "## Parallel Delegation Complete\n\n**Task:** {task}\n**Parallel agents:** {}\n\n{}",
                results.len(),
                consolidated
            ))
        } else {
            let mut output = format!(
                "## Parallel Delegation Complete\n\n**Task:** {task}\n**Parallel agents:** {}\n\n",
                results.len()
            );
            for (i, directive, result) in &results {
                output.push_str(&format!("---\n### Agent {}: {}\n\n{}\n", i, directive, result));
            }
            Ok(output)
        }
    }
}

/// Use LLM to analyse a task and split it into parallel directives.
async fn auto_split_task(
    task: &str,
    max_parallel: usize,
    base_url: &str,
    api_key: &str,
    model: &str,
) -> Result<Vec<(String, AccessTier)>> {
    let client = crate::llm::provider::LlmClient::new(
        api_key.to_string(),
        model.to_string(),
        Some(base_url.to_string()),
    );

    let sys_msg = zesdex_domain::core::ChatMessage::system(
        format!(
            "You are a task decomposition expert. Split the following task into {max_parallel} \
             independent sub-tasks that can run in parallel. Each sub-task must be self-contained \
             and produce useful output independently.\n\n\
             Output your response as a JSON array of objects, each with:\n\
             - \"directive\": a clear, detailed instruction for a subagent\n\
             - \"access\": one of \"read\", \"write\", or \"full\"\n\n\
             IMPORTANT: Return ONLY valid JSON, no other text. Example:\n\
             [{{\"directive\": \"Create the User model with fields...\", \"access\": \"write\"}}]"
        )
    );

    let user_msg =
        zesdex_domain::core::ChatMessage::user(format!("Task: {task}\n\nSplit into {max_parallel} parallel directives:"));

    match client
        .chat_with_tools_non_streaming(&[sys_msg, user_msg], None, Some(2048), Some(0.4), None)
    {
        Ok((response, _)) => {
            let text = response.content.unwrap_or_default();
            parse_directives_json(&text, max_parallel)
        }
        Err(e) => {
            warn!(error = %e, "auto-split LLM call failed, using fallback splitting");
            Ok(fallback_split(task, max_parallel))
        }
    }
}

/// Parse JSON directives from LLM response.
fn parse_directives_json(text: &str, max_parallel: usize) -> Result<Vec<(String, AccessTier)>> {
    // Find JSON array in the response
    let json_start = text.find('[');
    let json_end = text.rfind(']');

    let json_str = match (json_start, json_end) {
        (Some(start), Some(end)) if end > start => &text[start..=end],
        _ => {
            // Try parsing the whole text as JSON
            text.trim()
        }
    };

    let parsed: Vec<Value> = serde_json::from_str(json_str)
        .map_err(|e| anyhow::anyhow!("failed to parse directives JSON: {e}"))?;

    let directives: Vec<(String, AccessTier)> = parsed
        .into_iter()
        .take(max_parallel)
        .filter_map(|v| {
            let directive = v.get("directive")?.as_str()?.to_string();
            let access_str = v.get("access").and_then(|a| a.as_str()).unwrap_or("write");
            let access = match access_str {
                "read" => AccessTier::Read,
                "full" => AccessTier::Full,
                _ => AccessTier::Write,
            };
            Some((directive, access))
        })
        .collect();

    Ok(directives)
}

/// Fallback splitting when LLM call fails.
fn fallback_split(task: &str, max_parallel: usize) -> Vec<(String, AccessTier)> {
    // Simple heuristic split by common patterns
    let mut directives = Vec::new();

    if task.contains("backend") || task.contains("api") || task.contains("server") {
        directives.push((
            format!("Implement the backend/API components for: {task}"),
            AccessTier::Write,
        ));
    }

    if task.contains("frontend") || task.contains("ui") || task.contains("client") {
        directives.push((
            format!("Implement the frontend/UI components for: {task}"),
            AccessTier::Write,
        ));
    }

    if task.contains("test") || task.contains("unit") {
        directives.push((
            format!("Write unit tests for: {task}"),
            AccessTier::Read,
        ));
    }

    if directives.is_empty() {
        // Just split the task into generic chunks
        for i in 0..max_parallel {
            directives.push((
                format!("Part {} of parallel task: {task}", i + 1),
                AccessTier::Write,
            ));
        }
    }

    directives.truncate(max_parallel);
    directives
}

/// Use LLM to consolidate multiple agent results into a single response.
fn consolidate_results(
    results: &[(usize, String, String)],
    base_url: &str,
    api_key: &str,
    model: &str,
) -> Result<String> {
    let client = crate::llm::provider::LlmClient::new(
        api_key.to_string(),
        model.to_string(),
        Some(base_url.to_string()),
    );

    let mut summary = String::new();
    for (i, directive, result) in results {
        summary.push_str(&format!("### Agent {}: {}\n{}\n\n", i, directive, result));
    }

    let sys_msg = zesdex_domain::core::ChatMessage::system(
        "You are a synthesis expert. Consolidate the following parallel agent outputs \
         into a single coherent response. Remove duplication, reconcile conflicts, and \
         present the unified result in a well-structured format."
            .to_string(),
    );

    let user_msg = zesdex_domain::core::ChatMessage::user(format!(
        "Consolidate the following parallel agent outputs:\n\n{summary}"
    ));

    match client.chat_with_tools_non_streaming(&[sys_msg, user_msg], None, Some(2048), Some(0.3), None) {
        Ok((response, _)) => Ok(response.content.unwrap_or_else(|| summary.clone())),
        Err(e) => {
            warn!(error = %e, "consolidation LLM call failed, using raw concatenation");
            Ok(summary)
        }
    }
}