swarm-engine-llm 0.1.6

//! PromptBuilder - ResolvedContext からプロンプトを生成
//!
//! Core層の構造化されたコンテキストからLLM用プロンプトを生成する。
//!
//! # 設計
//!
//! ```text
//! Core Layer                        LLM Layer
//!     │                                 │
//! ContextStore ──┐                      │
//!                │                      │
//! ContextView ───┼─→ ContextResolver ───┼─→ ResolvedContext
//!                                       │
//!                                       ▼
//!                              PromptBuilder.build()
//!                                       │
//!                                       ▼
//!                               String (prompt)
//! ```
//!
//! # Scope 対応
//!
//! ResolvedContext の内容に応じて、プロンプトのセクションを動的に構築する。
//!
//! - `self_last_output` がある場合 → "Last Result" セクションを出力
//! - `visible_workers` が空でない && 自分以外がいる場合 → "Team Status" セクションを出力
//! - `visible_workers` に自分がいる場合 → "Your Status" に詳細を出力
//!
//! これにより、WorkerScope::Minimal の場合は最小限のプロンプトになり、
//! WorkerScope::WithTeamDetail の場合は詳細なプロンプトになる。

use swarm_engine_core::agent::{
    ActionCandidate, ContextTarget, ManagerInstruction, ResolvedContext, WorkerCtx,
    WorkerDecisionRequest,
};
use swarm_engine_core::types::WorkerId;

use crate::json_prompt::action_selection_template;

// ============================================================================
// PromptBuilder
// ============================================================================

/// プロンプトビルダー
///
/// ResolvedContext から LLM 用プロンプトを生成する。
/// ResolvedContext の内容に応じてセクションを動的に構築する。
#[derive(Debug, Clone, Default)]
pub struct PromptBuilder;

impl PromptBuilder {
    /// 新規作成
    pub fn new() -> Self {
        Self
    }

    /// ResolvedContext からプロンプトを生成（Scope 対応）
    ///
    /// ResolvedContext の内容に応じて、必要なセクションのみを含むプロンプトを生成する。
    ///
    /// # セクション構成
    ///
    /// - Task: 常に出力（タスク説明、進捗）
    /// - Manager's Instruction: `manager_instruction` がある場合のみ
    /// - Last Result: `self_last_output` がある場合のみ
    /// - Your Status: `visible_workers` に自分がいる場合のみ
    /// - Team Status: `visible_workers` に自分以外がいる場合のみ
    /// - Available Actions: 常に出力
    /// - Instructions: 常に出力
    pub fn build(&self, context: &ResolvedContext) -> String {
        // コンテンツセクションを構築
        let content = self.build_content_sections(context);

        // 共通テンプレートを使用
        action_selection_template().build(&content)
    }

    /// コンテンツセクションを構築（テンプレートに渡す部分）
    fn build_content_sections(&self, context: &ResolvedContext) -> String {
        let mut sections = Vec::new();

        // Task（常に出力）
        sections.push(self.format_task(context));

        // Manager's Instruction（manager_instruction がある場合のみ）
        if let Some(ref instruction) = context.manager_instruction {
            sections.push(self.format_manager_instruction(instruction));
        }

        // Last Result（self_last_output がある場合のみ）
        if let Some(ref output) = context.self_last_output {
            sections.push(format!("## Last Result\n{}", output));
        }

        // Your Status / Team Status（visible_workers の内容に応じて）
        let worker_id = match &context.target {
            ContextTarget::Worker(id) => Some(*id),
            ContextTarget::Manager(_) => None,
        };

        if !context.visible_workers.is_empty() {
            // 自分の情報があれば Your Status を出力
            if let Some(wid) = worker_id {
                if let Some(my_ctx) = context.visible_workers.iter().find(|w| w.id == wid) {
                    sections.push(self.format_your_status(my_ctx));
                }
            }

            // 自分以外がいれば Team Status を出力
            let others: Vec<_> = context
                .visible_workers
                .iter()
                .filter(|w| worker_id != Some(w.id))
                .collect();

            if !others.is_empty() {
                sections.push(self.format_team_status(&others, &context.escalations));
            }
        }

        // Available Actions（常に出力）
        sections.push(self.format_candidates(&context.candidates));

        // Response format instruction
        sections.push("Required fields: tool, target, args, confidence".to_string());

        sections.join("\n\n")
    }

