collet 0.1.1

//! Phase 2: parallel task execution via work-stealing queues.

use std::collections::{HashMap, HashSet};

use tokio::sync::mpsc;
use tokio::task::JoinSet;
use tokio_util::sync::CancellationToken;

use crate::agent::r#loop::AgentEvent;
use crate::agent::subagent::{SubagentResult, SubagentTask};
use crate::agent::swarm::knowledge::{
    BlackboardKind, CheckpointResult, ExecutionCheckpoint, save_checkpoint,
};

use super::*;

impl SwarmCoordinator {
    /// Execute tasks using a work-stealing approach with checkpoint persistence.
    ///
    /// Periodically saves execution progress so the coordinator can recover
    /// from crashes. Each task completion updates the checkpoint.
    pub(super) async fn execute_with_work_stealing_checkpointed(
        &self,
        tasks: Vec<SwarmTask>,
        system_prompt: &str,
        event_tx: &mpsc::UnboundedSender<AgentEvent>,
        cancel: &CancellationToken,
        checkpoint: &mut ExecutionCheckpoint,
    ) -> Vec<SubagentResult> {
        let mut task_map: HashMap<String, SwarmTask> =
            tasks.iter().map(|t| (t.id.clone(), t.clone())).collect();
        let mut all_ids: HashSet<String> = tasks.iter().map(|t| t.id.clone()).collect();

        // Pre-populate completed set from checkpoint (for recovery)
        let mut completed: HashSet<String> =
            checkpoint.completed_task_ids.iter().cloned().collect();
        completed.extend(checkpoint.failed_task_ids.iter().cloned());

        let mut all_results: Vec<SubagentResult> = Vec::new();
        let mut join_set: JoinSet<(String, SubagentResult)> = JoinSet::new();
        let mut scheduler = TaskScheduler::new(&tasks, &completed, &all_ids, &task_map);
        let mut worktree_map: HashMap<String, String> = HashMap::new();
        let mut last_heartbeat = std::time::Instant::now();

        'outer: loop {
            if cancel.is_cancelled() {
                // Cancel all in-flight workers so paused ones don't block forever.
                for id in &scheduler.in_flight_ids {
                    self.knowledge.cancel_worker(id).await;
                }
                break;
            }

            // Periodic heartbeat: update checkpoint timestamp
            if last_heartbeat.elapsed().as_millis() as u64 > HEARTBEAT_INTERVAL_MS {
                checkpoint.heartbeat();
                let _ = save_checkpoint(checkpoint, &self.checkpoint_path());
                last_heartbeat = std::time::Instant::now();
            }

            // Build shared resources once (lazy — first call builds, rest reuse)
            let shared = self.ensure_shared_resources().await;

            // Dequeue ready tasks from the incrementally-maintained scheduler.
            // drain_ready handles dep checks, file-conflict checks, and batch dedup.
            let available_slots = self.hive_config.max_agents.saturating_sub(join_set.len());
            let dispatch_ids = scheduler.drain_ready(available_slots, &task_map, &completed);

            // Batch-score knowledge for all tasks in this wave (single lock acquisition)
            let batch_queries: Vec<(String, Vec<String>)> = dispatch_ids
                .iter()
                .map(|id| {
                    let t = &task_map[id];
                    (t.prompt.clone(), t.target_files.clone())
                })
                .collect();
            let knowledge_hints = self
                .knowledge
                .batch_relevant_summaries(&batch_queries, 15)
                .await;

            let kb_index_size = self.knowledge.index_size().await;
            let kb_worker_count = self.knowledge.active_worker_ids().await.len();
            tracing::trace!(
                index_entries = kb_index_size,
                active_workers = kb_worker_count,
                "Knowledge: scored wave batch"
            );
            for (task_id, knowledge_hint) in dispatch_ids.into_iter().zip(knowledge_hints) {
                let task = &task_map[&task_id];

                let prompt = if knowledge_hint.is_empty() {
                    task.prompt.clone()
                } else {
                    format!(
                        "{}\n\n## Shared Knowledge (filtered for this task)\n\n{knowledge_hint}",
                        task.prompt
                    )
                };

                // Realtime tier: post task claim to blackboard
                if self.hive_config.mode.has_realtime() {
                    self.knowledge
                        .post_to_blackboard(
                            &format!("claim:{}", task.id),
                            &task.prompt,
                            &task.id,
                            BlackboardKind::Claim,
                        )
                        .await;
                }

                // Post task status to blackboard for visibility
                self.knowledge
                    .post_to_blackboard(
                        &format!("task:{}:status", task.id),
                        "running",
                        "coordinator",
                        BlackboardKind::Status,
                    )
                    .await;

                let _ = event_tx.send(AgentEvent::PhaseChange {
                    label: format!("{} — dispatching task {}...", self.mode_label(), task.id),
                });

                // Allocate a git worktree only when isolation is genuinely needed:
                //   1. worktree is enabled in config
                //   2. this task touches specific files (tasks with empty target_files
                //      are typically read-only or exploratory — no isolation needed)
                //   3. another task is already in-flight (single-agent execution
                //      doesn't need isolation — it's the only writer)
                let worktree_path = if self.hive_config.worktree
                    && !task.target_files.is_empty()
                    && !join_set.is_empty()
                {
                    self.create_worktree(&task_id).await
                } else {
                    None
                };

                // Agent selection is disabled — always dispatch without an assigned agent.
                let resolved_agent_name: Option<String> = None;
                let roster: Vec<crate::config::AgentDef> = Vec::new();

                // Create per-worker cancel token and iteration budget for external control.
                let worker_cancel = CancellationToken::new();
                let worker_budget =
                    crate::agent::guard::IterationBudget::new(self.config.max_iterations);
                // Publish handles and get instruction receiver for this worker.
                let instruction_rx = self
                    .knowledge
                    .register_worker_handles(&task.id, worker_cancel.clone(), worker_budget.clone())
                    .await;

                let subagent_task = SubagentTask {
                    id: task.id.clone(),
                    prompt,
                    agent_name: resolved_agent_name,
                    available_agents: roster,
                    model_override: self.hive_config.worker_model.clone(),
                    working_dir_override: worktree_path.clone(),
                    outer_event_tx: Some(event_tx.clone()),
                    cancel_token: Some(worker_cancel),
                    iteration_budget: Some(worker_budget),
                    instruction_rx: Some(instruction_rx),
                };

                let client = self.client.clone();
                let config = self.config.clone();
                let sp = system_prompt.to_string();
                let wd = self.working_dir.clone();
                let lsp = self.lsp_manager.clone();
                let resources = self.worker_resources(shared);
                let task_id_c = task_id.clone();
                join_set.spawn(async move {
                    let inner = tokio::spawn(async move {
                        crate::agent::subagent::spawn_with_resources(
                            subagent_task,
                            client,
                            config,
                            sp,
                            wd,
                            lsp,
                            resources,
                        )
                        .await
                    });
                    let result = match inner.await {
                        Ok(r) => r,
                        Err(e) => {
                            tracing::error!("Worker {task_id_c} panicked: {e}");
                            SubagentResult {
                                id: task_id_c.clone(),
                                success: false,
                                response: format!("Worker panicked: {e}"),
                                modified_files: Vec::new(),
                                tool_calls: 0,
                                input_tokens: 0,
                                output_tokens: 0,
                                continuation_hint: None,
                            }
                        }
                    };
                    (task_id_c, result)
                });
                scheduler.mark_dispatched(&task_id, &task.target_files.clone());
                if let Some(wt) = worktree_path {
                    worktree_map.insert(task_id, wt);
                }
            }

            // If nothing in flight and nothing ready, we're done
            if join_set.is_empty() {
                break;
            }

            // Event-driven wait: block until a task finishes or all-paused is detected.
            let (finished_id, mut result) = 'wait: loop {
                tokio::select! {
                    Some(join_res) = join_set.join_next() => {
                        let (id, r) = join_res.expect("inner spawn handles panics");
                        scheduler.remove_in_flight(&id);
                        break 'wait (id, r);
                    }
                    _ = tokio::time::sleep(tokio::time::Duration::from_secs(2)) => {
                        let ids: Vec<String> = scheduler.in_flight_ids.iter().cloned().collect();
                        if self.knowledge.all_workers_paused(&ids).await {
                            tracing::info!(
                                "All {} in-flight workers are paused — waiting for resume",
                                ids.len()
                            );
                            tokio::select! {
                                _ = self.knowledge.wait_for_resume() => {}
                                _ = tokio::time::sleep(tokio::time::Duration::from_secs(5)) => {}
                            }
                            continue 'outer;
                        }
                    }
                }
            };

