swarm-engine-core 0.1.6

//! TaskContext - Analyzer が生成するタスク状況
//!
//! # 設計
//!
//! Analyzer が SwarmState を分析して TaskContext を生成。
//! Manager は TaskContext を見て Request を生成する。
//!
//! ```text
//! SwarmState → [Analyzer] → TaskContext → [Manager] → BatchDecisionRequest
//! ```
//!
//! # 拡張性
//!
//! ベース情報（tick, workers, success_rate 等）に加え、
//! `metadata: HashMap<String, Value>` で任意の追加情報を格納可能。
//! 軽量LLMで分析した結果等を入れることを想定。

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

use serde_json::Value;

use crate::actions::ActionsConfig;
use crate::agent::Guidance;
use crate::state::Escalation;
use crate::types::WorkerId;

// ============================================================================
// WorkerSummary - Worker 状態の要約
// ============================================================================

/// Worker 状態の要約
#[derive(Debug, Clone)]
pub struct WorkerSummary {
    /// Worker ID
    pub id: WorkerId,
    /// 連続失敗数
    pub consecutive_failures: u32,
    /// 最新アクション名
    pub last_action: Option<String>,
    /// 最新アクションの成功/失敗
    pub last_success: Option<bool>,
    /// 最新アクションの出力（Environment からの結果）
    pub last_output: Option<String>,
    /// 履歴の長さ
    pub history_len: usize,
    /// Escalation 中かどうか
    pub has_escalation: bool,
}

impl WorkerSummary {
    pub fn new(id: WorkerId) -> Self {
        Self {
            id,
            consecutive_failures: 0,
            last_action: None,
            last_success: None,
            last_output: None,
            history_len: 0,
            has_escalation: false,
        }
    }

    /// 連続失敗数を設定
    pub fn with_failures(mut self, count: u32) -> Self {
        self.consecutive_failures = count;
        self
    }

    /// 最新アクションを設定
    pub fn with_last_action(mut self, action: impl Into<String>, success: bool) -> Self {
        self.last_action = Some(action.into());
        self.last_success = Some(success);
        self
    }

    /// 履歴の長さを設定
    pub fn with_history_len(mut self, len: usize) -> Self {
        self.history_len = len;
        self
    }

    /// Escalation を設定
    pub fn with_escalation(mut self, has_escalation: bool) -> Self {
        self.has_escalation = has_escalation;
        self
    }
}

// ============================================================================
// TaskContext - タスク状況
// ============================================================================

/// タスク状況（Analyzer が生成、Manager が消費）
///
/// # 構成
///
/// - **ベース情報**: tick, workers, success_rate, progress, escalations, available_actions
/// - **探索情報**: v2_guidances, excluded_actions
/// - **拡張 KV**: metadata で任意の追加情報を格納
///
/// # 使用例
///
/// ```ignore
/// let context = TaskContext::new(tick)
///     .with_worker(WorkerSummary::new(WorkerId(0)).with_phase(WorkerPhase::Working))
///     .with_progress(0.5)
///     .insert("llm_summary", "探索フェーズ中");
/// ```
#[derive(Debug, Clone)]
pub struct TaskContext {
    // === ベース情報 ===
    /// 現在の tick
    pub tick: u64,
    /// 各 Worker の状態サマリ
    pub workers: HashMap<WorkerId, WorkerSummary>,
    /// 成功率 (0.0 - 1.0)
    pub success_rate: f64,
    /// 進捗 (0.0 - 1.0)
    pub progress: f64,
    /// Escalation 一覧（WorkerId, Escalation）
    pub escalations: Vec<(WorkerId, Escalation)>,
    /// 利用可能なアクション
    pub available_actions: Option<ActionsConfig>,

