swarm-engine-core 0.1.6

//! OrchestratorBuilder - Orchestrator のビルダーパターン実装

use std::time::Duration;

use crate::agent::{Analyzer, BatchInvoker, ManagementStrategy, ManagerAgent, WorkerAgent};
use crate::environment::EnvironmentBox;
use crate::events::{ActionEventPublisher, LifecycleHook};
use crate::exploration::{
    AdaptiveOperatorProvider, DependencyGraphProvider, ExplorationSpaceV2, NodeRules,
    OperatorProvider,
};
use crate::extensions::Extensions;
use crate::learn::{LearningSnapshot, OfflineModel, OptimalParamsModel};

use super::config::SwarmConfig;
use super::Orchestrator;

/// Orchestrator ビルダー
pub struct OrchestratorBuilder {
    workers: Vec<Box<dyn WorkerAgent>>,
    managers: Vec<Box<dyn ManagerAgent>>,
    analyzer: Option<Box<dyn Analyzer>>,
    batch_invoker: Option<Box<dyn BatchInvoker>>,
    dependency_provider: Option<Box<dyn DependencyGraphProvider>>,
    operator_provider: Option<Box<dyn OperatorProvider<NodeRules>>>,
    action_collector: Option<ActionEventPublisher>,
    lifecycle_hook: Option<Box<dyn LifecycleHook>>,
    config: SwarmConfig,
    extensions: Extensions,
}

impl OrchestratorBuilder {
    pub fn new() -> Self {
        Self {
            workers: Vec::new(),
            managers: Vec::new(),
            analyzer: None,
            batch_invoker: None,
            dependency_provider: None,
            operator_provider: None,
            action_collector: None,
            lifecycle_hook: None,
            config: SwarmConfig::default(),
            extensions: Extensions::new(),
        }
    }

    /// Analyzer を設定
    pub fn analyzer(mut self, analyzer: impl Analyzer + 'static) -> Self {
        self.analyzer = Some(Box::new(analyzer));
        self
    }

    pub fn add_worker(mut self, worker: impl WorkerAgent + 'static) -> Self {
        self.workers.push(Box::new(worker));
        self
    }

    /// Box化済みの Worker を追加（eval等の動的生成向け）
    pub fn add_worker_boxed(mut self, worker: Box<dyn WorkerAgent>) -> Self {
        self.workers.push(worker);
        self
    }

    /// Manager を追加（複数追加可能）
    pub fn manager(mut self, manager: impl ManagerAgent + 'static) -> Self {
        self.managers.push(Box::new(manager));
        self
    }

    /// Box化済みの Manager を追加（動的生成向け）
    pub fn add_manager_boxed(mut self, manager: Box<dyn ManagerAgent>) -> Self {
        self.managers.push(manager);
        self
    }

    pub fn batch_invoker(mut self, invoker: impl BatchInvoker + 'static) -> Self {
        self.batch_invoker = Some(Box::new(invoker));
        self
    }

    /// Box化済みの BatchInvoker を設定（動的生成向け）
    pub fn batch_invoker_boxed(mut self, invoker: Box<dyn BatchInvoker>) -> Self {
        self.batch_invoker = Some(invoker);
        self
    }

    /// DependencyGraphProvider を設定
    ///
    /// run_task() 時に ExplorationSpace がない場合、
    /// このプロバイダーを使って DependencyGraph を自動生成する。
    ///
    /// # Example
    ///
    /// ```ignore
    /// // BatchInvoker を Provider として使用
    /// let provider = BatchInvokerProvider::new(batch_invoker);
    /// builder.dependency_provider(provider)
    /// ```
    pub fn dependency_provider(mut self, provider: impl DependencyGraphProvider + 'static) -> Self {
        self.dependency_provider = Some(Box::new(provider));
        self
    }

    /// Box化済みの DependencyGraphProvider を設定（動的生成向け）
    pub fn dependency_provider_boxed(mut self, provider: Box<dyn DependencyGraphProvider>) -> Self {
        self.dependency_provider = Some(provider);
        self
    }

    /// OperatorProvider を設定
    ///
    /// 探索時の Selection 戦略を制御する Provider を設定。
    /// デフォルトは `AdaptiveProvider`（エラー率に応じて動的切替）。
    ///
    /// # Example
    ///
    /// ```ignore
    /// use swarm_engine_core::exploration::{AdaptiveProvider, ConfigBasedProvider};
    ///
    /// // Adaptive（デフォルト推奨）- エラー率に応じて UCB1/Greedy/Thompson を自動切替
    /// builder.operator_provider(AdaptiveProvider::default())
    ///
    /// // 固定 Selection
    /// builder.operator_provider(ConfigBasedProvider::ucb1(1.41))
    /// builder.operator_provider(ConfigBasedProvider::fifo())
    /// ```
    pub fn operator_provider(
        mut self,
        provider: impl OperatorProvider<NodeRules> + 'static,
    ) -> Self {
        self.operator_provider = Some(Box::new(provider));
        self
    }

