swarm_engine_core/learn/daemon/

mod.rs

//! Learning Daemon - 継続的学習プロセス
//!
//! ## 概要
//!
//! Swarm セッションを跨いで動作する長寿命プロセス。
//! イベントを受信し、データを永続化し、Trigger 発火時に学習を実行する。
//!
//! ## アーキテクチャ
//!
//! ```text
//! EventSource → DataSink → Store
//!                   ↓
//!             TriggerLoop → Processor → Applier
//! ```
//!
//! ## 使用例
//!
//! ```ignore
//! use swarm_engine_core::learn::daemon::{LearningDaemon, DaemonConfig};
//!
//! let config = DaemonConfig::new("my-scenario")
//!     .trigger(TriggerBuilder::every_n_episodes(100))
//!     .auto_apply(true);
//!
//! let daemon = LearningDaemon::new(config)?;
//! daemon.run().await?;
//! ```

mod applier;
mod processor;
mod sink;
mod subscriber;

pub use applier::{Applier, ApplierConfig, ApplierError, ApplyMode, ApplyResult};
pub use processor::{ProcessResult, Processor, ProcessorConfig, ProcessorError, ProcessorMode};
pub use sink::{DataSink, DataSinkError, DataSinkStats};
pub use subscriber::{ActionEventSubscriber, EventSubscriberConfig, LearningEventSubscriber};

use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;

use tokio::sync::mpsc;
use tokio::time::interval;

use crate::learn::learn_model::{LearnModel, WorkerDecisionSequenceLearn};
use crate::learn::lora::{
    LoraTrainer, LoraTrainerConfig, ModelApplicator, NoOpApplicator, TrainedModel,
};
use crate::learn::record::{LearnStatsRecord, Record};
use crate::learn::snapshot::LearningStore;
use crate::learn::store::{
    EpisodeStore, FileEpisodeStore, FileRecordStore, InMemoryEpisodeStore, InMemoryRecordStore,
    RecordStore, RecordStoreError, StoreError,
};
use crate::learn::trigger::{TrainTrigger, TriggerBuilder, TriggerContext};
use crate::learn::LearnStats;
use crate::util::epoch_millis;

// ============================================================================
// DaemonError
// ============================================================================

/// Daemon のエラー型
#[derive(Debug)]
pub enum DaemonError {
    /// DataSink エラー
    Sink(DataSinkError),
    /// Processor エラー
    Processor(ProcessorError),
    /// Applier エラー
    Applier(ApplierError),
    /// IO エラー
    Io(std::io::Error),
    /// 設定エラー
    Config(String),
    /// Shutdown
    Shutdown,
}

impl std::fmt::Display for DaemonError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Sink(e) => write!(f, "Sink error: {}", e),
            Self::Processor(e) => write!(f, "Processor error: {}", e),
            Self::Applier(e) => write!(f, "Applier error: {}", e),
            Self::Io(e) => write!(f, "IO error: {}", e),
            Self::Config(msg) => write!(f, "Config error: {}", msg),
            Self::Shutdown => write!(f, "Daemon shutdown"),
        }
    }
}

impl std::error::Error for DaemonError {}

impl From<DataSinkError> for DaemonError {
    fn from(e: DataSinkError) -> Self {
        Self::Sink(e)
    }
}

impl From<ProcessorError> for DaemonError {
    fn from(e: ProcessorError) -> Self {
        Self::Processor(e)
    }
}

impl From<ApplierError> for DaemonError {
    fn from(e: ApplierError) -> Self {
        Self::Applier(e)
    }
}

impl From<std::io::Error> for DaemonError {
    fn from(e: std::io::Error) -> Self {
        Self::Io(e)
    }
}

impl From<RecordStoreError> for DaemonError {
    fn from(e: RecordStoreError) -> Self {
        Self::Sink(DataSinkError::RecordStore(e))
    }
}

impl From<StoreError> for DaemonError {
    fn from(e: StoreError) -> Self {
        Self::Sink(DataSinkError::EpisodeStore(e))
    }
}

// ============================================================================
// DaemonConfig
// ============================================================================

