swarm_engine_eval/environments/

search.rs

//! SearchEnvironment - ファイル検索環境
//!
//! 複数の検索結果から正解を見つけるシナリオ。
//! SelectionLogic が複数の frontier から選択する動作を検証するための環境。
//!
//! # パターン
//!
//! ```text
//! SearchFiles → [5 files] → ReadFile → (Success/Fail) → Analyze → Done
//! ```
//!
//! # 設計
//!
//! - `SearchFiles`: 検索を実行し、複数のファイルパスを返す
//! - `ReadFile(target)`: 指定されたファイルを読む（正解は1つだけ成功）
//! - `Analyze`: 正解ファイルの内容を分析して完了
//!
//! これにより、SearchFiles 実行後に複数の frontier が生成され、
//! SelectionLogic がスコアベースで選択する機会が生まれる。

use std::collections::HashSet;
use std::sync::RwLock;

use swarm_engine_core::agent::WorkResult;
use swarm_engine_core::environment::Environment;
use swarm_engine_core::types::{Action, WorkerId};

// ============================================================================
// SearchEnvironment
// ============================================================================

/// ファイル検索環境
///
/// 検索結果から正解ファイルを見つけるシナリオ。
pub struct SearchEnvironment {
    /// 検索結果として返すファイル一覧
    files: Vec<String>,
    /// 正解ファイル（これだけが成功する）
    target_file: String,
    /// 正解ファイルの内容
    target_content: String,
    /// 内部状態
    state: RwLock<SearchState>,
}

#[derive(Debug, Default)]
struct SearchState {
    /// 検索を実行したか
    searched: bool,
    /// 読んだファイル一覧
    read_files: HashSet<String>,
    /// 正解ファイルを読んだか
    found_target: bool,
    /// 完了した Worker
    completed: Vec<WorkerId>,
}

impl SearchEnvironment {
    /// 新しい環境を作成
    pub fn new(
        files: Vec<String>,
        target_file: impl Into<String>,
        target_content: impl Into<String>,
    ) -> Self {
        Self {
            files,
            target_file: target_file.into(),
            target_content: target_content.into(),
            state: RwLock::new(SearchState::default()),
        }
    }

    /// 基本シナリオ: 5ファイルから1つの正解を探す
    pub fn basic_scenario() -> Self {
        let files = vec![
            "src/config.rs".to_string(),
            "src/utils.rs".to_string(),
            "src/handler.rs".to_string(), // target
            "src/types.rs".to_string(),
            "src/error.rs".to_string(),
        ];
        Self::new(
            files,
            "src/handler.rs",
            "fn handle_request() { /* implementation */ }",
        )
    }

    /// 中規模シナリオ: 10ファイルから1つの正解を探す
    pub fn medium_scenario() -> Self {
        let files = vec![
            "src/lib.rs".to_string(),
            "src/config.rs".to_string(),
            "src/utils.rs".to_string(),
            "src/handler.rs".to_string(),
            "src/types.rs".to_string(),
            "src/error.rs".to_string(),
            "src/api/mod.rs".to_string(),
            "src/api/routes.rs".to_string(), // target
            "src/api/middleware.rs".to_string(),
            "src/db/connection.rs".to_string(),
        ];
        Self::new(
            files,
            "src/api/routes.rs",
            "pub fn register_routes(app: &mut App) { /* routes */ }",
        )
    }

    /// 大規模シナリオ: 20ファイルから1つの正解を探す
    pub fn large_scenario() -> Self {
        let files = vec![
            "src/lib.rs".to_string(),
            "src/main.rs".to_string(),
            "src/config.rs".to_string(),
            "src/utils.rs".to_string(),
            "src/handler.rs".to_string(),
            "src/types.rs".to_string(),
            "src/error.rs".to_string(),
            "src/api/mod.rs".to_string(),
            "src/api/routes.rs".to_string(),
            "src/api/middleware.rs".to_string(),
            "src/db/mod.rs".to_string(),
            "src/db/connection.rs".to_string(),
            "src/db/queries.rs".to_string(),
            "src/services/mod.rs".to_string(),
            "src/services/auth.rs".to_string(), // target
            "src/services/user.rs".to_string(),
            "src/services/payment.rs".to_string(),
            "src/models/mod.rs".to_string(),
            "src/models/user.rs".to_string(),
            "src/models/order.rs".to_string(),
        ];
        Self::new(
            files,
            "src/services/auth.rs",
            "pub fn authenticate(token: &str) -> Result<User> { /* auth logic */ }",
        )
    }