            // Clean up worker control handles for the completed worker.
            self.knowledge.remove_worker_handles(&finished_id).await;

            // Merge worktree changes back into main working dir
            if let Some(wt_path) = worktree_map.remove(&finished_id) {
                let merged_files = self.merge_and_remove_worktree(&wt_path).await;
                // Include actual worktree changes in the result
                result.modified_files.extend(merged_files);
            }

            // Record results in knowledge base with content snapshots for conflict resolution
            for path in &result.modified_files {
                let snapshot = Self::read_content_snapshot(path, &self.working_dir).await;
                // Track runtime file access for prediction accuracy and conflict detection.
                self.knowledge.track_file_access(&result.id, path).await;

                // Distinguish new file creation from edits by checking whether the path
                // existed before this agent ran. We approximate this by checking whether
                // the knowledge base already has a record for this file (meaning another
                // agent created/edited it first) vs. a fresh path.
                let full_path = std::path::Path::new(&self.working_dir).join(path);
                let mod_type = if full_path.exists() {
                    crate::agent::swarm::knowledge::ModificationType::Edited {
                        old_hash: 0,
                        new_hash: 0,
                    }
                } else {
                    // File was deleted by this agent (exists in modified_files but not on disk).
                    crate::agent::swarm::knowledge::ModificationType::Deleted
                };

                self.knowledge
                    .record_file_modification(&result.id, path, mod_type, 0, snapshot)
                    .await;
            }