/// Daemon の設定
#[derive(Debug, Clone)]
pub struct DaemonConfig {
    /// シナリオ名
    pub scenario: String,
    /// データディレクトリ
    pub data_dir: PathBuf,
    /// Trigger チェック間隔
    pub check_interval: Duration,
    /// 処理モード
    pub processor_mode: ProcessorMode,
    /// Offline 分析に使用するセッション数
    pub max_sessions: usize,
    /// 自動適用
    pub auto_apply: bool,
    /// LoRA 設定（None の場合は LoRA 学習しない）
    pub lora_config: Option<LoraTrainerConfig>,
}

impl DaemonConfig {
    /// 新しい設定を作成
    pub fn new(scenario: impl Into<String>) -> Self {
        Self {
            scenario: scenario.into(),
            data_dir: default_data_dir(),
            check_interval: Duration::from_secs(10),
            processor_mode: ProcessorMode::OfflineOnly,
            max_sessions: 20,
            auto_apply: false,
            lora_config: None,
        }
    }

    /// データディレクトリを設定
    pub fn data_dir(mut self, path: impl Into<PathBuf>) -> Self {
        self.data_dir = path.into();
        self
    }

    /// Trigger チェック間隔を設定
    pub fn check_interval(mut self, interval: Duration) -> Self {
        self.check_interval = interval;
        self
    }

    /// 処理モードを設定
    pub fn processor_mode(mut self, mode: ProcessorMode) -> Self {
        self.processor_mode = mode;
        self
    }

    /// 最大セッション数を設定
    pub fn max_sessions(mut self, n: usize) -> Self {
        self.max_sessions = n;
        self
    }

    /// 自動適用を設定
    pub fn auto_apply(mut self, enabled: bool) -> Self {
        self.auto_apply = enabled;
        self
    }

    /// LoRA 設定を追加
    pub fn with_lora(mut self, config: LoraTrainerConfig) -> Self {
        self.lora_config = Some(config);
        self
    }
}

fn default_data_dir() -> PathBuf {
    dirs::data_dir()
        .unwrap_or_else(|| PathBuf::from("."))
        .join("swarm-engine")
        .join("learning")
}

// ============================================================================
// DaemonStats
// ============================================================================

/// Daemon の統計情報
#[derive(Debug, Clone, Default)]
pub struct DaemonStats {
    /// 受信した Record 数
    pub records_received: usize,
    /// 生成した Episode 数
    pub episodes_created: usize,
    /// 実行した学習回数
    pub trainings_completed: usize,
    /// 適用したモデル数
    pub models_applied: usize,
    /// 最終学習時刻（Unix timestamp ms）
    pub last_train_at: Option<u64>,
    /// 起動時刻（Unix timestamp ms）
    pub started_at: u64,
}

// ============================================================================
// LearningDaemon
// ============================================================================

/// 継続的学習プロセス
pub struct LearningDaemon {
    /// 設定
    config: DaemonConfig,
    /// DataSink
    sink: DataSink,
    /// Trigger
    trigger: Arc<dyn TrainTrigger>,
    /// Processor
    processor: Processor,
    /// Applier（オプション）
    applier: Option<Applier>,
    /// LearningStore（session 保存用）
    learning_store: LearningStore,
    /// 統計情報
    stats: DaemonStats,
    /// 最終学習時の Episode 数
    last_train_count: usize,
    /// Record 受信チャンネル
    record_rx: mpsc::Receiver<Vec<Record>>,
    /// Record 送信チャンネル（外部公開用）
    record_tx: mpsc::Sender<Vec<Record>>,
    /// Shutdown チャンネル
    shutdown_rx: mpsc::Receiver<()>,
    /// Shutdown 送信チャンネル（外部公開用）
    shutdown_tx: mpsc::Sender<()>,
}

impl LearningDaemon {
    /// 新しい Daemon を作成（In-Memory Store）
    pub fn new(config: DaemonConfig, trigger: Arc<dyn TrainTrigger>) -> Result<Self, DaemonError> {
        let record_store: Arc<dyn RecordStore> = Arc::new(InMemoryRecordStore::new());
        let episode_store: Arc<dyn EpisodeStore> = Arc::new(InMemoryEpisodeStore::new());
        let learn_model: Arc<dyn LearnModel> = Arc::new(WorkerDecisionSequenceLearn::new());

        Self::with_stores(config, trigger, record_store, episode_store, learn_model)
    }