    // === 探索情報 ===
    /// ExplorationSpaceV2 から生成された Guidance
    ///
    /// select_nodes() → Guidance 変換で直接生成。
    /// Manager はこれをそのまま Worker に配布できる。
    pub v2_guidances: Option<Vec<Guidance>>,
    /// 除外すべきアクション（成功済み/クローズ済み）- プロンプトから除外
    pub excluded_actions: Vec<String>,

    // === Manager 指示情報 ===
    /// 前回の Guidance（Worker ID -> Arc<Guidance>）
    ///
    /// Orchestrator が Manager.prepare() 呼び出し前に設定。
    /// Manager は prepare() でこれを使って ResolvedContext に ManagerInstruction を埋め込む。
    /// Arc で共有することでクローン時のディープコピーを回避。
    pub previous_guidances: HashMap<WorkerId, Arc<Guidance>>,

    /// Goal Action（Terminal Action）を達成した Worker
    /// これらのWorkerは Manager.prepare() でリクエスト対象から除外される
    pub done_workers: HashSet<WorkerId>,

    // === 拡張用 KV ===
    /// 追加メタデータ（軽量LLM分析結果等）
    pub metadata: HashMap<String, Value>,
}

impl TaskContext {
    /// 新しい TaskContext を作成
    pub fn new(tick: u64) -> Self {
        Self {
            tick,
            workers: HashMap::new(),
            success_rate: 0.0,
            progress: 0.0,
            escalations: Vec::new(),
            available_actions: None,
            v2_guidances: None,
            excluded_actions: Vec::new(),
            previous_guidances: HashMap::new(),
            done_workers: HashSet::new(),
            metadata: HashMap::new(),
        }
    }

    // === Builder methods ===

    /// Worker サマリを追加
    pub fn with_worker(mut self, summary: WorkerSummary) -> Self {
        self.workers.insert(summary.id, summary);
        self
    }

    /// 成功率を設定
    pub fn with_success_rate(mut self, rate: f64) -> Self {
        self.success_rate = rate;
        self
    }

    /// 進捗を設定
    pub fn with_progress(mut self, progress: f64) -> Self {
        self.progress = progress;
        self
    }

    /// Escalation を追加
    pub fn with_escalation(mut self, worker_id: WorkerId, escalation: Escalation) -> Self {
        self.escalations.push((worker_id, escalation));
        self
    }

    /// 利用可能なアクションを設定
    pub fn with_actions(mut self, actions: ActionsConfig) -> Self {
        self.available_actions = Some(actions);
        self
    }

    /// 前回の Guidance を設定（Arc で共有）
    pub fn with_previous_guidances(mut self, guidances: HashMap<WorkerId, Arc<Guidance>>) -> Self {
        self.previous_guidances = guidances;
        self
    }

    /// 単一 Worker の前回 Guidance を追加
    pub fn with_previous_guidance(mut self, worker_id: WorkerId, guidance: Arc<Guidance>) -> Self {
        self.previous_guidances.insert(worker_id, guidance);
        self
    }

    // === Metadata (KV) methods ===

    /// メタデータを追加
    pub fn insert<V: Into<Value>>(mut self, key: impl Into<String>, value: V) -> Self {
        self.metadata.insert(key.into(), value.into());
        self
    }

    /// メタデータを追加（mutable）
    pub fn set<V: Into<Value>>(&mut self, key: impl Into<String>, value: V) {
        self.metadata.insert(key.into(), value.into());
    }

    /// メタデータを取得
    pub fn get(&self, key: &str) -> Option<&Value> {
        self.metadata.get(key)
    }

    /// メタデータを文字列として取得
    pub fn get_str(&self, key: &str) -> Option<&str> {
        self.metadata.get(key).and_then(|v| v.as_str())
    }

    /// メタデータを数値として取得
    pub fn get_f64(&self, key: &str) -> Option<f64> {
        self.metadata.get(key).and_then(|v| v.as_f64())
    }

    /// メタデータを整数として取得
    pub fn get_i64(&self, key: &str) -> Option<i64> {
        self.metadata.get(key).and_then(|v| v.as_i64())
    }