            // Log file prediction accuracy
            if let Some(task) = task_map.get(&finished_id) {
                let (correct, missed, extra) = self
                    .knowledge
                    .file_prediction_accuracy(&finished_id, &task.target_files)
                    .await;
                if missed > 0 || extra > 0 {
                    tracing::info!(
                        "File prediction for {}: {} correct, {} missed, {} extra",
                        finished_id,
                        correct,
                        missed,
                        extra
                    );
                }
            }

            // Post completion summary to knowledge base
            self.knowledge
                .post_fact(
                    &result.id,
                    &format!("result:{}", result.id),
                    &format!(
                        "Agent {} {}: modified {} files, {} tool calls",
                        result.id,
                        if result.success {
                            "succeeded"
                        } else {
                            "failed"
                        },
                        result.modified_files.len(),
                        result.tool_calls,
                    ),
                )
                .await;

            // Realtime tier: broadcast completion so pending agents pick it up in their knowledge hint
            if self.hive_config.mode.has_realtime() && result.success {
                self.knowledge
                    .announce(
                        &result.id,
                        &format!(
                            "Agent '{}' completed. Modified files: [{}]. Summary: {}",
                            result.id,
                            result.modified_files.join(", "),
                            truncate(&result.response, 400),
                        ),
                    )
                    .await;
            }

            // Update blackboard task status
            self.knowledge
                .post_to_blackboard(
                    &format!("task:{}:status", finished_id),
                    if result.success {
                        "completed"
                    } else {
                        "failed"
                    },
                    "coordinator",
                    BlackboardKind::Status,
                )
                .await;

            // Update checkpoint with completed task
            checkpoint.record_completion(CheckpointResult {
                id: result.id.clone(),
                success: result.success,
                response: result.response.clone(),
                modified_files: result.modified_files.clone(),
                tool_calls: result.tool_calls,
                input_tokens: result.input_tokens,
                output_tokens: result.output_tokens,
            });
            let _ = save_checkpoint(checkpoint, &self.checkpoint_path());

