swarm-engine-core 0.1.6

//! SwarmEngine Orchestrator
//!
//! メインループを管理し、Tick駆動でAgent Swarmを実行する。
//!
//! # DependencyGraph 自動判別
//!
//! DependencyGraph（タスク依存関係グラフ）は以下の優先順位で**自動的に**取得・生成される：
//!
//! 1. **DependencyGraphProvider** - 明示的に設定されたプロバイダーを使用
//! 2. **BatchInvoker.plan_dependencies()** - LLM がタスクとアクション一覧から依存関係を推論
//! 3. **Extensions に登録された DependencyGraph** - 静的な依存関係グラフ
//!
//! シナリオファイル（TOML）で `dependency_graph` を設定しなくても、
//! `batch_invoker` が設定されていれば LLM が自動生成を試みる。
//! 生成に失敗した場合（LLM が適切なグラフを出力できない等）はエラーとなる。
//!
//! ## 例：自動生成が成功するケース
//!
//! ```text
//! Task: "Diagnose and restart the failing service"
//! Actions: ["CheckStatus", "ReadLogs", "Diagnose", "Restart"]
//! → LLM が CheckStatus → ReadLogs → Diagnose → Restart の依存関係を推論
//! ```
//!
//! ## 例：自動生成が失敗するケース
//!
//! ```text
//! Task: "Navigate through the maze"
//! Actions: ["Move", "Look", "Wait"]
//! → LLM が適切な依存関係を推論できない場合がある
//! → シナリオファイルで明示的に dependency_graph を設定する必要がある
//! ```

mod adapter;
mod builder;
mod config;
mod execution;
mod lifecycle;
mod manager;
mod merge;
pub mod termination;

pub use adapter::WorkResultAdapter;
pub use builder::OrchestratorBuilder;
pub use config::{SwarmConfig, SwarmResult};
pub use termination::{TerminationConfig, TerminationJudge, TerminationVerdict};

use std::sync::Arc;
use std::time::Instant;

use tracing::{debug, info};

use crate::actions::ActionDef;
use crate::agent::{Analyzer, BatchInvoker, DefaultAnalyzer, ManagerAgent, WorkerAgent};
use crate::async_task::AsyncTaskSystem;
use crate::error::SwarmError;
use crate::events::{ActionEventPublisher, LearningEventChannel, LifecycleHook};
use crate::exploration::{
    ActionNodeData, AdaptiveOperatorProvider, ConfigurableSpace, DependencyGraphProvider,
    MapNodeState, NodeRules, OperatorProvider,
};
use crate::state::SwarmState;
use crate::types::{SwarmTask, WorkerId};

/// Orchestrator - メインループ管理
pub struct Orchestrator {
    /// 共有状態
    pub(crate) state: SwarmState,
    /// Worker Agents
    pub(crate) workers: Vec<Box<dyn WorkerAgent>>,
    /// Manager Agents（複数対応）
    pub(crate) managers: Vec<Box<dyn ManagerAgent>>,
    /// Analyzer（SwarmState → TaskContext）
    pub(crate) analyzer: Box<dyn Analyzer>,
    /// Batch Invoker（LLM 推論、オプション）
    pub(crate) batch_invoker: Option<Box<dyn BatchInvoker>>,
    /// DependencyGraph Provider（ExplorationSpace 自動生成用）
    pub(crate) dependency_provider: Option<Box<dyn DependencyGraphProvider>>,
    /// 非同期タスクシステム
    pub(crate) async_system: AsyncTaskSystem,
    /// 設定
    pub(crate) config: SwarmConfig,
    /// Termination Judge - Single Source of Truth for termination decisions
    pub(crate) termination_judge: TerminationJudge,
    /// Manager 最終起動Tick（Manager ID -> Tick）
    pub(crate) last_manager_ticks: std::collections::HashMap<crate::agent::ManagerId, u64>,
    /// 現在の Guidance（Manager から配布、次 Tick でクリア）
    /// Arc で共有することでクローン時のディープコピーを回避
    pub(crate) current_guidances: std::collections::HashMap<WorkerId, Arc<crate::agent::Guidance>>,
    /// Worker のパーティション割り当て（Manager ID -> Worker IDs）
    /// None の場合は全 Manager が全 Worker を担当
    pub(crate) worker_assignments:
        Option<std::collections::HashMap<crate::agent::ManagerId, Vec<WorkerId>>>,