    /// 新しい Daemon を作成（File Store）
    pub fn with_file_stores(
        config: DaemonConfig,
        trigger: Arc<dyn TrainTrigger>,
    ) -> Result<Self, DaemonError> {
        // ディレクトリ作成
        std::fs::create_dir_all(&config.data_dir)?;

        let record_store: Arc<dyn RecordStore> =
            Arc::new(FileRecordStore::new(config.data_dir.join("records"))?);
        let episode_store: Arc<dyn EpisodeStore> =
            Arc::new(FileEpisodeStore::new(config.data_dir.join("episodes"))?);
        let learn_model: Arc<dyn LearnModel> = Arc::new(WorkerDecisionSequenceLearn::new());

        Self::with_stores(config, trigger, record_store, episode_store, learn_model)
    }

    /// カスタム Store で Daemon を作成
    pub fn with_stores(
        config: DaemonConfig,
        trigger: Arc<dyn TrainTrigger>,
        record_store: Arc<dyn RecordStore>,
        episode_store: Arc<dyn EpisodeStore>,
        learn_model: Arc<dyn LearnModel>,
    ) -> Result<Self, DaemonError> {
        // DataSink
        let sink = DataSink::new(
            record_store,
            Arc::clone(&episode_store),
            Arc::clone(&learn_model),
        );

        // Processor
        let processor_config = ProcessorConfig::new(&config.scenario)
            .mode(config.processor_mode)
            .max_sessions(config.max_sessions);

        let mut processor = Processor::new(processor_config);

        // LearningStore for offline analysis and session persistence
        let learning_store = LearningStore::new(&config.data_dir)?;
        // Clone for processor (processor needs its own reference)
        let learning_store_for_processor = LearningStore::new(&config.data_dir)?;
        processor = processor.with_learning_store(learning_store_for_processor);

        // LoRA Trainer (if configured)
        if let Some(lora_config) = &config.lora_config {
            let trainer = LoraTrainer::new(lora_config.clone(), episode_store);
            processor = processor
                .with_lora_trainer(trainer)
                .with_learn_model(learn_model);
        }

        // Applier
        let applier = if config.auto_apply {
            let applier_config = ApplierConfig::default().auto_apply();
            // TODO: Real applicator based on config
            let applicator: Arc<dyn ModelApplicator> = Arc::new(NoOpApplicator::new());
            Some(Applier::new(applier_config, applicator))
        } else {
            None
        };

        // Channels
        let (record_tx, record_rx) = mpsc::channel(1000);
        let (shutdown_tx, shutdown_rx) = mpsc::channel(1);

        Ok(Self {
            config,
            sink,
            trigger,
            processor,
            applier,
            learning_store,
            stats: DaemonStats {
                started_at: epoch_millis(),
                ..Default::default()
            },
            last_train_count: 0,
            record_rx,
            record_tx,
            shutdown_rx,
            shutdown_tx,
        })
    }

    /// Record 送信チャンネルを取得
    pub fn record_sender(&self) -> mpsc::Sender<Vec<Record>> {
        self.record_tx.clone()
    }

    /// Shutdown 送信チャンネルを取得
    pub fn shutdown_sender(&self) -> mpsc::Sender<()> {
        self.shutdown_tx.clone()
    }

    /// 設定を取得
    pub fn config(&self) -> &DaemonConfig {
        &self.config
    }

    /// 統計情報を取得
    pub fn stats(&self) -> &DaemonStats {
        &self.stats
    }

