swarm-engine-core 0.1.6

//! LoRA Trainer - LoRA 学習の実行
//!
//! ## 概要
//!
//! Episode → TrainingData → LoRA 学習 → TrainedModel
//!
//! ## 設計
//!
//! 現在は `lora/train.py` をサブプロセスとして呼び出す。
//! 将来的には Rust native 実装への移行も可能。
//!
//! ## 使用例
//!
//! ```ignore
//! use swarm_engine_core::learn::lora::{LoraTrainer, LoraTrainerConfig};
//! use swarm_engine_core::learn::EpisodeStore;
//!
//! let config = LoraTrainerConfig::default()
//!     .base_model("LiquidAI/LFM2.5-1.2B-Instruct")
//!     .lora_rank(16);
//!
//! let trainer = LoraTrainer::new(config, episode_store);
//! let model = trainer.train(&learn_model, None).await?;
//! ```

use std::io::Write as IoWrite;
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::sync::Arc;

use tokio::process::Command;

use crate::learn::episode::{EpisodeId, Outcome};
use crate::learn::learn_model::LearnModel;
use crate::learn::store::{EpisodeDto, EpisodeFilter, EpisodeStore, StoreError};
use crate::learn::training::TrainingData;
use crate::util::{epoch_millis, epoch_millis_for_ordering};

// ============================================================================
// LoraTrainerConfig
// ============================================================================

/// LoRA Trainer の設定
#[derive(Debug, Clone)]
pub struct LoraTrainerConfig {
    /// ベースモデル (HuggingFace ID or path)
    pub base_model: String,
    /// LoRA rank
    pub lora_rank: u32,
    /// LoRA alpha
    pub lora_alpha: f32,
    /// Dropout rate
    pub lora_dropout: f32,
    /// 学習エポック数
    pub epochs: u32,
    /// バッチサイズ
    pub batch_size: u32,
    /// 勾配蓄積ステップ
    pub gradient_accumulation: u32,
    /// 学習率
    pub learning_rate: f32,
    /// 最大シーケンス長
    pub max_seq_length: u32,
    /// train.py のパス
    pub train_script: PathBuf,
    /// 出力ディレクトリ（アダプタ保存先）
    pub output_dir: PathBuf,
    /// 学習データ一時ファイル
    pub data_dir: PathBuf,
    /// Python 実行パス
    pub python_path: PathBuf,
}

impl Default for LoraTrainerConfig {
    fn default() -> Self {
        Self {
            base_model: "LiquidAI/LFM2.5-1.2B-Instruct".to_string(),
            lora_rank: 16,
            lora_alpha: 32.0,
            lora_dropout: 0.05,
            epochs: 3,
            batch_size: 4,
            gradient_accumulation: 4,
            learning_rate: 2e-4,
            max_seq_length: 2048,
            train_script: PathBuf::from("lora/train.py"),
            output_dir: PathBuf::from("lora/adapters"),
            data_dir: PathBuf::from("lora/data"),
            python_path: PathBuf::from("python3"),
        }
    }
}

impl LoraTrainerConfig {
    /// ベースモデルを設定
    pub fn base_model(mut self, model: impl Into<String>) -> Self {
        self.base_model = model.into();
        self
    }

    /// LoRA rank を設定
    pub fn lora_rank(mut self, rank: u32) -> Self {
        self.lora_rank = rank;
        self
    }

    /// LoRA alpha を設定
    pub fn lora_alpha(mut self, alpha: f32) -> Self {
        self.lora_alpha = alpha;
        self
    }

    /// エポック数を設定
    pub fn epochs(mut self, epochs: u32) -> Self {
        self.epochs = epochs;
        self
    }

    /// バッチサイズを設定
    pub fn batch_size(mut self, size: u32) -> Self {
        self.batch_size = size;
        self
    }

    /// 学習率を設定
    pub fn learning_rate(mut self, lr: f32) -> Self {
        self.learning_rate = lr;
        self
    }

    /// train.py のパスを設定
    pub fn train_script(mut self, path: impl Into<PathBuf>) -> Self {
        self.train_script = path.into();
        self
    }

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

    /// Python パスを設定
    pub fn python_path(mut self, path: impl Into<PathBuf>) -> Self {
        self.python_path = path.into();
        self
    }
}

// ============================================================================
// TrainedModel
// ============================================================================

/// 学習済みモデル
#[derive(Debug, Clone)]
pub struct TrainedModel {
    /// モデル ID
    pub id: LoraModelId,
    /// ベースモデル
    pub base_model: String,
    /// アダプタのパス
    pub adapter_path: PathBuf,
    /// 使用した LearnModel の名前
    pub learn_model_name: String,
    /// 学習に使用した Episode の ID リスト
    pub episode_ids: Vec<EpisodeId>,
    /// 学習データサンプル数
    pub sample_count: usize,
    /// 作成日時（Unix timestamp ms）
    pub created_at: u64,
    /// 学習メトリクス
    pub metrics: Option<TrainingMetrics>,
}