            // If the worker hit an iteration limit, re-queue a continuation task
            // rather than treating the work as done.  The continuation hint contains
            // a summary of what was accomplished and what still needs to be finished.
            if let Some(ref hint) = result.continuation_hint {
                let cont_id = format!("{}-cont", finished_id);
                if !completed.contains(&cont_id) && !scheduler.in_flight_ids.contains(&cont_id) {
                    let _ = event_tx.send(AgentEvent::PhaseChange {
                        label: format!(
                            "{} — worker {} hit iteration limit; re-queuing continuation",
                            self.mode_label(),
                            finished_id
                        ),
                    });
                    let original_task = task_map.get(&finished_id);
                    let cont_task = SwarmTask {
                        id: cont_id.clone(),
                        prompt: hint.clone(),
                        role: original_task.map(|t| t.role.clone()).unwrap_or_default(),
                        agent_name: original_task.and_then(|t| t.agent_name.clone()),
                        target_files: original_task
                            .map(|t| t.target_files.clone())
                            .unwrap_or_default(),
                        dependencies: vec![finished_id.clone()],
                    };
                    task_map.insert(cont_id.clone(), cont_task.clone());
                    all_ids.insert(cont_id.clone());
                    // Register with scheduler (after task_map/all_ids are updated)
                    scheduler.add_task(&cont_task, &completed, &task_map, &all_ids);
                }
            }