    // ========================================================================
    // Exploration V2 (新アーキテクチャ)
    // ========================================================================
    /// ExplorationSpace V2 - 新しい探索空間
    /// GraphMap + Operator ベースの統合レイヤー
    pub(crate) space_v2: Option<ConfigurableSpace<NodeRules>>,

    /// OperatorProvider（Selection 動的切り替え用）
    ///
    /// AdaptiveProvider: エラー率に応じて UCB1 → Greedy/Thompson を自動切替
    /// ConfigBasedProvider: 固定の Selection を使用
    pub(crate) operator_provider: Box<dyn OperatorProvider<NodeRules>>,

    /// ActionEventPublisher（行動イベント配信、オプション）
    pub(crate) action_collector: Option<ActionEventPublisher>,

    /// LearnedProvider（学習済みデータへのアクセス、オプション）
    pub(crate) learned_provider: Option<crate::learn::SharedLearnedProvider>,

    /// LifecycleHook（開始・終了時のフック）
    pub(crate) lifecycle_hook: Option<Box<dyn LifecycleHook>>,
}

impl Orchestrator {
    /// 新規作成
    pub fn new(
        workers: Vec<Box<dyn WorkerAgent>>,
        config: SwarmConfig,
        runtime: tokio::runtime::Handle,
    ) -> Self {
        let agent_count = workers.len();
        let termination_config = TerminationConfig::with_max_ticks(config.max_ticks);
        Self {
            state: SwarmState::new(agent_count),
            workers,
            managers: Vec::new(),
            analyzer: Box::new(DefaultAnalyzer::new()),
            batch_invoker: None,
            dependency_provider: None,
            async_system: AsyncTaskSystem::new(runtime),
            config,
            termination_judge: TerminationJudge::new(termination_config, agent_count),
            last_manager_ticks: std::collections::HashMap::new(),
            current_guidances: std::collections::HashMap::new(),
            worker_assignments: None,
            space_v2: None,
            operator_provider: Box::new(AdaptiveOperatorProvider::default()),
            action_collector: None,
            learned_provider: None,
            lifecycle_hook: None,
        }
    }

    /// Analyzer を設定
    pub fn with_analyzer(mut self, analyzer: Box<dyn Analyzer>) -> Self {
        self.analyzer = analyzer;
        self
    }

    /// Manager を追加
    pub fn add_manager(mut self, manager: Box<dyn ManagerAgent>) -> Self {
        self.managers.push(manager);
        self
    }

    /// BatchInvoker を設定
    pub fn with_batch_invoker(mut self, invoker: Box<dyn BatchInvoker>) -> Self {
        self.batch_invoker = Some(invoker);
        self
    }

    /// DependencyGraphProvider を設定
    ///
    /// run_task() 時に ExplorationSpace がない場合、
    /// このプロバイダーを使って DependencyGraph を自動生成する。
    pub fn with_dependency_provider(mut self, provider: Box<dyn DependencyGraphProvider>) -> Self {
        self.dependency_provider = Some(provider);
        self
    }

    /// Worker パーティショニングを有効化
    ///
    /// 各 Manager に Worker を均等に割り当てる。
    /// Manager ごとに担当 Worker のみの TaskContext を渡すため、
    /// 大規模システムでのスケーラビリティが向上する。
    ///
    /// # 呼び出しタイミング
    ///
    /// 全ての Manager を追加した後に呼び出すこと。
    pub fn enable_partitioning(&mut self) {
        if self.managers.is_empty() {
            return;
        }

        let worker_count = self.workers.len();
        let manager_count = self.managers.len();
        let workers_per_manager = worker_count.div_ceil(manager_count);

        let mut assignments = std::collections::HashMap::new();
        let all_worker_ids: Vec<WorkerId> = (0..worker_count).map(WorkerId).collect();

        for (i, manager) in self.managers.iter().enumerate() {
            let start = i * workers_per_manager;
            let end = ((i + 1) * workers_per_manager).min(worker_count);
            let assigned: Vec<WorkerId> = all_worker_ids[start..end].to_vec();
            assignments.insert(manager.id(), assigned);
        }

        self.worker_assignments = Some(assignments);
    }

    /// 特定の Manager に割り当てられた Worker IDs を取得
    pub(crate) fn get_assigned_workers(
        &self,
        manager_id: crate::agent::ManagerId,
    ) -> Option<Vec<WorkerId>> {
        self.worker_assignments
            .as_ref()
            .and_then(|assignments| assignments.get(&manager_id).cloned())
    }