/// LoRA モデル ID
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct LoraModelId(String);

impl LoraModelId {
    /// 新しい ID を生成
    pub fn new() -> Self {
        use std::sync::atomic::{AtomicU32, Ordering};
        static COUNTER: AtomicU32 = AtomicU32::new(0);

        let ts = epoch_millis_for_ordering();
        let counter = COUNTER.fetch_add(1, Ordering::Relaxed);
        Self(format!("lora-{}-{:08x}", ts, counter))
    }

    /// 文字列から生成
    pub fn parse(s: &str) -> Self {
        Self(s.to_string())
    }

    /// 文字列として取得
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

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

impl std::fmt::Display for LoraModelId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.0)
    }
}

/// 学習メトリクス
#[derive(Debug, Clone, Default)]
pub struct TrainingMetrics {
    /// 最終 loss
    pub final_loss: Option<f64>,
    /// 学習時間（秒）
    pub training_time_secs: Option<u64>,
    /// 使用した GPU メモリ（MB）
    pub gpu_memory_mb: Option<u64>,
}

// ============================================================================
// LoraTrainerError
// ============================================================================

/// LoRA Trainer エラー
#[derive(Debug)]
pub enum LoraTrainerError {
    /// Store エラー
    Store(StoreError),
    /// データが空
    EmptyData(String),
    /// IO エラー
    Io(std::io::Error),
    /// スクリプトが見つからない
    ScriptNotFound(PathBuf),
    /// 学習プロセスエラー
    ProcessFailed { exit_code: i32, stderr: String },
    /// その他
    Other(String),
}

impl std::fmt::Display for LoraTrainerError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Store(e) => write!(f, "Store error: {}", e),
            Self::EmptyData(msg) => write!(f, "Empty data: {}", msg),
            Self::Io(e) => write!(f, "IO error: {}", e),
            Self::ScriptNotFound(p) => write!(f, "Script not found: {}", p.display()),
            Self::ProcessFailed { exit_code, stderr } => {
                write!(f, "Training failed (exit {}): {}", exit_code, stderr)
            }
            Self::Other(msg) => write!(f, "{}", msg),
        }
    }
}

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

impl From<StoreError> for LoraTrainerError {
    fn from(e: StoreError) -> Self {
        Self::Store(e)
    }
}

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

// ============================================================================
// LoraTrainer
// ============================================================================

/// LoRA 学習を実行する
pub struct LoraTrainer {
    /// 設定
    config: LoraTrainerConfig,
    /// Episode ストア
    episode_store: Arc<dyn EpisodeStore>,
}

impl LoraTrainer {
    /// 新しい LoraTrainer を作成
    pub fn new(config: LoraTrainerConfig, episode_store: Arc<dyn EpisodeStore>) -> Self {
        Self {
            config,
            episode_store,
        }
    }

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

    /// Episode ストアを取得
    pub fn episode_store(&self) -> &Arc<dyn EpisodeStore> {
        &self.episode_store
    }

    /// LoRA 学習を実行（非同期）
    ///
    /// サブプロセスとして `train.py` を実行するため、長時間ブロックを回避。
    ///
    /// # Arguments
    /// * `learn_model` - Episode → TrainingData 変換を行う LearnModel
    /// * `filter` - Episode フィルタ（None で全件）
    ///
    /// # Returns
    /// 学習済みモデル
    pub async fn train(
        &self,
        learn_model: &dyn LearnModel,
        filter: Option<EpisodeFilter>,
    ) -> Result<TrainedModel, LoraTrainerError> {
        let started_at = std::time::Instant::now();

        // 1. Episode 取得
        tracing::info!(
            learn_model = learn_model.name(),
            "Fetching episodes for training"
        );
        let filter = filter.unwrap_or_default();
        let episodes = self.episode_store.query(&filter)?;

        if episodes.is_empty() {
            return Err(LoraTrainerError::EmptyData(
                "No episodes found for training".into(),
            ));
        }

        let episode_ids: Vec<_> = episodes.iter().map(|e| e.id.clone()).collect();
        tracing::info!(episode_count = episodes.len(), "Episodes fetched");

        // 2. TrainingData に変換
        tracing::info!("Converting episodes to training data");
        let training_data: Vec<TrainingData> = episodes
            .iter()
            .filter_map(|ep| episode_dto_to_training_data(ep, learn_model.name()).ok())
            .collect();

        if training_data.is_empty() {
            return Err(LoraTrainerError::EmptyData(
                "No training data generated from episodes".into(),
            ));
        }

        let sample_count = training_data.len();
        tracing::info!(sample_count, "Training data prepared");

        // 3. データをファイルに書き出し
        let data_path = self.write_training_data(&training_data, learn_model.name())?;
        tracing::info!(path = %data_path.display(), "Training data written");

        // 4. LoRA 学習実行（非同期）
        let timestamp = epoch_millis() / 1000; // 秒単位
        let adapter_name = format!("{}-{}", learn_model.name(), timestamp);
        let adapter_path = self.run_lora_training(&data_path, &adapter_name).await?;

        let elapsed = started_at.elapsed();
        tracing::info!(
            elapsed_secs = elapsed.as_secs(),
            adapter = %adapter_path.display(),
            "Training completed"
        );

        // 5. TrainedModel を作成
        let model = TrainedModel {
            id: LoraModelId::new(),
            base_model: self.config.base_model.clone(),
            adapter_path,
            learn_model_name: learn_model.name().to_string(),
            episode_ids,
            sample_count,
            created_at: epoch_millis(),
            metrics: Some(TrainingMetrics {
                final_loss: None, // TODO: train.py の出力からパース
                training_time_secs: Some(elapsed.as_secs()),
                gpu_memory_mb: None,
            }),
        };

        Ok(model)
    }