    /// ResolvedContext から WorkerDecisionRequest を生成
    ///
    /// 既存の LLMDecider との互換性のため。
    pub fn to_request(&self, context: &ResolvedContext) -> WorkerDecisionRequest {
        let worker_id = match context.target {
            ContextTarget::Worker(id) => id,
            ContextTarget::Manager(id) => WorkerId(id.0), // Manager の場合は ID を流用
        };

        // プロンプトを生成（query として使用）
        let prompt = self.build(context);

        WorkerDecisionRequest {
            worker_id,
            query: prompt,
            context: context.clone(),
            lora: None,
        }
    }

    // ========================================================================
    // Section Formatters
    // ========================================================================

    /// タスク情報をフォーマット
    fn format_task(&self, context: &ResolvedContext) -> String {
        let mut lines = Vec::new();
        lines.push("## Task".to_string());

        if let Some(ref task) = context.global.task_description {
            lines.push(task.clone());
        } else {
            lines.push("Continue current work".to_string());
        }

        // 進捗情報
        lines.push(format!(
            "Progress: {:.1}% | Tick: {}/{}",
            context.global.progress * 100.0,
            context.global.tick,
            context.global.max_ticks,
        ));

        if let Some(ref hint) = context.global.hint {
            lines.push(format!("Hint: {}", hint));
        }

        lines.join("\n")
    }

    /// Manager からの指示をフォーマット
    fn format_manager_instruction(&self, instruction: &ManagerInstruction) -> String {
        let mut lines = Vec::new();
        lines.push("## Manager's Instruction".to_string());

        // 指示テキスト
        if let Some(ref text) = instruction.instruction {
            lines.push(text.clone());
        }

        // 推奨アクション
        if let Some(ref action) = instruction.suggested_action {
            if let Some(ref target) = instruction.suggested_target {
                lines.push(format!("Suggested: {} -> {}", action, target));
            } else {
                lines.push(format!("Suggested: {}", action));
            }
        }

        // 探索ヒント
        if let Some(ref hint) = instruction.exploration_hint {
            lines.push(format!("Exploration: {}", hint));
        }

        lines.join("\n")
    }

    /// 自分のステータスをフォーマット（詳細版）
    fn format_your_status(&self, ctx: &WorkerCtx) -> String {
        let mut lines = Vec::new();
        lines.push("## Your Status".to_string());
        lines.push(format!("Worker {} (you)", ctx.id.0));

        // ステータス
        let status = if ctx.has_escalation {
            "ESCALATED"
        } else {
            "active"
        };
        lines.push(format!("  Status: {}", status));

        // 失敗カウント
        if ctx.consecutive_failures > 0 {
            lines.push(format!(
                "  Consecutive Failures: {} (consider different approach)",
                ctx.consecutive_failures
            ));
        }

        // 最新アクション
        if let Some(ref action) = ctx.last_action {
            let result = ctx
                .last_success
                .map_or("unknown", |s| if s { "SUCCESS" } else { "FAILED" });
            lines.push(format!("  Last Action: {} -> {}", action, result));
        }

        // last_output（metadata から取得）
        if let Some(output) = ctx.metadata.get("last_output") {
            if let Some(output_str) = output.as_str() {
                // 長すぎる場合は切り詰め
                let truncated = if output_str.len() > 500 {
                    format!("{}...(truncated)", &output_str[..500])
                } else {
                    output_str.to_string()
                };
                lines.push(format!("  Last Result: {}", truncated));
            }
        }

        // 履歴
        lines.push(format!("  Actions Taken: {}", ctx.history_len));

        lines.join("\n")
    }

    /// チームステータスをフォーマット
    fn format_team_status(
        &self,
        others: &[&WorkerCtx],
        escalations: &[(WorkerId, swarm_engine_core::state::Escalation)],
    ) -> String {
        let mut lines = Vec::new();
        lines.push("## Team Status".to_string());

        for ctx in others {
            lines.push(self.format_worker_brief(ctx));
        }

        // Escalation 情報を追加
        if !escalations.is_empty() {
            lines.push(String::new());
            lines.push("** Escalations (need attention) **".to_string());
            for (wid, esc) in escalations {
                lines.push(format!("  Worker {}: {:?}", wid.0, esc.reason));
            }
        }

        lines.join("\n")
    }

    /// Worker 簡易表示（チーム用）
    fn format_worker_brief(&self, ctx: &WorkerCtx) -> String {
        let status = if ctx.has_escalation { "ESC" } else { "ok" };
        let last = ctx.last_action.as_deref().unwrap_or("idle");
        let result = ctx.last_success.map_or("", |s| if s { "+" } else { "-" });

        format!(
            "Worker {}: [{}] last={}{} failures={}",
            ctx.id.0, status, last, result, ctx.consecutive_failures
        )
    }