    /// DependencyGraph への参照を取得
    ///
    /// ExplorationSpaceV2 に設定されている DependencyGraph を返す。
    /// 学習データ保存時に action_order を抽出するために使用。
    pub fn dependency_graph(&self) -> Option<&crate::exploration::DependencyGraph> {
        self.space_v2
            .as_ref()
            .and_then(|space| space.dependency_graph())
    }

    /// タスクを設定してメインループを実行
    ///
    /// SwarmTask を Extensions に登録してから run() を実行する。
    /// Manager は `state.shared.extensions.get::<SwarmTask>()` でタスクを読み取れる。
    ///
    /// # Errors
    ///
    /// - `SwarmError::MissingDependencyGraph`: DependencyGraph を生成できなかった
    ///
    /// # Example
    ///
    /// ```ignore
    /// let task = SwarmTask::new("Find the auth handler")
    ///     .with_target_path("/path/to/repo");
    ///
    /// let result = orchestrator.run_task(task)?;
    /// ```
    pub fn run_task(&mut self, task: SwarmTask) -> Result<SwarmResult, SwarmError> {
        // タスクを Extensions に登録
        self.state.shared.extensions.insert(task.clone());

        // ActionsConfig からアクション定義一覧を取得
        let actions: Vec<ActionDef> = self
            .state
            .shared
            .extensions
            .get::<crate::actions::ActionsConfig>()
            .map(|cfg| cfg.all_actions().cloned().collect())
            .unwrap_or_default();

        // ExplorationSpace がない場合、DependencyGraphProvider で自動生成
        // DependencyGraph は必須。生成できない場合はエラー。
        self.ensure_exploration_space(&task.goal, &actions)?;

        // initial_context を取得（V2 の initialize() 用）
        // task.context["initial_context"] が配列なら使用、なければ target_service を使用
        let initial_contexts: Vec<String> = task
            .context
            .get("initial_context")
            .and_then(|v| v.as_array())
            .map(|arr| {
                arr.iter()
                    .filter_map(|v| v.as_str().map(|s| s.to_string()))
                    .collect()
            })
            .unwrap_or_else(|| {
                // フォールバック: target_service を単一の初期コンテキストとして使用
                task.context
                    .get("target_service")
                    .and_then(|v| v.as_str())
                    .map(|s| vec![s.to_string()])
                    .unwrap_or_default()
            });

        // ========================================================================
        // ExplorationSpaceV2 の初期化
        // ========================================================================
        if let Some(ref mut space_v2) = self.space_v2 {
            // ルートノード作成
            let root_id = space_v2.create_root(ActionNodeData::new("root"));
            debug!(root_id = ?root_id, "ExplorationSpaceV2 root node created");

            // initial_context に基づいて初期ノードを展開
            if !initial_contexts.is_empty() {
                let ctx_refs: Vec<&str> = initial_contexts.iter().map(|s| s.as_str()).collect();
                let results = space_v2.initialize(&ctx_refs);
                info!(
                    initial_contexts = ?initial_contexts,
                    expanded_nodes = results.len(),
                    "ExplorationSpaceV2: initial nodes expanded via Rules"
                );

                // 初期ノード展開後、ルートノードを Close
                // これによりルートがフロンティアから除外され、初期アクションが選択対象になる
                if !results.is_empty() {
                    space_v2.map_mut().set_state(root_id, MapNodeState::Closed);
                    debug!(root_id = ?root_id, "ExplorationSpaceV2 root node closed after initialization");
                }
            }
        }

        Ok(self.run())
    }