    /// Box化済みの OperatorProvider を設定（動的生成向け）
    pub fn operator_provider_boxed(
        mut self,
        provider: Box<dyn OperatorProvider<NodeRules>>,
    ) -> Self {
        self.operator_provider = Some(provider);
        self
    }

    /// ActionEventPublisher を設定
    ///
    /// 行動イベントを broadcast channel で配信する。
    /// 統計集計・永続化は Subscriber が担当。
    ///
    /// # Example
    ///
    /// ```ignore
    /// let (publisher, rx) = ActionEventPublisher::new(1024);
    ///
    /// // Subscriber を起動
    /// let (stats_sub, stats) = StatsSubscriber::with_new_stats(publisher.subscribe());
    /// tokio::spawn(stats_sub.run());
    ///
    /// // Builder に設定
    /// builder.action_collector(publisher)
    /// ```
    pub fn action_collector(mut self, collector: ActionEventPublisher) -> Self {
        self.action_collector = Some(collector);
        self
    }

    /// LifecycleHook を設定
    ///
    /// Swarm の開始・終了時に同期的にフックを実行する。
    /// Learning の自動実行などに使用。
    ///
    /// # Hook vs Sink
    ///
    /// - **Hook**: 同期 callback（本メソッドで設定）
    /// - **Sink**: async trait（`pipeline::EventSink`）
    ///
    /// # Example
    ///
    /// ```ignore
    /// use swarm_engine_core::events::LearningLifecycleHook;
    /// use swarm_engine_core::learn::TriggerBuilder;
    ///
    /// // 10 回の Eval 後に Learn を実行
    /// let hook = LearningLifecycleHook::new(learning_path)
    ///     .with_trigger(TriggerBuilder::every_n_episodes(10));
    ///
    /// builder.lifecycle_hook(Box::new(hook))
    /// ```
    pub fn lifecycle_hook(mut self, hook: Box<dyn LifecycleHook>) -> Self {
        self.lifecycle_hook = Some(hook);
        self
    }

    pub fn config(mut self, config: SwarmConfig) -> Self {
        self.config = config;
        self
    }

    pub fn tick_duration(mut self, duration: Duration) -> Self {
        self.config.tick_duration = duration;
        self
    }

    pub fn max_ticks(mut self, max: u64) -> Self {
        self.config.max_ticks = max;
        self
    }

    /// Manager 起動戦略を設定
    ///
    /// # Example
    ///
    /// ```ignore
    /// // 毎 Tick 起動（V2 フロー向け）
    /// builder.management_strategy(ManagementStrategy::EveryTick)
    ///
    /// // 5 Tick ごとに起動
    /// builder.management_strategy(ManagementStrategy::FixedInterval { interval: 5 })
    /// ```
    pub fn management_strategy(mut self, strategy: ManagementStrategy) -> Self {
        self.config.management_strategy = strategy;
        self
    }