    /// カスタムシナリオ: ファイル数と正解位置を指定
    pub fn custom_scenario(file_count: usize, target_index: usize, seed: u64) -> Self {
        let file_count = file_count.clamp(2, 50);
        let target_index = target_index.min(file_count - 1);

        // 簡易乱数生成
        let mut rng_state = seed;
        let mut next_rand = || {
            rng_state = rng_state.wrapping_mul(6364136223846793005).wrapping_add(1);
            rng_state
        };

        let prefixes = ["src", "lib", "core", "api", "services", "models", "utils"];
        let suffixes = [
            "mod", "types", "error", "handler", "config", "utils", "impl",
        ];

        let files: Vec<String> = (0..file_count)
            .map(|i| {
                let prefix = prefixes[(next_rand() as usize) % prefixes.len()];
                let suffix = suffixes[(next_rand() as usize) % suffixes.len()];
                format!("{}/{}_{}.rs", prefix, suffix, i)
            })
            .collect();

        let target_file = files[target_index].clone();

        Self::new(
            files,
            target_file,
            "// TARGET FILE CONTENT\nfn target_function() { /* found! */ }",
        )
    }

    // ------------------------------------------------------------------------
    // Action Handlers
    // ------------------------------------------------------------------------

    fn handle_search_files(&self, _worker_id: WorkerId, action: &Action) -> WorkResult {
        let query = action
            .params
            .args
            .get("query")
            .or(action.params.target.as_ref())
            .cloned()
            .unwrap_or_else(|| "*.rs".to_string());

        let mut state = self.state.write().unwrap();
        state.searched = true;

        // JSON 配列として返す（discovery として認識される）
        let files_json: Vec<serde_json::Value> = self
            .files
            .iter()
            .map(|f| serde_json::Value::String(f.clone()))
            .collect();

        let result = serde_json::json!({
            "query": query,
            "count": self.files.len(),
            "files": files_json,
            "message": "Use ReadFile to examine each file."
        });

        WorkResult::env_success_structured(result)
    }

    fn handle_read_file(&self, _worker_id: WorkerId, action: &Action) -> WorkResult {
        let file_path = action
            .params
            .args
            .get("file")
            .or(action.params.target.as_ref())
            .cloned()
            .unwrap_or_default();

        if file_path.is_empty() {
            return WorkResult::env_failure("ReadFile requires a file path");
        }

        let mut state = self.state.write().unwrap();

        // ファイルが検索結果に含まれているか確認
        // Note: searched チェックを緩和し、initial_context からの直接 ReadFile を許可
        if !self.files.contains(&file_path) {
            return WorkResult::env_failure(format!(
                "File '{}' not found. Available files: {:?}",
                file_path, self.files
            ));
        }

        state.read_files.insert(file_path.clone());

        // 正解ファイルかどうか
        if file_path == self.target_file {
            state.found_target = true;
            WorkResult::env_success(format!(
                "=== {} ===\n{}\n\n[TARGET FOUND] This file contains the target content!",
                file_path, self.target_content
            ))
        } else {
            // 不正解ファイル - 失敗として記録
            WorkResult::env_failure(format!(
                "=== {} ===\n// Empty or irrelevant content\n// Not the target file",
                file_path
            ))
        }
    }

    fn handle_analyze(&self, worker_id: WorkerId, action: &Action) -> WorkResult {
        let file_path = action
            .params
            .args
            .get("file")
            .or(action.params.target.as_ref())
            .cloned()
            .unwrap_or_default();

        if file_path.is_empty() {
            return WorkResult::env_failure("Analyze requires a file path");
        }

        let mut state = self.state.write().unwrap();

        // 正解ファイルを見つけていない場合はエラー
        if !state.found_target {
            return WorkResult::env_failure("Cannot analyze without finding the target file first. Read files to find the target.");
        }

        // 正解ファイルを分析しているか
        if file_path != self.target_file {
            return WorkResult::env_failure(format!(
                "Cannot analyze '{}'. The target file is different.",
                file_path
            ));
        }

        // 成功！
        if !state.completed.contains(&worker_id) {
            state.completed.push(worker_id);
        }

        WorkResult::done_success(format!(
            "=== Analysis Complete ===\nFile: {}\nContent analyzed successfully!\n\nTask completed!",
            file_path
        ))
    }
}