    /// 学習データをファイルに書き出し
    fn write_training_data(
        &self,
        data: &[TrainingData],
        learn_model_name: &str,
    ) -> Result<PathBuf, LoraTrainerError> {
        // 出力ディレクトリ作成
        std::fs::create_dir_all(&self.config.data_dir)?;

        let filename = format!("{}.jsonl", learn_model_name);
        let path = self.config.data_dir.join(filename);

        let mut file = std::fs::File::create(&path)?;

        for td in data {
            // TrainingData を学習用フォーマットに変換
            let json_str = training_data_to_json(td)?;
            writeln!(file, "{}", json_str)?;
        }

        Ok(path)
    }

    /// train.py を実行（非同期）
    async fn run_lora_training(
        &self,
        data_path: &Path,
        adapter_name: &str,
    ) -> Result<PathBuf, LoraTrainerError> {
        // スクリプト存在確認
        if !self.config.train_script.exists() {
            return Err(LoraTrainerError::ScriptNotFound(
                self.config.train_script.clone(),
            ));
        }

        let output_path = self.config.output_dir.join(adapter_name);

        // コマンド構築
        let mut cmd = Command::new(&self.config.python_path);
        cmd.arg(&self.config.train_script)
            .arg("--data")
            .arg(data_path)
            .arg("--output")
            .arg(&output_path)
            .arg("--model")
            .arg(&self.config.base_model)
            .arg("--rank")
            .arg(self.config.lora_rank.to_string())
            .arg("--alpha")
            .arg(self.config.lora_alpha.to_string())
            .arg("--dropout")
            .arg(self.config.lora_dropout.to_string())
            .arg("--epochs")
            .arg(self.config.epochs.to_string())
            .arg("--batch-size")
            .arg(self.config.batch_size.to_string())
            .arg("--grad-accum")
            .arg(self.config.gradient_accumulation.to_string())
            .arg("--lr")
            .arg(self.config.learning_rate.to_string())
            .arg("--max-seq-length")
            .arg(self.config.max_seq_length.to_string())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped());

        tracing::info!(
            script = %self.config.train_script.display(),
            data = %data_path.display(),
            output = %output_path.display(),
            "Starting LoRA training"
        );

        // 非同期実行
        let output = cmd.output().await?;

        if !output.status.success() {
            let stderr = String::from_utf8_lossy(&output.stderr);
            return Err(LoraTrainerError::ProcessFailed {
                exit_code: output.status.code().unwrap_or(-1),
                stderr: stderr.to_string(),
            });
        }

        // stdout をログ出力
        let stdout = String::from_utf8_lossy(&output.stdout);
        for line in stdout.lines() {
            tracing::debug!(line, "train.py output");
        }

        Ok(output_path)
    }
}

// ============================================================================
// Helpers
// ============================================================================