    /// メタデータを真偽値として取得
    pub fn get_bool(&self, key: &str) -> Option<bool> {
        self.metadata.get(key).and_then(|v| v.as_bool())
    }

    // === Query methods ===

    /// Escalation があるか
    pub fn has_escalations(&self) -> bool {
        !self.escalations.is_empty()
    }

    /// 特定の Worker が Escalation 中かどうか
    pub fn has_escalation_for(&self, worker_id: WorkerId) -> bool {
        self.escalations.iter().any(|(id, _)| *id == worker_id)
    }

    /// 特定 Worker の情報を取得
    pub fn worker(&self, id: WorkerId) -> Option<&WorkerSummary> {
        self.workers.get(&id)
    }

    /// Escalation 中の Worker 数を取得
    pub fn escalated_worker_count(&self) -> usize {
        self.workers.values().filter(|w| w.has_escalation).count()
    }

    /// Worker ID 一覧を取得
    pub fn worker_ids(&self) -> Vec<WorkerId> {
        self.workers.keys().copied().collect()
    }
}

impl TaskContext {
    // === Exploration methods ===

    /// 探索が有効か（v2_guidances がある場合）
    pub fn has_exploration(&self) -> bool {
        self.v2_guidances.is_some()
    }

    /// 指定した Worker のみを含むフィルタ済み TaskContext を作成
    ///
    /// Manager のパーティショニングで使用。
    /// 各 Manager は担当する Worker のみの TaskContext を受け取る。
    pub fn filter_for_workers(&self, worker_ids: &[WorkerId]) -> TaskContext {
        use std::collections::HashSet;
        let worker_set: HashSet<WorkerId> = worker_ids.iter().copied().collect();

        // フィルタ済み workers
        let filtered_workers: HashMap<WorkerId, WorkerSummary> = self
            .workers
            .iter()
            .filter(|(id, _)| worker_set.contains(id))
            .map(|(id, summary)| (*id, summary.clone()))
            .collect();

        // フィルタ済み escalations
        let filtered_escalations: Vec<(WorkerId, Escalation)> = self
            .escalations
            .iter()
            .filter(|(id, _)| worker_set.contains(id))
            .cloned()
            .collect();

        // フィルタ済み previous_guidances（Arc::clone は軽量）
        let filtered_guidances: HashMap<WorkerId, Arc<Guidance>> = self
            .previous_guidances
            .iter()
            .filter(|(id, _)| worker_set.contains(id))
            .map(|(id, g)| (*id, Arc::clone(g)))
            .collect();

        // フィルタ済み done_workers
        let filtered_done_workers: HashSet<WorkerId> = self
            .done_workers
            .iter()
            .filter(|id| worker_set.contains(id))
            .copied()
            .collect();

        TaskContext {
            tick: self.tick,
            workers: filtered_workers,
            success_rate: self.success_rate,
            progress: self.progress,
            escalations: filtered_escalations,
            available_actions: self.available_actions.clone(),
            v2_guidances: self.v2_guidances.clone(),
            excluded_actions: self.excluded_actions.clone(),
            previous_guidances: filtered_guidances,
            done_workers: filtered_done_workers,
            metadata: self.metadata.clone(),
        }
    }
}

impl Default for TaskContext {
    fn default() -> Self {
        Self::new(0)
    }
}

