swarm-engine-core 0.1.6

//! Lifecycle Hooks - Orchestrator ライフサイクルフック
//!
//! Swarm の開始・終了時に同期的に呼び出されるコールバック。
//!
//! # Hook vs Sink の命名規則
//!
//! | 名前 | 実行モデル | 用途 |
//! |------|----------|------|
//! | **Hook** | 同期 callback | Orchestrator lifecycle など、同期的に呼ばれる処理 |
//! | **Sink** | async trait | Pipeline 終端など、非同期でデータを消費する処理 |
//!
//! # 設計
//!
//! - `LifecycleEvent`: Swarm の状態変化を表すイベント
//! - `LifecycleHook`: 同期的に呼び出される callback trait
//! - `LearningLifecycleHook`: TrainTrigger + Learn 実行の実装
//!
//! # 使用例
//!
//! ```ignore
//! use swarm_engine_core::events::{LifecycleHook, LearningLifecycleHook};
//! use swarm_engine_core::learn::TriggerBuilder;
//!
//! // 10 回の Eval 実行後に Learn を実行
//! let hook = LearningLifecycleHook::new(learning_path)
//!     .with_trigger(TriggerBuilder::every_n_episodes(10));
//!
//! let orchestrator = OrchestratorBuilder::new()
//!     .lifecycle_hook(Box::new(hook))
//!     .build(runtime);
//! ```

/// Learn 実行コールバック型
type LearnCallback = Box<dyn Fn(&str, &SwarmState) + Send + Sync>;

use std::path::PathBuf;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;

use crate::learn::{AlwaysTrigger, TrainTrigger, TriggerContext};
use crate::orchestrator::SwarmResult;
use crate::state::SwarmState;

// ============================================================================
// LifecycleEvent
// ============================================================================

/// Orchestrator のライフサイクルイベント
#[derive(Debug, Clone)]
pub enum LifecycleEvent {
    /// Swarm 開始
    Started {
        /// Worker 数
        worker_count: usize,
    },
    /// Swarm 終了
    Terminated {
        /// 実行結果
        result: SwarmResult,
        /// 終了時の統計（シリアライズ可能な形式）
        stats: TerminationStats,
    },
}

/// 終了時の統計情報（シリアライズ可能）
#[derive(Debug, Clone, Default)]
pub struct TerminationStats {
    /// 総 Tick 数
    pub total_ticks: u64,
    /// 総アクション数
    pub total_actions: u64,
    /// 成功アクション数
    pub successful_actions: u64,
    /// 失敗アクション数
    pub failed_actions: u64,
    /// シナリオ名（あれば）
    pub scenario: Option<String>,
    /// GroupId（あれば）
    pub group_id: Option<String>,
}

impl TerminationStats {
    /// SwarmState から統計を抽出
    pub fn from_state(state: &SwarmState) -> Self {
        Self {
            total_ticks: state.shared.tick,
            total_actions: state.shared.stats.total_visits() as u64,
            successful_actions: state.shared.stats.total_successes() as u64,
            failed_actions: state.shared.stats.total_failures() as u64,
            scenario: None,
            group_id: None,
        }
    }

    /// シナリオ名を設定
    pub fn with_scenario(mut self, scenario: impl Into<String>) -> Self {
        self.scenario = Some(scenario.into());
        self
    }

    /// GroupId を設定
    pub fn with_group_id(mut self, group_id: impl Into<String>) -> Self {
        self.group_id = Some(group_id.into());
        self
    }
}

// ============================================================================
// LifecycleHook trait
// ============================================================================

/// Orchestrator ライフサイクルの同期コールバック
///
/// Swarm の開始・終了時に**同期的に**呼び出される。
/// Learning の自動実行、統計の永続化などに使用。
///
/// # Hook vs Sink
///
/// - **Hook**: 同期的に呼ばれる callback（本 trait）
/// - **Sink**: 非同期でデータを消費する処理（`pipeline::EventSink`）
///
/// # 注意
///
/// `on_terminate` 内で重い処理を行う場合は `tokio::task::spawn_blocking`
/// を使用してブロッキングを回避すること。
pub trait LifecycleHook: Send + Sync {
    /// Swarm 開始時に呼ばれる
    fn on_start(&mut self, worker_count: usize) {
        let _ = worker_count;
    }

    /// Swarm 終了時に呼ばれる
    ///
    /// `state` は終了時点の SwarmState への参照。
    /// Learning 実行などはここで行う。
    fn on_terminate(&mut self, state: &SwarmState, result: &SwarmResult);

    /// Hook の名前（デバッグ用）
    fn name(&self) -> &str {
        "lifecycle_hook"
    }
}

// ============================================================================
// LearningLifecycleHook
// ============================================================================