            completed.insert(finished_id.clone());
            scheduler.on_completed(&finished_id, &completed, &task_map, &all_ids);
            all_results.push(result);
        }

        all_results
    }

    /// Execute tasks using a work-stealing approach.
    ///
    /// Instead of rigid wave barriers, tasks are dispatched as soon as their
    /// dependencies complete. This eliminates straggler delays within waves.
    /// Alternative scheduler using work-stealing instead of rigid wave barriers.
    /// Use via `run_work_stealing()` as an alternative entry point to `run()`.
    pub async fn execute_with_work_stealing(
        &self,
        tasks: Vec<SwarmTask>,
        system_prompt: &str,
        event_tx: &mpsc::UnboundedSender<AgentEvent>,
        cancel: &CancellationToken,
    ) -> Vec<SubagentResult> {
        let mut task_map: HashMap<String, SwarmTask> =
            tasks.iter().map(|t| (t.id.clone(), t.clone())).collect();
        let mut all_ids: HashSet<String> = tasks.iter().map(|t| t.id.clone()).collect();

        // Track which tasks have completed
        let mut completed: HashSet<String> = HashSet::new();
        let mut all_results: Vec<SubagentResult> = Vec::new();
        let mut join_set: JoinSet<(String, SubagentResult)> = JoinSet::new();
        let mut scheduler = TaskScheduler::new(&tasks, &completed, &all_ids, &task_map);

        'outer: loop {
            if cancel.is_cancelled() {
                // Cancel all in-flight workers so paused ones don't block forever.
                for id in &scheduler.in_flight_ids {
                    self.knowledge.cancel_worker(id).await;
                }
                break;
            }

            // Build shared resources once (lazy — first call builds, rest reuse)
            let shared = self.ensure_shared_resources().await;

            // Dequeue ready tasks from the incrementally-maintained scheduler.
            // drain_ready handles dep checks, file-conflict checks, and batch dedup.
            let available_slots = self.hive_config.max_agents.saturating_sub(join_set.len());
            let dispatch_ids = scheduler.drain_ready(available_slots, &task_map, &completed);

            // Batch-score knowledge for all tasks in this wave (single lock acquisition)
            let batch_queries: Vec<(String, Vec<String>)> = dispatch_ids
                .iter()
                .map(|id| {
                    let t = &task_map[id];
                    (t.prompt.clone(), t.target_files.clone())
                })
                .collect();
            let knowledge_hints = self
                .knowledge
                .batch_relevant_summaries(&batch_queries, 15)
                .await;

            for (task_id, knowledge_hint) in dispatch_ids.into_iter().zip(knowledge_hints) {
                let task = &task_map[&task_id];

                // Use targeted single-task relevance scoring for work-stealing path;
                // fall back to broader context summary if no relevant entries found.
                let effective_hint = if knowledge_hint.is_empty() {
                    let targeted = self
                        .knowledge
                        .build_relevant_summary(&task.prompt, &task.target_files, 15)
                        .await;
                    if targeted.is_empty() {
                        self.knowledge.build_context_summary().await
                    } else {
                        targeted
                    }
                } else {
                    knowledge_hint
                };

                // Check blackboard for existing claims / status on target files.
                for target_file in &task.target_files {
                    if let Some(summary) = self.knowledge.file_already_read(target_file).await {
                        tracing::debug!(
                            file = %target_file, read_by = %summary.agent_id,
                            "Target file already read by another agent (work-stealing)"
                        );
                    }
                }
                let claim_key = format!("claim:{}", task.id);
                if let Some(existing_claim) = self.knowledge.read_blackboard(&claim_key).await {
                    tracing::debug!(claim_author = %existing_claim.author, "Existing blackboard claim for task");
                }
                let status_entries = self
                    .knowledge
                    .blackboard_by_kind(crate::agent::swarm::knowledge::BlackboardKind::Status)
                    .await;
                tracing::trace!(
                    status_count = status_entries.len(),
                    "Blackboard status at work-stealing dispatch"
                );
                let prompt = if effective_hint.is_empty() {
                    task.prompt.clone()
                } else {
                    format!(
                        "{}\n\n## Shared Knowledge (filtered for this task)\n\n{effective_hint}",
                        task.prompt
                    )
                };

                // Realtime tier: post task claim to blackboard
                if self.hive_config.mode.has_realtime() {
                    self.knowledge
                        .post_to_blackboard(
                            &format!("claim:{}", task.id),
                            &task.prompt,
                            &task.id,
                            BlackboardKind::Claim,
                        )
                        .await;
                }

                let _ = event_tx.send(AgentEvent::PhaseChange {
                    label: format!("{} — dispatching task {}...", self.mode_label(), task.id),
                });

                // Agent selection is disabled — always dispatch without an assigned agent.
                let resolved_agent_name: Option<String> = None;
                let roster: Vec<crate::config::AgentDef> = Vec::new();
                // Create per-worker cancel token and iteration budget for external control.
                let worker_cancel = CancellationToken::new();
                let worker_budget =
                    crate::agent::guard::IterationBudget::new(self.config.max_iterations);

                // Publish handles and get instruction receiver for this worker.
                let instruction_rx = self
                    .knowledge
                    .register_worker_handles(&task.id, worker_cancel.clone(), worker_budget.clone())
                    .await;

                let subagent_task = SubagentTask {
                    id: task.id.clone(),
                    prompt,
                    agent_name: resolved_agent_name,
                    available_agents: roster,
                    model_override: self.hive_config.worker_model.clone(),
                    working_dir_override: None,
                    outer_event_tx: Some(event_tx.clone()),
                    cancel_token: Some(worker_cancel),
                    iteration_budget: Some(worker_budget),
                    instruction_rx: Some(instruction_rx),
                };

                let client = self.client.clone();
                let config = self.config.clone();
                let sp = system_prompt.to_string();
                let wd = self.working_dir.clone();
                let lsp = self.lsp_manager.clone();
                let resources = self.worker_resources(shared);
                let task_id_c = task_id.clone();
                join_set.spawn(async move {
                    let inner = tokio::spawn(async move {
                        crate::agent::subagent::spawn_with_resources(
                            subagent_task,
                            client,
                            config,
                            sp,
                            wd,
                            lsp,
                            resources,
                        )
                        .await
                    });
                    let result = match inner.await {
                        Ok(r) => r,
                        Err(e) => {
                            tracing::error!("Worker {task_id_c} panicked: {e}");
                            SubagentResult {
                                id: task_id_c.clone(),
                                success: false,
                                response: format!("Worker panicked: {e}"),
                                modified_files: Vec::new(),
                                tool_calls: 0,
                                input_tokens: 0,
                                output_tokens: 0,
                                continuation_hint: None,
                            }
                        }
                    };
                    (task_id_c, result)
                });
                scheduler.mark_dispatched(&task_id, &task.target_files.clone());
            }

            // If nothing in flight and nothing ready, we're done
            if join_set.is_empty() {
                break;
            }

            // Event-driven wait: block until a task finishes or all-paused is detected.
            let (finished_id, result) = 'wait: loop {
                tokio::select! {
                    Some(join_res) = join_set.join_next() => {
                        let (id, r) = join_res.expect("inner spawn handles panics");
                        scheduler.remove_in_flight(&id);
                        break 'wait (id, r);
                    }
                    _ = tokio::time::sleep(tokio::time::Duration::from_secs(2)) => {
                        let ids: Vec<String> = scheduler.in_flight_ids.iter().cloned().collect();
                        if self.knowledge.all_workers_paused(&ids).await {
                            tracing::info!(
                                "All {} in-flight workers are paused — waiting for resume",
                                ids.len()
                            );
                            tokio::select! {
                                _ = self.knowledge.wait_for_resume() => {}
                                _ = tokio::time::sleep(tokio::time::Duration::from_secs(5)) => {}
                            }
                            continue 'outer;
                        }
                    }
                }
            };

            // Clean up worker control handles for the completed worker.
            self.knowledge.remove_worker_handles(&finished_id).await;

            // Record results in knowledge base with content snapshots for conflict resolution
            for path in &result.modified_files {
                let snapshot = Self::read_content_snapshot(path, &self.working_dir).await;

                // Classify modification type: Created if no prior record exists in the
                // knowledge base, Deleted if the file is absent from disk, Edited otherwise.
                let full_path = std::path::Path::new(&self.working_dir).join(path);
                let existing_mods = self.knowledge.file_modifications(path).await;
                let mod_type = if !full_path.exists() {
                    crate::agent::swarm::knowledge::ModificationType::Deleted
                } else if existing_mods.is_empty() {
                    crate::agent::swarm::knowledge::ModificationType::Created
                } else {
                    crate::agent::swarm::knowledge::ModificationType::Edited {
                        old_hash: 0,
                        new_hash: 0,
                    }
                };

                self.knowledge
                    .record_file_modification(&result.id, path, mod_type, 0, snapshot)
                    .await;
            }

            // Log file prediction accuracy
            if let Some(task) = task_map.get(&finished_id) {
                let (correct, missed, extra) = self
                    .knowledge
                    .file_prediction_accuracy(&finished_id, &task.target_files)
                    .await;
                if missed > 0 || extra > 0 {
                    tracing::info!(
                        "File prediction for {}: {} correct, {} missed, {} extra",
                        finished_id,
                        correct,
                        missed,
                        extra
                    );
                }
            }

            // Post completion summary
            self.knowledge
                .post_fact(
                    &result.id,
                    &format!("result:{}", result.id),
                    &format!(
                        "Agent {} {}: modified {} files, {} tool calls",
                        result.id,
                        if result.success {
                            "succeeded"
                        } else {
                            "failed"
                        },
                        result.modified_files.len(),
                        result.tool_calls,
                    ),
                )
                .await;

            // Re-queue as continuation task if the worker hit an iteration limit.
            if let Some(ref hint) = result.continuation_hint {
                let cont_id = format!("{}-cont", finished_id);
                if !completed.contains(&cont_id) && !scheduler.in_flight_ids.contains(&cont_id) {
                    let _ = event_tx.send(AgentEvent::PhaseChange {
                        label: format!(
                            "{} — worker {} hit iteration limit; re-queuing continuation",
                            self.mode_label(),
                            finished_id
                        ),
                    });
                    let original_task = task_map.get(&finished_id);
                    let cont_task = SwarmTask {
                        id: cont_id.clone(),
                        prompt: hint.clone(),
                        role: original_task.map(|t| t.role.clone()).unwrap_or_default(),
                        agent_name: original_task.and_then(|t| t.agent_name.clone()),
                        target_files: original_task
                            .map(|t| t.target_files.clone())
                            .unwrap_or_default(),
                        dependencies: vec![finished_id.clone()],
                    };
                    task_map.insert(cont_id.clone(), cont_task.clone());
                    all_ids.insert(cont_id.clone());
                    // Register with scheduler (after task_map/all_ids are updated)
                    scheduler.add_task(&cont_task, &completed, &task_map, &all_ids);
                }
            }

            completed.insert(finished_id.clone());
            scheduler.on_completed(&finished_id, &completed, &task_map, &all_ids);
            all_results.push(result);
        }

        all_results
    }
}