/// EpisodeDto → TrainingData 変換
fn episode_dto_to_training_data(
    dto: &EpisodeDto,
    learn_model_name: &str,
) -> Result<TrainingData, LoraTrainerError> {
    // システムプロンプト
    let system_prompt = format!(
        "You are an intelligent agent using the {} strategy. Your task is to make optimal decisions.",
        learn_model_name
    );

    // ユーザープロンプト（Episode のメタデータを含む）
    let user_prompt = format!(
        "Episode ID: {}\nLearn Model: {}\nMetadata: {:?}",
        dto.id, dto.learn_model, dto.metadata
    );

    // アシスタント応答（Outcome に基づく）
    let response = match &dto.outcome {
        Outcome::Success { score } => {
            format!("Decision successful with score {:.2}", score)
        }
        Outcome::Failure { reason } => {
            format!("Decision failed: {}", reason)
        }
        Outcome::Timeout { partial_score } => match partial_score {
            Some(score) => format!("Timeout with partial score {:.2}", score),
            None => "Timeout without progress".to_string(),
        },
        Outcome::Unknown => "Outcome unknown".to_string(),
    };

    // SFT 形式の TrainingData を作成
    let training_data = TrainingData::sft(&system_prompt, &user_prompt, &response)
        .with_episode_id(dto.id.to_string())
        .with_model(learn_model_name);

    // Outcome が Success なら score を追加
    let training_data = if let Outcome::Success { score } = &dto.outcome {
        training_data.with_outcome_score(*score)
    } else {
        training_data
    };

    Ok(training_data)
}

/// TrainingData → JSON 文字列（学習用フォーマット）
fn training_data_to_json(td: &TrainingData) -> Result<String, LoraTrainerError> {
    // train.py が期待するフォーマット: {"conversations": [...]}
    let conversation = td.to_conversation();

    let turns: Vec<serde_json::Value> = conversation
        .conversations
        .iter()
        .map(|turn| {
            serde_json::json!({
                "role": match turn.role {
                    crate::learn::training::ConversationRole::System => "system",
                    crate::learn::training::ConversationRole::User => "user",
                    crate::learn::training::ConversationRole::Assistant => "assistant",
                },
                "content": turn.content,
            })
        })
        .collect();

    let json_value = serde_json::json!({
        "conversations": turns
    });

    serde_json::to_string(&json_value)
        .map_err(|e| LoraTrainerError::Other(format!("JSON serialization error: {}", e)))
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::learn::store::InMemoryEpisodeStore;

    #[test]
    fn test_trainer_config_builder() {
        let config = LoraTrainerConfig::default()
            .base_model("test-model")
            .lora_rank(32)
            .lora_alpha(64.0)
            .epochs(5)
            .batch_size(8)
            .learning_rate(1e-4);

        assert_eq!(config.base_model, "test-model");
        assert_eq!(config.lora_rank, 32);
        assert_eq!(config.lora_alpha, 64.0);
        assert_eq!(config.epochs, 5);
        assert_eq!(config.batch_size, 8);
        assert!((config.learning_rate - 1e-4).abs() < 1e-10);
    }

    #[test]
    fn test_model_id() {
        let id1 = LoraModelId::new();
        let id2 = LoraModelId::new();

        // IDs should be unique (different timestamps or rand)
        // Note: In fast tests, they might have same timestamp but different rand
        assert!(!id1.as_str().is_empty());
        assert!(!id2.as_str().is_empty());
    }

    #[test]
    fn test_trainer_creation() {
        let config = LoraTrainerConfig::default();
        let store = Arc::new(InMemoryEpisodeStore::new());
        let trainer = LoraTrainer::new(config, store);

        assert_eq!(trainer.config().base_model, "LiquidAI/LFM2.5-1.2B-Instruct");
        assert_eq!(trainer.config().lora_rank, 16);
    }

    #[tokio::test]
    async fn test_train_empty_store() {
        use crate::learn::learn_model::WorkerTaskLearn;

        let config = LoraTrainerConfig::default();
        let store = Arc::new(InMemoryEpisodeStore::new());
        let trainer = LoraTrainer::new(config, store);

        let learn_model = WorkerTaskLearn::new();
        let result = trainer.train(&learn_model, None).await;

        assert!(result.is_err());
        match result {
            Err(LoraTrainerError::EmptyData(_)) => {}
            _ => panic!("Expected EmptyData error"),
        }
    }

    #[test]
    fn test_episode_dto_to_training_data() {
        use crate::learn::episode::EpisodeMetadata;

        let dto = EpisodeDto {
            id: EpisodeId::new(),
            learn_model: "test".to_string(),
            outcome: Outcome::success(0.95),
            metadata: EpisodeMetadata::new(),
            record_ids: vec![],
        };

        let td = episode_dto_to_training_data(&dto, "test-model").unwrap();
        assert!(td.is_sft());
    }

    #[test]
    fn test_training_data_to_json() {
        let td = TrainingData::sft(
            "You are a helpful assistant.",
            "What is 2+2?",
            "2+2 equals 4.",
        );

        let json = training_data_to_json(&td).unwrap();
        assert!(json.contains("conversations"));
        assert!(json.contains("system"));
        assert!(json.contains("user"));
        assert!(json.contains("assistant"));
    }
}