    /// Extension (動的リソース) を追加
    ///
    /// Worker から `state.shared.extensions.get::<T>()` でアクセス可能
    pub fn extension<T: Send + Sync + 'static>(mut self, resource: T) -> Self {
        self.extensions.insert(resource);
        self
    }

    /// Extensions を一括設定
    pub fn extensions(mut self, extensions: Extensions) -> Self {
        self.extensions = extensions;
        self
    }

    /// EvalEnvironment を設定
    ///
    /// GenericWorker がアクション実行時に使用する環境を設定。
    /// 設定されていない場合は従来の execute_action にフォールバック。
    ///
    /// # Example
    ///
    /// ```ignore
    /// // デフォルト環境（Glob/Grep/Read）
    /// builder.environment(Box::new(DefaultEnvironment::new()))
    ///
    /// // カスタム環境（迷路）
    /// builder.environment(Box::new(MazeEnvironment::from_map(map)))
    /// ```
    pub fn environment(mut self, env: EnvironmentBox) -> Self {
        self.extensions.insert(env);
        self
    }

    /// ExplorationSpaceV2 を有効化（DependencyGraph があれば自動作成）
    ///
    /// build() 時に DependencyGraph が Extensions にあれば、
    /// ExplorationSpaceV2 を自動作成する。
    pub fn with_exploration(self) -> Self {
        // フラグは不要。DependencyGraph があれば build() で自動作成
        self
    }

    /// OfflineModel を適用（旧 API、後方互換）
    ///
    /// Offline 学習で生成されたモデルを適用し、最適化されたパラメータで
    /// Selection 戦略を構成する。
    ///
    /// # Example
    ///
    /// ```ignore
    /// let model = store.load_offline_model("my-scenario")?;
    /// builder.with_offline_model(model)
    /// ```
    pub fn with_offline_model(mut self, model: OfflineModel) -> Self {
        // OfflineModel のパラメータから AdaptiveOperatorProvider を構築
        let provider = AdaptiveOperatorProvider::default()
            .with_ucb1_c(model.parameters.ucb1_c)
            .with_maturity_threshold(model.strategy_config.maturity_threshold)
            .with_error_rate_threshold(model.strategy_config.error_rate_threshold);

        self.operator_provider = Some(Box::new(provider));

        // OfflineModel を Extensions に登録（Selection 側でも参照可能に）
        self.extensions.insert(model);

        self
    }

    /// OptimalParamsModel を適用
    ///
    /// OptimalParamsModel からパラメータを適用し、Selection 戦略を構成する。
    ///
    /// # Example
    ///
    /// ```ignore
    /// use swarm_engine_core::learn::{OptimalParamsModel, param_keys};
    ///
    /// let mut model = OptimalParamsModel::new();
    /// model.set_param(param_keys::UCB1_C, 1.5);
    /// builder.with_optimal_params_model(model)
    /// ```
    pub fn with_optimal_params_model(mut self, model: OptimalParamsModel) -> Self {
        use crate::learn::{param_keys, Parametric};

        // Parametric trait を使ってパラメータを取得
        let maturity_threshold = model
            .get_param(param_keys::MATURITY_THRESHOLD)
            .and_then(|v| v.as_i64())
            .unwrap_or(model.strategy_config.maturity_threshold as i64)
            as u32;

        let error_rate_threshold = model
            .get_param(param_keys::ERROR_RATE_THRESHOLD)
            .and_then(|v| v.as_f64())
            .unwrap_or(model.strategy_config.error_rate_threshold);

        // OptimalParamsModel のパラメータから AdaptiveOperatorProvider を構築
        let provider = AdaptiveOperatorProvider::default()
            .with_ucb1_c(model.ucb1_c())
            .with_maturity_threshold(maturity_threshold)
            .with_error_rate_threshold(error_rate_threshold);

        self.operator_provider = Some(Box::new(provider));

        // OptimalParamsModel を Extensions に登録
        self.extensions.insert(model);

        self
    }

    pub fn build(self, runtime: tokio::runtime::Handle) -> Orchestrator {
        let mut orchestrator = Orchestrator::new(self.workers, self.config.clone(), runtime);

        // OperatorProvider を設定（デフォルト: AdaptiveProvider）
        if let Some(provider) = self.operator_provider {
            orchestrator.operator_provider = provider;
        }
        // else: Orchestrator::new() で AdaptiveProvider::default() が設定済み

        // Extensions を State に設定
        orchestrator.state.shared.extensions = self.extensions;

        // SwarmConfig を Extensions に登録（Analyzer が max_ticks 等を参照可能に）
        orchestrator.state.shared.extensions.insert(self.config);

        // LearningSnapshot があれば LearnedProvider を作成
        if let Some(prior) = orchestrator
            .state
            .shared
            .extensions
            .remove::<LearningSnapshot>()
        {
            let mut stats = crate::learn::LearnStats::default();
            stats.load_prior(&prior);
            let provider = crate::learn::LearnStatsProvider::new(stats);
            orchestrator.learned_provider = Some(std::sync::Arc::new(provider));
        }

        // DependencyGraph があれば V2 ExplorationSpace を作成
        // ただし、dependency_provider が設定されている場合は使用しない
        // （provider が提供するグラフを優先するため、ensure_exploration_space() で処理）
        if self.dependency_provider.is_none() {
            if let Some(graph) = orchestrator
                .state
                .shared
                .extensions
                .remove::<crate::exploration::DependencyGraph>()
            {
                let rules: NodeRules = (&graph).into();
                let operator = orchestrator.operator_provider.provide(rules, None);
                orchestrator.space_v2 =
                    Some(ExplorationSpaceV2::new(operator).with_dependency_graph(graph));
            }
        }

        // Analyzer を設定
        if let Some(analyzer) = self.analyzer {
            orchestrator = orchestrator.with_analyzer(analyzer);
        }

        for manager in self.managers {
            orchestrator = orchestrator.add_manager(manager);
        }

        if let Some(invoker) = self.batch_invoker {
            orchestrator = orchestrator.with_batch_invoker(invoker);
        }

        if let Some(provider) = self.dependency_provider {
            orchestrator = orchestrator.with_dependency_provider(provider);
        }

        // ActionEventPublisher を設定
        if let Some(collector) = self.action_collector {
            orchestrator.action_collector = Some(collector);
        }

        // LifecycleHook を設定
        if let Some(hook) = self.lifecycle_hook {
            orchestrator.lifecycle_hook = Some(hook);
        }

        orchestrator
    }
}

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