    /// メインループを実行
    pub async fn run(&mut self) -> Result<(), DaemonError> {
        tracing::info!(
            scenario = %self.config.scenario,
            data_dir = %self.config.data_dir.display(),
            trigger = self.trigger.name(),
            "Learning daemon started"
        );

        let mut check_interval = interval(self.config.check_interval);

        loop {
            tokio::select! {
                // Shutdown signal
                _ = self.shutdown_rx.recv() => {
                    tracing::info!("Shutdown signal received, draining remaining records...");

                    // Drain remaining records before shutdown
                    // Give subscribers a moment to flush their batches
                    tokio::time::sleep(std::time::Duration::from_millis(100)).await;

                    // Process any remaining records in the channel
                    while let Ok(records) = self.record_rx.try_recv() {
                        if let Err(e) = self.handle_records(records).await {
                            tracing::warn!("Error processing records during shutdown: {}", e);
                        }
                    }

                    tracing::info!(
                        records_received = self.stats.records_received,
                        episodes_created = self.stats.episodes_created,
                        "Shutdown complete"
                    );
                    return Ok(());
                }

                // Record batch received
                Some(records) = self.record_rx.recv() => {
                    self.handle_records(records).await?;
                }

                // Periodic trigger check
                _ = check_interval.tick() => {
                    self.check_and_train().await?;
                }
            }
        }
    }

    /// Record を処理
    async fn handle_records(&mut self, records: Vec<Record>) -> Result<(), DaemonError> {
        if records.is_empty() {
            return Ok(());
        }

        let count = records.len();

        // LearnStatsRecord を検出して LearningStore に保存
        for record in &records {
            if let Record::LearnStats(stats_record) = record {
                self.save_stats_to_learning_store(stats_record);
            }
        }

        let episode_ids = self.sink.ingest(records)?;

        self.stats.records_received += count;
        self.stats.episodes_created += episode_ids.len();

        tracing::debug!(
            records = count,
            episodes = episode_ids.len(),
            "Processed records"
        );

        Ok(())
    }

    /// LearnStatsRecord を LearningSnapshot に変換して保存
    fn save_stats_to_learning_store(&self, stats_record: &LearnStatsRecord) {
        use crate::learn::snapshot::{LearningSnapshot, SnapshotMetadata, SNAPSHOT_VERSION};
        use crate::learn::{EpisodeTransitions, NgramStats, SelectionPerformance};
        use crate::online_stats::ActionStats;
        use std::collections::HashMap;

        // LearnStats JSON をデシリアライズ
        let learn_stats: Option<LearnStats> = serde_json::from_str(&stats_record.stats_json).ok();

        // SnapshotMetadata を作成（実際のフィールドのみ使用）
        let metadata = SnapshotMetadata {
            scenario_name: Some(stats_record.scenario.clone()),
            task_description: None,
            created_at: stats_record.timestamp_ms / 1000, // ms → s
            session_count: 1,
            total_episodes: 1,
            total_actions: stats_record.total_actions as u32,
            phase: None,
            group_id: None,
        };

        // LearnStats からデータを抽出（あれば）
        let (
            episode_transitions,
            action_stats,
            ngram_stats,
            selection_performance,
            contextual_stats,
        ) = if let Some(ref stats) = learn_stats {
            // EpisodeTransitions を直接コピー
            let transitions = stats.episode_transitions.clone();

            // NgramStats を直接コピー
            let ngram = stats.ngram_stats.clone();

            // SelectionPerformance を直接コピー
            let selection = stats.selection_performance.clone();

            // contextual_stats を ActionStats に変換
            let mut ctx_stats: HashMap<(String, String), ActionStats> = HashMap::new();
            for ((prev, action), ctx) in &stats.contextual_stats {
                ctx_stats.insert(
                    (prev.clone(), action.clone()),
                    ActionStats {
                        visits: ctx.visits,
                        successes: ctx.successes,
                        failures: ctx.failures,
                        ..Default::default()
                    },
                );
            }

            // action_stats は空（LearnStats には直接のアクション統計はない）
            let action_stats: HashMap<String, ActionStats> = HashMap::new();

            (transitions, action_stats, ngram, selection, ctx_stats)
        } else {
            (
                EpisodeTransitions::default(),
                HashMap::new(),
                NgramStats::default(),
                SelectionPerformance::default(),
                HashMap::new(),
            )
        };

        // LearningSnapshot を作成
        let snapshot = LearningSnapshot {
            version: SNAPSHOT_VERSION,
            metadata,
            episode_transitions,
            ngram_stats,
            selection_performance,
            contextual_stats,
            action_stats,
        };

        // LearningStore に保存
        match self
            .learning_store
            .save_session(&stats_record.scenario, &snapshot)
        {
            Ok(session_id) => {
                tracing::info!(
                    scenario = %stats_record.scenario,
                    session_id = %session_id.0,
                    success = stats_record.is_success(),
                    "Saved session to LearningStore"
                );
            }
            Err(e) => {
                tracing::warn!(
                    scenario = %stats_record.scenario,
                    error = %e,
                    "Failed to save session to LearningStore"
                );
            }
        }
    }