    /// メインループを実行
    pub fn run(&mut self) -> SwarmResult {
        let start = Instant::now();
        let worker_count = self.workers.len();
        info!(worker_count = worker_count, "system_start");

        // LifecycleHook: on_start
        if let Some(ref mut hook) = self.lifecycle_hook {
            hook.on_start(worker_count);
        }

        loop {
            let tick_start = Instant::now();
            let current_tick = self.state.shared.tick;

            // Update termination judge's tick
            self.termination_judge.set_tick(current_tick);

            // Tick 開始イベント（telemetry 用）
            info!(tick = current_tick, "tick_start");

            // LearningEventChannel の current_tick を更新
            LearningEventChannel::global().set_tick(current_tick);

            // TickStart ActionEvent を発行
            {
                let event = crate::events::ActionEventBuilder::new(
                    current_tick,
                    crate::types::WorkerId::MANAGER,
                    "tick_start",
                )
                .result(crate::events::ActionEventResult::success())
                .build();
                self.state.shared.stats.record(&event);
                if let Some(ref collector) = self.action_collector {
                    collector.record(event);
                }
            }

            // 1. 非同期結果があれば取り込み
            self.collect_async_results();

            // 2. Manager 起動判定 & Guidance 更新
            // Note: manager_activation と llm_errors は ActionEvent で記録される
            // (manager.rs で llm_invoke イベントを発行)
            if self.should_run_manager() {
                let _ = self.run_manager();
            } else {
                // Manager が起動しない Tick でも、ExplorationSpaceV2 があれば Guidance を自動生成
                // これにより、LLM 呼び出しなしで探索を継続できる
                self.generate_exploration_guidances();
            }

            // 3. Execute Phase: 各 Worker が並列実行
            let results = self.execute_workers();

            // 4. Merge Phase: 結果を State に反映
            self.merge_results(&results);

            // 4.5. SharedData クリーンアップ（メモリリーク防止）
            self.state.shared.shared_data.cleanup_env_entries();

            // 5. Tick更新
            self.state.advance_tick();
            // Update termination judge's tick for correct termination check
            self.termination_judge.set_tick(self.state.shared.tick);

            // 6. 次Tickまで待機
            let elapsed = tick_start.elapsed();
            if elapsed < self.config.tick_duration {
                std::thread::sleep(self.config.tick_duration - elapsed);
            }

            // 平均Tick時間を更新（EMA with α=0.1）
            // EMA = α × current + (1-α) × EMA_old = (current + 9 × EMA_old) / 10
            let current_ns = elapsed.as_nanos() as u64;
            let prev_avg = self.state.shared.avg_tick_duration_ns;
            self.state.shared.avg_tick_duration_ns = if prev_avg == 0 {
                current_ns // 初回は現在値をそのまま使用
            } else {
                (current_ns + 9 * prev_avg) / 10
            };

            // LearningEventChannel の sync_buffer をクリア
            // （LearningEvent は LearningEventSubscriber が処理するため、ここでは drain のみ）
            let _ = LearningEventChannel::global().drain_sync();

            // TickEnd ActionEvent を発行（duration を metadata に含める）
            {
                let event = crate::events::ActionEventBuilder::new(
                    current_tick,
                    crate::types::WorkerId::MANAGER,
                    "tick_end",
                )
                .duration(elapsed)
                .result(crate::events::ActionEventResult::success())
                .context(
                    crate::events::ActionContext::new()
                        .with_metadata("duration_ns", elapsed.as_nanos().to_string()),
                )
                .build();
                self.state.shared.stats.record(&event);
                if let Some(ref collector) = self.action_collector {
                    collector.record(event);
                }
            }

            // Tick 完了イベント（telemetry 用）
            info!(
                tick = current_tick,
                duration_ns = elapsed.as_nanos() as u64,
                total_actions = self.state.shared.stats.total_visits(),
                successful_actions = self.state.shared.stats.total_successes(),
                failed_actions = self.state.shared.stats.total_failures(),
                active_workers = self.workers.len() as u64,
                "tick_complete"
            );

            // 終了判定
            if self.should_terminate() {
                break;
            }
        }

        let total_duration = start.elapsed();

        // システム停止イベント（telemetry 用）
        info!(
            total_ticks = self.state.shared.tick,
            total_duration_ms = total_duration.as_millis() as u64,
            "system_stop"
        );

        let result = SwarmResult {
            total_ticks: self.state.shared.tick,
            total_duration,
            completed: true,
        };

        // LifecycleHook: on_terminate
        if let Some(ref mut hook) = self.lifecycle_hook {
            hook.on_terminate(&self.state, &result);
        }

        result
    }

    /// 外部から終了を要求
    pub fn request_terminate(&mut self) {
        self.termination_judge.request_terminate("External request");
    }

    /// TerminationJudge への参照を取得
    pub fn termination_judge(&self) -> &TerminationJudge {
        &self.termination_judge
    }

    /// TerminationJudge への可変参照を取得
    pub fn termination_judge_mut(&mut self) -> &mut TerminationJudge {
        &mut self.termination_judge
    }

    /// State への参照を取得
    pub fn state(&self) -> &SwarmState {
        &self.state
    }

    /// AsyncTaskSystem への参照を取得
    pub fn async_system(&self) -> &AsyncTaskSystem {
        &self.async_system
    }

    /// LearnedProvider への参照を取得
    pub fn learned_provider(&self) -> Option<&crate::learn::SharedLearnedProvider> {
        self.learned_provider.as_ref()
    }
}