/// Learning を自動実行する LifecycleHook
///
/// Swarm 終了時に TrainTrigger をチェックし、条件を満たせば Learn を実行。
///
/// # Eval カウント
///
/// 内部で Eval 実行回数をカウントし、CountTrigger と連携。
/// 例: 10 回の Eval 後に Learn を実行。
///
/// # 非同期 Sink との違い
///
/// | 項目 | LearningLifecycleHook | LearningSink (pipeline) |
/// |------|----------------------|------------------------|
/// | 呼び出し | Orchestrator 終了時 | ファイル監視イベント時 |
/// | 実行モデル | 同期 callback | async trait |
/// | 用途 | Eval 終了後の即時学習 | Daemon モードでの継続学習 |
///
/// # 将来拡張
///
/// - IPC/WebSocket で外部 Learner プロセスと通信
/// - DPO 学習の自動実行
pub struct LearningLifecycleHook {
    /// 学習データのパス
    learning_path: PathBuf,
    /// TrainTrigger
    trigger: Arc<dyn TrainTrigger>,
    /// Eval 実行カウント（グローバル）
    eval_count: Arc<AtomicUsize>,
    /// 最後に Learn を実行した時の eval_count
    last_learn_count: usize,
    /// シナリオ名（Learn 実行時に使用）
    scenario: Option<String>,
    /// Learn 実行コールバック（カスタマイズ用）
    learn_callback: Option<LearnCallback>,
}

impl LearningLifecycleHook {
    /// 新しい LearningLifecycleHook を作成
    ///
    /// デフォルトは AlwaysTrigger（毎回 Learn 実行）。
    pub fn new(learning_path: impl Into<PathBuf>) -> Self {
        Self {
            learning_path: learning_path.into(),
            trigger: Arc::new(AlwaysTrigger),
            eval_count: Arc::new(AtomicUsize::new(0)),
            last_learn_count: 0,
            scenario: None,
            learn_callback: None,
        }
    }

    /// TrainTrigger を設定
    ///
    /// # Example
    ///
    /// ```ignore
    /// use swarm_engine_core::learn::TriggerBuilder;
    ///
    /// // 10 回の Eval 後に Learn 実行
    /// let hook = LearningLifecycleHook::new(path)
    ///     .with_trigger(TriggerBuilder::every_n_episodes(10));
    /// ```
    pub fn with_trigger(mut self, trigger: Arc<dyn TrainTrigger>) -> Self {
        self.trigger = trigger;
        self
    }

    /// シナリオ名を設定
    pub fn with_scenario(mut self, scenario: impl Into<String>) -> Self {
        self.scenario = Some(scenario.into());
        self
    }

    /// Learn 実行コールバックを設定
    ///
    /// デフォルトの Learn 実行ロジックをオーバーライド。
    /// DPO 学習など、カスタムの学習処理を実行できる。
    pub fn with_learn_callback<F>(mut self, callback: F) -> Self
    where
        F: Fn(&str, &SwarmState) + Send + Sync + 'static,
    {
        self.learn_callback = Some(Box::new(callback));
        self
    }

    /// 外部から eval_count を共有するための Arc を取得
    ///
    /// 複数の Orchestrator 間で Eval カウントを共有する場合に使用。
    pub fn eval_count_handle(&self) -> Arc<AtomicUsize> {
        Arc::clone(&self.eval_count)
    }

    /// 共有された eval_count を設定
    pub fn with_shared_eval_count(mut self, count: Arc<AtomicUsize>) -> Self {
        self.eval_count = count;
        self
    }

    /// 現在の Eval カウントを取得
    pub fn current_eval_count(&self) -> usize {
        self.eval_count.load(Ordering::SeqCst)
    }

    /// Learning パスを取得
    pub fn learning_path(&self) -> &PathBuf {
        &self.learning_path
    }

    /// Trigger をチェック
    fn should_learn(&self) -> bool {
        let current = self.eval_count.load(Ordering::SeqCst);
        let ctx = TriggerContext::with_count(current).last_train_count(self.last_learn_count);
        self.trigger.should_train(&ctx).unwrap_or(false)
    }

    /// Learn を実行
    fn run_learn(&mut self, state: &SwarmState) {
        let scenario = self.scenario.as_deref().unwrap_or("unknown");

        tracing::info!(
            scenario = scenario,
            eval_count = self.current_eval_count(),
            trigger = self.trigger.name(),
            "Running learning after trigger condition met"
        );

        // カスタムコールバックがあれば使用
        if let Some(ref callback) = self.learn_callback {
            callback(scenario, state);
        } else {
            // デフォルト: LearningStore を使った offline learning
            self.run_default_learn(scenario);
        }

        // last_learn_count を更新
        self.last_learn_count = self.eval_count.load(Ordering::SeqCst);
    }

    /// デフォルトの Learn 実行
    fn run_default_learn(&self, scenario: &str) {
        use crate::learn::LearningStore;

        match LearningStore::new(&self.learning_path) {
            Ok(store) => match store.run_offline_learning(scenario, 20) {
                Ok(model) => {
                    tracing::info!(
                        scenario = scenario,
                        sessions = model.analyzed_sessions,
                        "Offline learning completed"
                    );
                }
                Err(e) => {
                    tracing::warn!(
                        scenario = scenario,
                        error = %e,
                        "Offline learning failed"
                    );
                }
            },
            Err(e) => {
                tracing::error!(
                    path = %self.learning_path.display(),
                    error = %e,
                    "Failed to create LearningStore"
                );
            }
        }
    }
}