    /// Trigger をチェックし、必要なら学習を実行
    async fn check_and_train(&mut self) -> Result<(), DaemonError> {
        // Trigger 判定用コンテキスト
        let current_count = self.sink.episode_count();
        let ctx = TriggerContext::with_count(current_count)
            .last_train_at(self.stats.last_train_at.unwrap_or(0))
            .last_train_count(self.last_train_count);

        if !self.trigger.should_train(&ctx).unwrap_or(false) {
            return Ok(());
        }

        tracing::info!(
            episode_count = current_count,
            trigger = self.trigger.name(),
            "Trigger fired, starting learning"
        );

        // 学習実行
        let result = self
            .processor
            .run(self.sink.episode_store().as_ref())
            .await?;

        // 統計更新
        self.stats.trainings_completed += 1;
        self.stats.last_train_at = Some(epoch_millis());
        self.last_train_count = current_count;

        // 自動適用
        if let Some(applier) = &mut self.applier {
            if let Some(model) = result.lora_model() {
                let apply_result = applier.apply(model).await?;
                if apply_result.is_applied() {
                    self.stats.models_applied += 1;
                }
            }
        }

        tracing::info!(
            trainings = self.stats.trainings_completed,
            models_applied = self.stats.models_applied,
            "Learning cycle completed"
        );

        Ok(())
    }

    /// 手動で学習を実行（Trigger 無視）
    pub async fn train_now(&mut self) -> Result<ProcessResult, DaemonError> {
        tracing::info!("Manual training triggered");

        let result = self
            .processor
            .run(self.sink.episode_store().as_ref())
            .await?;

        self.stats.trainings_completed += 1;
        self.stats.last_train_at = Some(epoch_millis());
        self.last_train_count = self.sink.episode_count();

        Ok(result)
    }

    /// 手動でモデルを適用（非同期）
    pub async fn apply_model(&mut self, model: &TrainedModel) -> Result<ApplyResult, DaemonError> {
        let applier = self
            .applier
            .as_mut()
            .ok_or_else(|| DaemonError::Config("Applier not configured".into()))?;

        let result = applier.apply_now(model).await?;
        if result.is_applied() {
            self.stats.models_applied += 1;
        }

        Ok(result)
    }
}

// ============================================================================
// Builder Pattern
// ============================================================================

/// LearningDaemon のビルダー
pub struct DaemonBuilder {
    config: DaemonConfig,
    trigger: Option<Arc<dyn TrainTrigger>>,
    record_store: Option<Arc<dyn RecordStore>>,
    episode_store: Option<Arc<dyn EpisodeStore>>,
    learn_model: Option<Arc<dyn LearnModel>>,
    applicator: Option<Arc<dyn ModelApplicator>>,
}

impl DaemonBuilder {
    /// 新しいビルダーを作成
    pub fn new(scenario: impl Into<String>) -> Self {
        Self {
            config: DaemonConfig::new(scenario),
            trigger: None,
            record_store: None,
            episode_store: None,
            learn_model: None,
            applicator: None,
        }
    }

    /// データディレクトリを設定
    pub fn data_dir(mut self, path: impl Into<PathBuf>) -> Self {
        self.config.data_dir = path.into();
        self
    }

    /// Trigger を設定
    pub fn trigger(mut self, trigger: Arc<dyn TrainTrigger>) -> Self {
        self.trigger = Some(trigger);
        self
    }

    /// 処理モードを設定
    pub fn processor_mode(mut self, mode: ProcessorMode) -> Self {
        self.config.processor_mode = mode;
        self
    }