impl Environment for SearchEnvironment {
    fn step(&self, worker_id: WorkerId, action: &Action) -> WorkResult {
        match action.name.to_lowercase().as_str() {
            "searchfiles" | "search_files" | "search" | "list" | "listfiles" | "list_files" => {
                self.handle_search_files(worker_id, action)
            }
            "readfile" | "read_file" | "read" | "cat" => self.handle_read_file(worker_id, action),
            "analyze" | "process" | "complete" => self.handle_analyze(worker_id, action),
            "continue" => WorkResult::env_success("Continuing..."),
            _ => WorkResult::unsupported(&action.name),
        }
    }

    fn reset(&self) {
        let mut state = self.state.write().unwrap();
        state.searched = false;
        state.read_files.clear();
        state.found_target = false;
        state.completed.clear();
    }

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

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

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashMap;

    fn is_success(result: &WorkResult) -> bool {
        match result {
            WorkResult::Acted { action_result, .. } => action_result.success,
            WorkResult::Done { success, .. } => *success,
            _ => false,
        }
    }

    fn is_done(result: &WorkResult) -> bool {
        matches!(result, WorkResult::Done { .. })
    }

    fn action(name: &str, target: Option<&str>) -> Action {
        Action {
            name: name.into(),
            params: swarm_engine_core::types::ActionParams {
                target: target.map(|s| s.into()),
                args: HashMap::new(),
                data: vec![],
            },
        }
    }

    #[test]
    fn test_search_files() {
        let env = SearchEnvironment::basic_scenario();
        let worker = WorkerId(0);

        let result = env.step(worker, &action("SearchFiles", None));
        assert!(is_success(&result));
        assert!(!is_done(&result));
    }

    #[test]
    fn test_read_file_without_search() {
        let env = SearchEnvironment::basic_scenario();
        let worker = WorkerId(0);

        // 検索なしでも ReadFile は動作する（initial_context からの直接アクセス対応）
        // ただし、正解ファイルのみ成功
        let result = env.step(worker, &action("ReadFile", Some("src/handler.rs")));
        assert!(is_success(&result)); // 正解ファイルは成功

        // 不正解ファイルは失敗
        let env2 = SearchEnvironment::basic_scenario();
        let result2 = env2.step(worker, &action("ReadFile", Some("src/config.rs")));
        assert!(!is_success(&result2));
    }

    #[test]
    fn test_read_target_file_success() {
        let env = SearchEnvironment::basic_scenario();
        let worker = WorkerId(0);

        // 検索
        env.step(worker, &action("SearchFiles", None));

        // 正解ファイルを読む
        let result = env.step(worker, &action("ReadFile", Some("src/handler.rs")));
        assert!(is_success(&result));
    }

    #[test]
    fn test_read_wrong_file_failure() {
        let env = SearchEnvironment::basic_scenario();
        let worker = WorkerId(0);

        // 検索
        env.step(worker, &action("SearchFiles", None));

        // 不正解ファイルを読む
        let result = env.step(worker, &action("ReadFile", Some("src/config.rs")));
        assert!(!is_success(&result)); // 失敗
    }

    #[test]
    fn test_full_search_flow() {
        let env = SearchEnvironment::basic_scenario();
        let worker = WorkerId(0);

        // 1. SearchFiles
        let result = env.step(worker, &action("SearchFiles", None));
        assert!(is_success(&result));
        assert!(!is_done(&result));

        // 2. ReadFile (wrong) - 失敗だがフローは続行可能
        let result = env.step(worker, &action("ReadFile", Some("src/config.rs")));
        assert!(!is_success(&result));
        assert!(!is_done(&result));

        // 3. ReadFile (target) - 成功
        let result = env.step(worker, &action("ReadFile", Some("src/handler.rs")));
        assert!(is_success(&result));
        assert!(!is_done(&result));

        // 4. Analyze - 完了！
        let result = env.step(worker, &action("Analyze", Some("src/handler.rs")));
        assert!(is_success(&result));
        assert!(is_done(&result));
    }

    #[test]
    fn test_analyze_requires_found_target() {
        let env = SearchEnvironment::basic_scenario();
        let worker = WorkerId(0);

        // 検索
        env.step(worker, &action("SearchFiles", None));

        // 正解を見つけずにAnalyzeしようとするとエラー
        let result = env.step(worker, &action("Analyze", Some("src/handler.rs")));
        assert!(!is_success(&result));
    }

    #[test]
    fn test_custom_scenario() {
        let env = SearchEnvironment::custom_scenario(10, 3, 42);
        let worker = WorkerId(0);

        // 検索
        let result = env.step(worker, &action("SearchFiles", None));
        assert!(is_success(&result));
    }
}