impl LifecycleHook for LearningLifecycleHook {
    fn on_start(&mut self, worker_count: usize) {
        tracing::debug!(
            worker_count = worker_count,
            eval_count = self.current_eval_count(),
            "LearningLifecycleHook: Swarm started"
        );
    }

    fn on_terminate(&mut self, state: &SwarmState, result: &SwarmResult) {
        // Eval カウントをインクリメント
        let new_count = self.eval_count.fetch_add(1, Ordering::SeqCst) + 1;

        tracing::debug!(
            eval_count = new_count,
            total_ticks = result.total_ticks,
            trigger = self.trigger.name(),
            "LearningLifecycleHook: Swarm terminated"
        );

        // Trigger チェック
        if self.should_learn() {
            self.run_learn(state);
        } else {
            tracing::debug!(
                eval_count = new_count,
                last_learn = self.last_learn_count,
                trigger = self.trigger.name(),
                "Trigger not met, skipping learning"
            );
        }
    }

    fn name(&self) -> &str {
        "learning_lifecycle_hook"
    }
}

// ============================================================================
// CompositeLifecycleHook
// ============================================================================

/// 複数の LifecycleHook を合成
///
/// 複数の Hook を登録し、順番に呼び出す。
///
/// # Example
///
/// ```ignore
/// let composite = CompositeLifecycleHook::new()
///     .with_hook(Box::new(LearningLifecycleHook::new(path)))
///     .with_hook(Box::new(CustomHook::new()));
/// ```
pub struct CompositeLifecycleHook {
    hooks: Vec<Box<dyn LifecycleHook>>,
}

impl CompositeLifecycleHook {
    /// 新しい CompositeLifecycleHook を作成
    pub fn new() -> Self {
        Self { hooks: Vec::new() }
    }

    /// Hook を追加
    pub fn with_hook(mut self, hook: Box<dyn LifecycleHook>) -> Self {
        self.hooks.push(hook);
        self
    }
}

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

impl LifecycleHook for CompositeLifecycleHook {
    fn on_start(&mut self, worker_count: usize) {
        for hook in &mut self.hooks {
            hook.on_start(worker_count);
        }
    }

    fn on_terminate(&mut self, state: &SwarmState, result: &SwarmResult) {
        for hook in &mut self.hooks {
            hook.on_terminate(state, result);
        }
    }

    fn name(&self) -> &str {
        "composite_lifecycle_hook"
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::AtomicBool;

    struct TestHook {
        started: Arc<AtomicBool>,
        terminated: Arc<AtomicBool>,
    }

    impl TestHook {
        fn new() -> (Self, Arc<AtomicBool>, Arc<AtomicBool>) {
            let started = Arc::new(AtomicBool::new(false));
            let terminated = Arc::new(AtomicBool::new(false));
            (
                Self {
                    started: Arc::clone(&started),
                    terminated: Arc::clone(&terminated),
                },
                started,
                terminated,
            )
        }
    }

    impl LifecycleHook for TestHook {
        fn on_start(&mut self, _worker_count: usize) {
            self.started.store(true, Ordering::SeqCst);
        }

        fn on_terminate(&mut self, _state: &SwarmState, _result: &SwarmResult) {
            self.terminated.store(true, Ordering::SeqCst);
        }
    }

    #[test]
    fn test_termination_stats_from_state() {
        let state = SwarmState::new(4);
        let stats = TerminationStats::from_state(&state);
        assert_eq!(stats.total_ticks, 0);
        assert!(stats.scenario.is_none());
    }

    #[test]
    fn test_learning_lifecycle_hook_eval_count() {
        let hook = LearningLifecycleHook::new("/tmp/test");
        assert_eq!(hook.current_eval_count(), 0);

        let handle = hook.eval_count_handle();
        handle.fetch_add(5, Ordering::SeqCst);
        assert_eq!(hook.current_eval_count(), 5);
    }

    #[test]
    fn test_composite_hook() {
        let (hook1, started1, terminated1) = TestHook::new();
        let (hook2, started2, terminated2) = TestHook::new();

        let mut composite = CompositeLifecycleHook::new()
            .with_hook(Box::new(hook1))
            .with_hook(Box::new(hook2));

        composite.on_start(4);
        assert!(started1.load(Ordering::SeqCst));
        assert!(started2.load(Ordering::SeqCst));

        let state = SwarmState::new(4);
        let result = SwarmResult {
            total_ticks: 10,
            total_duration: std::time::Duration::from_secs(1),
            completed: true,
        };
        composite.on_terminate(&state, &result);
        assert!(terminated1.load(Ordering::SeqCst));
        assert!(terminated2.load(Ordering::SeqCst));
    }
}