    /// 自動適用を有効化
    pub fn auto_apply(mut self) -> Self {
        self.config.auto_apply = true;
        self
    }

    /// RecordStore を設定
    pub fn record_store(mut self, store: Arc<dyn RecordStore>) -> Self {
        self.record_store = Some(store);
        self
    }

    /// EpisodeStore を設定
    pub fn episode_store(mut self, store: Arc<dyn EpisodeStore>) -> Self {
        self.episode_store = Some(store);
        self
    }

    /// LearnModel を設定
    pub fn learn_model(mut self, model: Arc<dyn LearnModel>) -> Self {
        self.learn_model = Some(model);
        self
    }

    /// ModelApplicator を設定
    pub fn applicator(mut self, applicator: Arc<dyn ModelApplicator>) -> Self {
        self.applicator = Some(applicator);
        self
    }

    /// LoRA 設定を追加
    pub fn with_lora(mut self, config: LoraTrainerConfig) -> Self {
        self.config.lora_config = Some(config);
        self
    }

    /// Daemon をビルド
    pub fn build(self) -> Result<LearningDaemon, DaemonError> {
        let trigger = self
            .trigger
            .unwrap_or_else(|| TriggerBuilder::default_watch());

        let record_store = self
            .record_store
            .unwrap_or_else(|| Arc::new(InMemoryRecordStore::new()));

        let episode_store = self
            .episode_store
            .unwrap_or_else(|| Arc::new(InMemoryEpisodeStore::new()));

        let learn_model = self
            .learn_model
            .unwrap_or_else(|| Arc::new(WorkerDecisionSequenceLearn::new()));

        LearningDaemon::with_stores(
            self.config,
            trigger,
            record_store,
            episode_store,
            learn_model,
        )
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::events::{ActionContext, ActionEventBuilder, ActionEventResult};
    use crate::learn::trigger::AlwaysTrigger;
    use crate::types::WorkerId;

    fn make_test_records(count: usize) -> Vec<Record> {
        (0..count)
            .map(|i| {
                let event = ActionEventBuilder::new(i as u64, WorkerId(0), &format!("Action{}", i))
                    .result(ActionEventResult::success())
                    .duration(std::time::Duration::from_millis(10))
                    .context(ActionContext::new())
                    .build();
                Record::from(&event)
            })
            .collect()
    }

    #[test]
    fn test_daemon_config_builder() {
        let config = DaemonConfig::new("test")
            .data_dir("/tmp/test")
            .check_interval(Duration::from_secs(30))
            .processor_mode(ProcessorMode::Full)
            .auto_apply(true);

        assert_eq!(config.scenario, "test");
        assert_eq!(config.data_dir, PathBuf::from("/tmp/test"));
        assert_eq!(config.check_interval, Duration::from_secs(30));
        assert_eq!(config.processor_mode, ProcessorMode::Full);
        assert!(config.auto_apply);
    }

    #[tokio::test]
    async fn test_daemon_creation() {
        let config = DaemonConfig::new("test");
        let trigger = TriggerBuilder::never();

        let daemon = LearningDaemon::new(config, trigger).unwrap();
        assert_eq!(daemon.config().scenario, "test");
        assert_eq!(daemon.stats().records_received, 0);
    }

    #[tokio::test]
    async fn test_daemon_record_ingestion() {
        let config = DaemonConfig::new("test");
        let trigger = TriggerBuilder::never(); // Never auto-train

        let mut daemon = LearningDaemon::new(config, trigger).unwrap();
        let sender = daemon.record_sender();

        // Send records
        let records = make_test_records(5);
        sender.send(records).await.unwrap();

        // Process one batch
        daemon.handle_records(make_test_records(3)).await.unwrap();

        assert_eq!(daemon.stats().records_received, 3);
    }

    #[tokio::test]
    async fn test_daemon_builder() {
        let daemon = DaemonBuilder::new("test-scenario")
            .data_dir("/tmp/test")
            .trigger(Arc::new(AlwaysTrigger))
            .processor_mode(ProcessorMode::OfflineOnly)
            .build()
            .unwrap();

        assert_eq!(daemon.config().scenario, "test-scenario");
    }
}