// ============================================================================
// Tests
// ============================================================================

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

    #[test]
    fn test_task_context_new() {
        let ctx = TaskContext::new(10);
        assert_eq!(ctx.tick, 10);
        assert!(ctx.workers.is_empty());
        assert_eq!(ctx.success_rate, 0.0);
        assert_eq!(ctx.progress, 0.0);
    }

    #[test]
    fn test_task_context_builder() {
        let ctx = TaskContext::new(5)
            .with_worker(WorkerSummary::new(WorkerId(0)))
            .with_worker(WorkerSummary::new(WorkerId(1)).with_escalation(true))
            .with_success_rate(0.8)
            .with_progress(0.5)
            .insert("key1", "value1")
            .insert("count", 42);

        assert_eq!(ctx.tick, 5);
        assert_eq!(ctx.workers.len(), 2);
        assert_eq!(ctx.success_rate, 0.8);
        assert_eq!(ctx.progress, 0.5);
        assert_eq!(ctx.get_str("key1"), Some("value1"));
        assert_eq!(ctx.get_i64("count"), Some(42));
    }

    #[test]
    fn test_worker_summary() {
        let summary = WorkerSummary::new(WorkerId(0))
            .with_failures(2)
            .with_last_action("read:/path", true)
            .with_history_len(10)
            .with_escalation(true);

        assert_eq!(summary.id, WorkerId(0));
        assert_eq!(summary.consecutive_failures, 2);
        assert_eq!(summary.last_action, Some("read:/path".to_string()));
        assert_eq!(summary.last_success, Some(true));
        assert_eq!(summary.history_len, 10);
        assert!(summary.has_escalation);
    }

    #[test]
    fn test_query_methods() {
        let ctx = TaskContext::new(0)
            .with_worker(WorkerSummary::new(WorkerId(0)))
            .with_worker(WorkerSummary::new(WorkerId(1)).with_escalation(true))
            .with_worker(WorkerSummary::new(WorkerId(2)));

        assert_eq!(ctx.escalated_worker_count(), 1);
        assert_eq!(ctx.worker_ids().len(), 3);
    }

    #[test]
    fn test_filter_for_workers() {
        // Setup: 4 workers with different states
        let ctx = TaskContext::new(10)
            .with_worker(WorkerSummary::new(WorkerId(0)).with_failures(1))
            .with_worker(WorkerSummary::new(WorkerId(1)).with_escalation(true))
            .with_worker(WorkerSummary::new(WorkerId(2)).with_history_len(5))
            .with_worker(WorkerSummary::new(WorkerId(3)).with_last_action("read", true))
            .with_escalation(WorkerId(1), Escalation::consecutive_failures(3, 5))
            .with_success_rate(0.75)
            .with_progress(0.5)
            .insert("meta_key", "meta_value");

        // Filter for workers 0 and 2 only
        let filtered = ctx.filter_for_workers(&[WorkerId(0), WorkerId(2)]);

        // Verify filtered context
        assert_eq!(filtered.tick, 10);
        assert_eq!(filtered.workers.len(), 2);
        assert!(filtered.workers.contains_key(&WorkerId(0)));
        assert!(filtered.workers.contains_key(&WorkerId(2)));
        assert!(!filtered.workers.contains_key(&WorkerId(1)));
        assert!(!filtered.workers.contains_key(&WorkerId(3)));

        // Worker 0 should have its failures preserved
        assert_eq!(
            filtered
                .workers
                .get(&WorkerId(0))
                .unwrap()
                .consecutive_failures,
            1
        );

        // Worker 2 should have its history_len preserved
        assert_eq!(filtered.workers.get(&WorkerId(2)).unwrap().history_len, 5);

        // Escalations should be filtered (worker 1's escalation should be excluded)
        assert!(filtered.escalations.is_empty());

        // Global metrics should be preserved
        assert_eq!(filtered.success_rate, 0.75);
        assert_eq!(filtered.progress, 0.5);

        // Metadata should be preserved
        assert_eq!(filtered.get_str("meta_key"), Some("meta_value"));
    }

    #[test]
    fn test_filter_for_workers_empty() {
        let ctx = TaskContext::new(5)
            .with_worker(WorkerSummary::new(WorkerId(0)))
            .with_worker(WorkerSummary::new(WorkerId(1)));

        let filtered = ctx.filter_for_workers(&[]);

        assert_eq!(filtered.tick, 5);
        assert!(filtered.workers.is_empty());
    }
}