    /// 候補をフォーマット（名前、説明、パラメータを含む）
    fn format_candidates(&self, candidates: &[ActionCandidate]) -> String {
        let mut lines = Vec::new();
        lines.push("## Available Actions".to_string());

        if candidates.is_empty() {
            lines.push("No actions available".to_string());
            return lines.join("\n");
        }

        for c in candidates {
            let params_str = if c.params.is_empty() {
                String::new()
            } else {
                let params: Vec<String> = c
                    .params
                    .iter()
                    .map(|p| {
                        let req = if p.required { " (required)" } else { "" };
                        format!("    - {}: {}{}", p.name, p.description, req)
                    })
                    .collect();
                format!("\n  params:\n{}", params.join("\n"))
            };
            let example_str = c
                .example
                .as_ref()
                .map_or(String::new(), |ex| format!("\n  example: {}", ex));
            lines.push(format!(
                "- {}: {}{}{}",
                c.name, c.description, params_str, example_str
            ));
        }

        lines.join("\n")
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use swarm_engine_core::agent::{ActionParam, GlobalContext};

    fn create_test_candidates() -> Vec<ActionCandidate> {
        vec![
            ActionCandidate {
                name: "Read".to_string(),
                description: "Read a file".to_string(),
                params: vec![ActionParam {
                    name: "path".to_string(),
                    description: "File path to read".to_string(),
                    required: true,
                }],
                example: None,
            },
            ActionCandidate {
                name: "Grep".to_string(),
                description: "Search for pattern".to_string(),
                params: vec![ActionParam {
                    name: "pattern".to_string(),
                    description: "Search pattern".to_string(),
                    required: true,
                }],
                example: None,
            },
        ]
    }

    fn create_minimal_context() -> ResolvedContext {
        let global = GlobalContext::new(10)
            .with_max_ticks(100)
            .with_progress(0.5)
            .with_task("Find the bug in authentication module");

        ResolvedContext::new(global, ContextTarget::Worker(WorkerId(0)))
            .with_self_last_output(Some("Found 3 files matching pattern".to_string()))
            .with_candidates(create_test_candidates())
    }

    fn create_detailed_context() -> ResolvedContext {
        let global = GlobalContext::new(10)
            .with_max_ticks(100)
            .with_progress(0.5)
            .with_task("Find the bug in authentication module");

        let worker0 = WorkerCtx::new(WorkerId(0))
            .with_last_action("read:src/auth.rs", true)
            .with_history_len(5);

        let worker1 = WorkerCtx::new(WorkerId(1))
            .with_last_action("grep:error", false)
            .with_failures(2)
            .with_escalation(true);

        ResolvedContext::new(global, ContextTarget::Worker(WorkerId(0)))
            .with_workers(vec![worker0, worker1])
            .with_candidates(create_test_candidates())
    }

    #[test]
    fn test_build_minimal_prompt() {
        let builder = PromptBuilder::new();
        let context = create_minimal_context();
        let prompt = builder.build(&context);

        // Task
        assert!(prompt.contains("## Task"));
        assert!(prompt.contains("Find the bug"));
        assert!(prompt.contains("Progress: 50.0%"));

        // Last Result（self_last_output があるので出力される）
        assert!(prompt.contains("## Last Result"));
        assert!(prompt.contains("Found 3 files matching pattern"));

        // Your Status（visible_workers が空なので出力されない）
        assert!(!prompt.contains("## Your Status"));

        // Team Status（visible_workers が空なので出力されない）
        assert!(!prompt.contains("## Team Status"));

        // Available Actions
        assert!(prompt.contains("## Available Actions"));
        assert!(prompt.contains("- Read: Read a file"));
    }

    #[test]
    fn test_build_detailed_prompt() {
        let builder = PromptBuilder::new();
        let context = create_detailed_context();
        let prompt = builder.build(&context);

        // Task
        assert!(prompt.contains("## Task"));

        // Your Status（visible_workers に自分がいるので出力される）
        assert!(prompt.contains("## Your Status"));
        assert!(prompt.contains("Worker 0 (you)"));

        // Team Status（自分以外がいるので出力される）
        assert!(prompt.contains("## Team Status"));
        assert!(prompt.contains("Worker 1"));
        assert!(prompt.contains("ESC")); // Escalation
    }

    #[test]
    fn test_to_request() {
        let builder = PromptBuilder::new();
        let context = create_minimal_context();
        let request = builder.to_request(&context);

        assert_eq!(request.worker_id, WorkerId(0));
        assert!(!request.query.is_empty());
        assert_eq!(request.context.candidates.len(), 2);
    }

    #[test]
    fn test_format_candidates_empty() {
        let builder = PromptBuilder::new();
        let result = builder.format_candidates(&[]);
        assert!(result.contains("No actions available"));
    }
}