swarm_engine_core/exploration/

dependency_planner.rs

//! Action Dependency Planner - アクション依存グラフの動的構築
//!
//! タスク開始時に LLM を使ってアクション間の依存関係を一度だけ計画する。
//! 毎 tick で LLM を呼ぶ代わりに、グラフに従って探索を進める。
//!
//! # コンセプト
//!
//! ```text
//! Task + AvailableActions → [LLM] → DependencyGraph
//!
//! DependencyGraph:
//!   - edges: [(Grep, Read, 0.95), (List, Grep, 0.60)]
//!   - start_nodes: [Grep]
//!   - terminal_nodes: [Read]
//!
//! 実行時:
//!   Worker は Graph に従うだけ → 毎 tick の LLM 不要
//! ```
//!
//! # 利点
//!
//! 1. **軽量 Meta-Planning**: タスク開始時に1回だけ LLM を呼ぶ
//! 2. **構造的な制約**: 逆流（Read → Grep）が起きない
//! 3. **高精度**: パターンが決まっているから LLM の判断が簡単
//!
//! # 使用例
//!
//! ```ignore
//! let planner = ActionDependencyPlanner::new();
//!
//! // タスク開始時に依存グラフを生成
//! let graph = planner.plan(
//!     "src/auth.rs の verify_token 関数を見つける",
//!     &["Grep", "List", "Read"],
//!     &llm_client,
//! ).await?;
//!
//! // 探索時はグラフに従う
//! let next_actions = graph.valid_next_actions("Grep"); // → ["Read"]
//! ```

use std::collections::{HashMap, HashSet};

use serde::{Deserialize, Serialize};

// ============================================================================
// Core Types
// ============================================================================

/// アクション間の依存エッジ
///
/// `from` アクションの成功後に `to` アクションを実行可能。
/// `confidence` は LLM の確信度（0.0〜1.0）。
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DependencyEdge {
    /// 元アクション名
    pub from: String,
    /// 先アクション名
    pub to: String,
    /// 確信度（0.0〜1.0）
    pub confidence: f64,
}

impl DependencyEdge {
    pub fn new(from: impl Into<String>, to: impl Into<String>, confidence: f64) -> Self {
        Self {
            from: from.into(),
            to: to.into(),
            confidence: confidence.clamp(0.0, 1.0),
        }
    }
}

/// アクション依存グラフ
///
/// タスクに対して LLM が生成したアクション間の依存関係。
/// 探索時はこのグラフに従ってアクションを選択する。
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct DependencyGraph {
    /// 依存エッジ（from → to の関係）
    edges: Vec<DependencyEdge>,
    /// 開始可能なアクション（依存なしで実行可能）
    start_nodes: HashSet<String>,
    /// 終端アクション（これが成功したらタスク完了の可能性）
    terminal_nodes: HashSet<String>,
    /// タスク説明（デバッグ用）
    task: String,
    /// 利用可能なアクション一覧（検証用）
    available_actions: Vec<String>,
    /// アクションごとのパラメータバリアント（action → (key, values)）
    ///
    /// 例: "Move" → ("target", ["north", "south", "east", "west"])
    #[serde(default)]
    param_variants: HashMap<String, (String, Vec<String>)>,
    /// Discover（NodeExpand）アクションの順序（キャッシュ用）
    #[serde(default)]
    discover_order: Vec<String>,
    /// NotDiscover（NodeStateChange）アクションの順序（キャッシュ用）
    #[serde(default)]
    not_discover_order: Vec<String>,
    /// 学習用記録（LLM 推論を学習データとして保存）
    #[serde(skip)]
    learn_record: Option<crate::learn::DependencyGraphRecord>,
}

impl DependencyGraph {
    /// 空のグラフを作成
    pub fn new() -> Self {
        Self::default()
    }

    /// Builder パターンでグラフを構築
    pub fn builder() -> DependencyGraphBuilder {
        DependencyGraphBuilder::new()
    }

    // ------------------------------------------------------------------------
    // Query API
    // ------------------------------------------------------------------------

    /// 指定アクションの後に実行可能なアクションを取得
    ///
    /// 確信度順にソートして返す。
    pub fn valid_next_actions(&self, current_action: &str) -> Vec<String> {
        let mut edges: Vec<_> = self
            .edges
            .iter()
            .filter(|e| e.from == current_action)
            .collect();

        // 確信度が高い順にソート（NaN対策で unwrap_or を使用）
        edges.sort_by(|a, b| {
            b.confidence
                .partial_cmp(&a.confidence)
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        edges.iter().map(|e| e.to.clone()).collect()
    }

    /// 開始可能なアクションを取得
    ///
    /// アルファベット順にソートして一貫した順序を保証。
    pub fn start_actions(&self) -> Vec<String> {
        let mut actions: Vec<_> = self.start_nodes.iter().cloned().collect();
        actions.sort();
        actions
    }

    /// 終端アクションを取得
    ///
    /// アルファベット順にソートして一貫した順序を保証。
    pub fn terminal_actions(&self) -> Vec<String> {
        let mut actions: Vec<_> = self.terminal_nodes.iter().cloned().collect();
        actions.sort();
        actions
    }

    /// 指定アクションが終端か
    pub fn is_terminal(&self, action: &str) -> bool {
        self.terminal_nodes.contains(action)
    }

    /// 指定アクションが開始ノードか
    pub fn is_start(&self, action: &str) -> bool {
        self.start_nodes.contains(action)
    }

    /// 指定アクションから指定アクションへの遷移が有効か
    pub fn can_transition(&self, from: &str, to: &str) -> bool {
        self.edges.iter().any(|e| e.from == from && e.to == to)
    }

    /// 遷移の確信度を取得
    pub fn transition_confidence(&self, from: &str, to: &str) -> Option<f64> {
        self.edges
            .iter()
            .find(|e| e.from == from && e.to == to)
            .map(|e| e.confidence)
    }

    /// 全エッジを取得
    pub fn edges(&self) -> &[DependencyEdge] {
        &self.edges
    }

    /// タスク説明を取得
    pub fn task(&self) -> &str {
        &self.task
    }

    /// 利用可能なアクション一覧を取得
    pub fn available_actions(&self) -> &[String] {
        &self.available_actions
    }

    /// 指定アクションのパラメータバリアントを取得
    pub fn param_variants(&self, action: &str) -> Option<(&str, &[String])> {
        self.param_variants
            .get(action)
            .map(|(key, values)| (key.as_str(), values.as_slice()))
    }

    /// 全パラメータバリアントを取得
    pub fn all_param_variants(&self) -> &HashMap<String, (String, Vec<String>)> {
        &self.param_variants
    }

    // ------------------------------------------------------------------------
    // Action Order (for caching)
    // ------------------------------------------------------------------------

    /// Discover アクションの順序を取得
    pub fn discover_order(&self) -> &[String] {
        &self.discover_order
    }

    /// NotDiscover アクションの順序を取得
    pub fn not_discover_order(&self) -> &[String] {
        &self.not_discover_order
    }

    /// アクション順序を設定
    pub fn set_action_order(&mut self, discover: Vec<String>, not_discover: Vec<String>) {
        self.discover_order = discover;
        self.not_discover_order = not_discover;
    }

    /// アクション順序が設定されているか
    pub fn has_action_order(&self) -> bool {
        !self.discover_order.is_empty() || !self.not_discover_order.is_empty()
    }

    // ------------------------------------------------------------------------
    // Learning Record
    // ------------------------------------------------------------------------

    /// 学習用記録を設定
    pub fn set_learn_record(&mut self, record: crate::learn::DependencyGraphRecord) {
        self.learn_record = Some(record);
    }

    /// 学習用記録を取得
    pub fn learn_record(&self) -> Option<&crate::learn::DependencyGraphRecord> {
        self.learn_record.as_ref()
    }

    /// 学習用記録を取り出し（所有権移動）
    pub fn take_learn_record(&mut self) -> Option<crate::learn::DependencyGraphRecord> {
        self.learn_record.take()
    }

    // ------------------------------------------------------------------------
    // Validation
    // ------------------------------------------------------------------------

    /// グラフが有効か検証
    ///
    /// - 少なくとも1つの開始ノードがある
    /// - 少なくとも1つの終端ノードがある
    /// - 全エッジのアクションが available_actions に含まれる
    /// - 全開始ノードが available_actions に含まれる
    /// - 全終端ノードが available_actions に含まれる
    pub fn validate(&self) -> Result<(), DependencyGraphError> {
        if self.start_nodes.is_empty() {
            return Err(DependencyGraphError::NoStartNodes);
        }

        if self.terminal_nodes.is_empty() {
            return Err(DependencyGraphError::NoTerminalNodes);
        }

        // 開始ノードが available_actions に含まれるか
        for node in &self.start_nodes {
            if !self.available_actions.contains(node) {
                return Err(DependencyGraphError::UnknownAction(node.clone()));
            }
        }

        // 終端ノードが available_actions に含まれるか
        for node in &self.terminal_nodes {
            if !self.available_actions.contains(node) {
                return Err(DependencyGraphError::UnknownAction(node.clone()));
            }
        }

        // エッジのアクションが available_actions に含まれるか
        for edge in &self.edges {
            if !self.available_actions.contains(&edge.from) {
                return Err(DependencyGraphError::UnknownAction(edge.from.clone()));
            }
            if !self.available_actions.contains(&edge.to) {
                return Err(DependencyGraphError::UnknownAction(edge.to.clone()));
            }
        }

        Ok(())
    }

    // ------------------------------------------------------------------------
    // Visualization
    // ------------------------------------------------------------------------

    /// グラフを Mermaid 形式で出力
    pub fn to_mermaid(&self) -> String {
        let mut lines = vec!["graph LR".to_string()];

        for edge in &self.edges {
            let label = format!("{:.0}%", edge.confidence * 100.0);
            lines.push(format!("    {} -->|{}| {}", edge.from, label, edge.to));
        }

        // 開始ノードをスタイル
        for start in &self.start_nodes {
            lines.push(format!("    style {} fill:#9f9", start));
        }

        // 終端ノードをスタイル
        for terminal in &self.terminal_nodes {
            lines.push(format!("    style {} fill:#f99", terminal));
        }

        lines.join("\n")
    }
}

/// グラフ構築エラー
#[derive(Debug, Clone, thiserror::Error)]
pub enum DependencyGraphError {
    #[error("No start nodes defined")]
    NoStartNodes,

    #[error("No terminal nodes defined")]
    NoTerminalNodes,

    #[error("Unknown action: {0}")]
    UnknownAction(String),

    #[error("Parse error: {0}")]
    ParseError(String),

    #[error("LLM error: {0}")]
    LlmError(String),
}

// ============================================================================
// Provider Trait
// ============================================================================

/// DependencyGraph を提供するトレイト
///
/// このトレイトを実装することで、様々な方法で DependencyGraph を生成できる:
/// - LLM による動的生成（BatchInvoker が実装）
/// - 静的な事前定義グラフ
/// - テスト用のモック
///
/// # Example
///
/// ```ignore
/// use swarm_engine_core::exploration::{DependencyGraph, DependencyGraphProvider};
///
/// struct StaticProvider {
///     graph: DependencyGraph,
/// }
///
/// impl DependencyGraphProvider for StaticProvider {
///     fn provide_graph(&self, _task: &str, _actions: &[String]) -> Option<DependencyGraph> {
///         Some(self.graph.clone())
///     }
/// }
/// ```
pub trait DependencyGraphProvider: Send + Sync {
    /// タスクとアクション一覧から DependencyGraph を生成
    ///
    /// # Arguments
    /// - `task`: タスクの説明（goal）
    /// - `available_actions`: 利用可能なアクション名の一覧
    ///
    /// # Returns
    /// - `Some(DependencyGraph)`: グラフが生成された場合
    /// - `None`: 生成をスキップする場合（LLM エラー等）
    fn provide_graph(&self, task: &str, available_actions: &[String]) -> Option<DependencyGraph>;
}

// ============================================================================
// Builder
// ============================================================================

/// DependencyGraph のビルダー
#[derive(Debug, Clone, Default)]
pub struct DependencyGraphBuilder {
    edges: Vec<DependencyEdge>,
    start_nodes: HashSet<String>,
    terminal_nodes: HashSet<String>,
    task: String,
    available_actions: Vec<String>,
    param_variants: HashMap<String, (String, Vec<String>)>,
}

impl DependencyGraphBuilder {
    pub fn new() -> Self {
        Self::default()
    }

    /// タスク説明を設定
    pub fn task(mut self, task: impl Into<String>) -> Self {
        self.task = task.into();
        self
    }

    /// 利用可能なアクションを設定
    pub fn available_actions<I, S>(mut self, actions: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        self.available_actions = actions.into_iter().map(|s| s.into()).collect();
        self
    }

    /// エッジを追加
    pub fn edge(mut self, from: impl Into<String>, to: impl Into<String>, confidence: f64) -> Self {
        self.edges.push(DependencyEdge::new(from, to, confidence));
        self
    }

    /// 開始ノードを追加
    pub fn start_node(mut self, action: impl Into<String>) -> Self {
        self.start_nodes.insert(action.into());
        self
    }

    /// 複数の開始ノードを追加
    pub fn start_nodes<I, S>(mut self, actions: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        self.start_nodes
            .extend(actions.into_iter().map(|s| s.into()));
        self
    }

    /// 終端ノードを追加
    pub fn terminal_node(mut self, action: impl Into<String>) -> Self {
        self.terminal_nodes.insert(action.into());
        self
    }

    /// 複数の終端ノードを追加
    pub fn terminal_nodes<I, S>(mut self, actions: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        self.terminal_nodes
            .extend(actions.into_iter().map(|s| s.into()));
        self
    }

    /// パラメータバリアントを追加
    ///
    /// 指定アクションの後続ノード展開時に、各バリアントごとにノードを生成する。
    pub fn param_variants<I, S>(
        mut self,
        action: impl Into<String>,
        key: impl Into<String>,
        values: I,
    ) -> Self
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        self.param_variants.insert(
            action.into(),
            (key.into(), values.into_iter().map(|s| s.into()).collect()),
        );
        self
    }

    /// グラフを構築
    pub fn build(self) -> DependencyGraph {
        DependencyGraph {
            edges: self.edges,
            start_nodes: self.start_nodes,
            terminal_nodes: self.terminal_nodes,
            task: self.task,
            available_actions: self.available_actions,
            param_variants: self.param_variants,
            discover_order: Vec::new(),
            not_discover_order: Vec::new(),
            learn_record: None,
        }
    }

    /// グラフを構築してバリデーション
    pub fn build_validated(self) -> Result<DependencyGraph, DependencyGraphError> {
        let graph = self.build();
        graph.validate()?;
        Ok(graph)
    }
}

// ============================================================================
// LLM Response Parsing
// ============================================================================

/// LLM からの依存グラフ生成レスポンス
///
/// LLM が JSON で返すことを想定。
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LlmDependencyResponse {
    /// 依存エッジ
    pub edges: Vec<LlmEdge>,
    /// 開始アクション
    pub start: Vec<String>,
    /// 終端アクション
    pub terminal: Vec<String>,
    /// 推論理由（オプション）
    #[serde(default)]
    pub reasoning: Option<String>,
}

/// LLM レスポンス内のエッジ
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LlmEdge {
    pub from: String,
    pub to: String,
    pub confidence: f64,
}

impl LlmDependencyResponse {
    /// DependencyGraph に変換
    pub fn into_graph(
        self,
        task: impl Into<String>,
        available_actions: Vec<String>,
    ) -> DependencyGraph {
        let mut builder = DependencyGraphBuilder::new()
            .task(task)
            .available_actions(available_actions)
            .start_nodes(self.start)
            .terminal_nodes(self.terminal);

        for edge in self.edges {
            builder = builder.edge(edge.from, edge.to, edge.confidence);
        }

        builder.build()
    }

    /// テキストからパース（矢印形式 or JSON）
    ///
    /// 優先順位:
    /// 1. 矢印形式（A -> B -> C）
    /// 2. JSON形式
    pub fn parse(text: &str) -> Result<Self, DependencyGraphError> {
        // 1. 矢印形式をパース
        if let Some(response) = Self::parse_arrow_format(text) {
            return Ok(response);
        }

        // 2. JSON形式を試行
        if let Ok(parsed) = serde_json::from_str(text) {
            return Ok(parsed);
        }

        // 3. テキストからJSONを抽出して再試行
        if let Some(json) = Self::extract_json(text) {
            serde_json::from_str(&json).map_err(|e| DependencyGraphError::ParseError(e.to_string()))
        } else {
            Err(DependencyGraphError::ParseError(format!(
                "No valid format found in response: {}",
                text.chars().take(200).collect::<String>()
            )))
        }
    }

    /// アクション順序をパース（矢印形式 or 番号リスト形式）
    ///
    /// 対応形式:
    /// - 矢印: "Grep -> Read -> List"
    /// - 番号: "1. Grep 2. Read 3. List"
    /// - コンマ: "Grep, Read, List"
    fn parse_arrow_format(text: &str) -> Option<Self> {
        // 矢印形式を試行
        if let Some(result) = Self::parse_arrow_only(text) {
            return Some(result);
        }

        // 番号リスト形式を試行（例: "1. Read 2. List 3. Grep"）
        if let Some(result) = Self::parse_numbered_list(text) {
            return Some(result);
        }

        None
    }

    /// 矢印形式のみをパース
    fn parse_arrow_only(text: &str) -> Option<Self> {
        let normalized = text.replace('→', "->");

        // 矢印を含む行を探す
        let arrow_line = normalized.lines().find(|line| line.contains("->"))?;

        let parts: Vec<&str> = arrow_line.split("->").collect();
        if parts.len() < 2 {
            return None;
        }

        // 最後の単語のみを抽出（"So the order is Read" -> "Read"）
        let actions_in_order: Vec<String> = parts
            .iter()
            .filter_map(|part| {
                let trimmed = part.trim();
                // 最後の単語を取得
                let last_word = trimmed.split_whitespace().last()?;
                // 英字のみを抽出
                let action: String = last_word.chars().filter(|c| c.is_alphabetic()).collect();
                if action.is_empty() {
                    None
                } else {
                    Some(action)
                }
            })
            .collect();

        if actions_in_order.len() < 2 {
            return None;
        }

        Self::build_response(actions_in_order)
    }

    /// 番号リスト形式をパース（例: "1. Read 2. List 3. Grep"）
    fn parse_numbered_list(text: &str) -> Option<Self> {
        let mut actions_in_order: Vec<String> = Vec::new();

        // "1." "2." "3." などを探す
        for i in 1..=10 {
            let pattern = format!("{}.", i);
            if let Some(pos) = text.find(&pattern) {
                // 番号の後の単語を取得
                let after = &text[pos + pattern.len()..];
                if let Some(word) = after.split_whitespace().next() {
                    // 英字のみを抽出
                    let action: String = word.chars().filter(|c| c.is_alphabetic()).collect();
                    if !action.is_empty() && !actions_in_order.contains(&action) {
                        actions_in_order.push(action);
                    }
                }
            }
        }

        if actions_in_order.len() < 2 {
            return None;
        }

        Self::build_response(actions_in_order)
    }

    /// LlmDependencyResponse を構築
    fn build_response(actions_in_order: Vec<String>) -> Option<Self> {
        let mut edges = Vec::new();
        for window in actions_in_order.windows(2) {
            edges.push(LlmEdge {
                from: window[0].clone(),
                to: window[1].clone(),
                confidence: 0.9,
            });
        }

        Some(Self {
            edges,
            start: vec![actions_in_order.first()?.clone()],
            terminal: vec![actions_in_order.last()?.clone()],
            reasoning: Some("Parsed from text format".to_string()),
        })
    }

    /// テキストからJSONオブジェクトを抽出
    fn extract_json(text: &str) -> Option<String> {
        // { ... } を探す（バランスを取って）
        let start = text.find('{')?;
        let chars: Vec<char> = text[start..].chars().collect();
        let mut depth = 0;
        let mut in_string = false;
        let mut escape_next = false;

        for (i, &ch) in chars.iter().enumerate() {
            if escape_next {
                escape_next = false;
                continue;
            }

            match ch {
                '\\' if in_string => escape_next = true,
                '"' => in_string = !in_string,
                '{' if !in_string => depth += 1,
                '}' if !in_string => {
                    depth -= 1;
                    if depth == 0 {
                        return Some(chars[..=i].iter().collect());
                    }
                }
                _ => {}
            }
        }

        None
    }
}

// ============================================================================
// Planner Trait
// ============================================================================

/// 依存グラフ生成プランナー trait
///
/// LLM を使った実装と、静的な実装の両方をサポート。
pub trait DependencyPlanner: Send + Sync {
    /// タスクとアクション一覧から依存グラフを生成
    ///
    /// 非同期処理は呼び出し側で管理（LLM 呼び出しは別レイヤー）。
    fn plan(
        &self,
        task: &str,
        available_actions: &[String],
    ) -> Result<DependencyGraph, DependencyGraphError>;

    /// プランナー名
    fn name(&self) -> &str;
}

/// 静的な依存グラフプランナー（LLM 不使用）
///
/// 事前定義されたパターンから依存グラフを生成。
/// テストや LLM なしでの動作確認に使用。
#[derive(Debug, Clone, Default)]
pub struct StaticDependencyPlanner {
    /// パターン名 → グラフ のマッピング
    patterns: HashMap<String, DependencyGraph>,
    /// デフォルトパターン名
    default_pattern: Option<String>,
}

impl StaticDependencyPlanner {
    pub fn new() -> Self {
        Self::default()
    }

    /// パターンを追加
    pub fn with_pattern(mut self, name: impl Into<String>, graph: DependencyGraph) -> Self {
        let name = name.into();
        if self.default_pattern.is_none() {
            self.default_pattern = Some(name.clone());
        }
        self.patterns.insert(name, graph);
        self
    }

    /// デフォルトパターンを設定
    pub fn with_default_pattern(mut self, name: impl Into<String>) -> Self {
        self.default_pattern = Some(name.into());
        self
    }

    /// ファイル探索パターンを追加（組み込み）
    ///
    /// Grep|List → Read のパターン。
    pub fn with_file_exploration_pattern(self) -> Self {
        let graph = DependencyGraph::builder()
            .task("File exploration")
            .available_actions(["Grep", "List", "Read"])
            .edge("Grep", "Read", 0.95)
            .edge("List", "Grep", 0.60)
            .edge("List", "Read", 0.40)
            .start_nodes(["Grep", "List"])
            .terminal_node("Read")
            .build();

        self.with_pattern("file_exploration", graph)
    }

    /// コード検索パターンを追加（組み込み）
    ///
    /// Grep → Read の強いパターン。
    pub fn with_code_search_pattern(self) -> Self {
        let graph = DependencyGraph::builder()
            .task("Code search")
            .available_actions(["Grep", "Read"])
            .edge("Grep", "Read", 0.95)
            .start_node("Grep")
            .terminal_node("Read")
            .build();

        self.with_pattern("code_search", graph)
    }
}

impl DependencyPlanner for StaticDependencyPlanner {
    fn plan(
        &self,
        task: &str,
        available_actions: &[String],
    ) -> Result<DependencyGraph, DependencyGraphError> {
        // デフォルトパターンを使用
        if let Some(pattern_name) = &self.default_pattern {
            if let Some(graph) = self.patterns.get(pattern_name) {
                let mut graph = graph.clone();
                graph.task = task.to_string();
                graph.available_actions = available_actions.to_vec();
                return Ok(graph);
            }
        }

        // パターンがない場合は単純な線形グラフを生成
        // 全アクションを順番に実行する想定
        if available_actions.is_empty() {
            return Err(DependencyGraphError::NoStartNodes);
        }

        let mut builder = DependencyGraphBuilder::new()
            .task(task)
            .available_actions(available_actions.to_vec())
            .start_node(&available_actions[0]);

        if available_actions.len() > 1 {
            for window in available_actions.windows(2) {
                builder = builder.edge(&window[0], &window[1], 0.80);
            }
            builder = builder.terminal_node(&available_actions[available_actions.len() - 1]);
        } else {
            builder = builder.terminal_node(&available_actions[0]);
        }

        Ok(builder.build())
    }

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

// ============================================================================
// LLM Prompt Generation
// ============================================================================

use crate::actions::ActionDef;

/// LLM 向けプロンプト生成
///
/// 分割クエリ方式：first/lastを個別に聞いて安定した結果を得る。
/// ActionDef の name と description を使ってヒントを生成。
pub struct DependencyPromptGenerator;

impl DependencyPromptGenerator {
    /// 依存グラフ生成用のプロンプトを生成
    ///
    /// 従来の全順序プロンプト（フォールバック用）
    pub fn generate_prompt(task: &str, actions: &[ActionDef]) -> String {
        let actions_list = actions
            .iter()
            .map(|a| a.name.as_str())
            .collect::<Vec<_>>()
            .join(", ");

        format!(
            r#"{task}
Steps: {actions_list}
The very first step is:"#
        )
    }

    /// First step を聞くプロンプト
    ///
    /// ActionDef の description から動詞を抽出してヒントを生成。
    /// 例: "Search for patterns" → "Which step SEARCHES?"
    pub fn generate_first_prompt(_task: &str, actions: &[ActionDef]) -> String {
        // アルファベット順にソート（順序バイアスを軽減）
        let mut sorted_actions: Vec<&ActionDef> = actions.iter().collect();
        sorted_actions.sort_by(|a, b| a.name.cmp(&b.name));

        let actions_list = sorted_actions
            .iter()
            .map(|a| a.name.as_str())
            .collect::<Vec<_>>()
            .join(", ");

        // description から動詞を抽出（最初の単語を大文字化）
        let descriptions: Vec<String> = sorted_actions
            .iter()
            .map(|a| format!("- {}: {}", a.name, a.description))
            .collect();
        let descriptions_block = descriptions.join("\n");

        // 最初のアクションの description から動詞を取得してヒントに
        let first_verb = sorted_actions
            .first()
            .map(|a| Self::extract_verb(&a.description))
            .unwrap_or_else(|| "CHECK".to_string());

        format!(
            r#"Steps: {actions_list}
{descriptions_block}
Which step {first_verb}S first?
Answer:"#
        )
    }

    /// Last step を聞くプロンプト
    ///
    /// ActionDef の name と description をそのまま渡す。
    /// description に "requires X first" 等の情報があれば LLM が判断できる。
    pub fn generate_last_prompt(_task: &str, actions: &[ActionDef]) -> String {
        // アルファベット順にソート（順序バイアスを軽減）
        let mut sorted_actions: Vec<&ActionDef> = actions.iter().collect();
        sorted_actions.sort_by(|a, b| a.name.cmp(&b.name));

        let actions_list = sorted_actions
            .iter()
            .map(|a| a.name.as_str())
            .collect::<Vec<_>>()
            .join(", ");

        let descriptions: Vec<String> = sorted_actions
            .iter()
            .map(|a| format!("- {}: {}", a.name, a.description))
            .collect();
        let descriptions_block = descriptions.join("\n");

        format!(
            r#"Steps: {actions_list}
{descriptions_block}
Which step should be done last?
Answer:"#
        )
    }

    /// ペア比較プロンプト（どちらが先か）
    pub fn generate_pair_prompt(task: &str, action_a: &str, action_b: &str) -> String {
        format!(
            r#"For {task}, which comes first: {action_a} or {action_b}?
Answer (one word):"#
        )
    }

    /// description から動詞を抽出（大文字化）
    ///
    /// 例: "Search for patterns" → "SEARCH"
    /// 例: "Reads file contents" → "READ"
    fn extract_verb(description: &str) -> String {
        description
            .split_whitespace()
            .next()
            .map(|w| {
                // 末尾の 's' を除去（"Reads" → "Read"）
                let word = w.trim_end_matches('s').trim_end_matches('S');
                word.to_uppercase()
            })
            .unwrap_or_else(|| "CHECK".to_string())
    }
}

// ============================================================================
// Graph Navigation Helper
// ============================================================================

/// 依存グラフに基づくアクション選択ヘルパー
///
/// 現在の探索状態と依存グラフを組み合わせて、
/// 次に実行すべきアクションを推奨する。
#[derive(Debug, Clone)]
pub struct GraphNavigator {
    graph: DependencyGraph,
    /// 実行済みアクション（成功したもの）
    completed_actions: HashSet<String>,
}

impl GraphNavigator {
    pub fn new(graph: DependencyGraph) -> Self {
        Self {
            graph,
            completed_actions: HashSet::new(),
        }
    }

    /// アクションを完了としてマーク
    pub fn mark_completed(&mut self, action: &str) {
        self.completed_actions.insert(action.to_string());
    }

    /// 次に実行すべきアクションを取得
    ///
    /// - まだ開始していない場合: 開始ノードを返す
    /// - 実行中の場合: 完了したアクションの後続を返す
    /// - 全て完了している場合: 空を返す
    pub fn suggest_next(&self) -> Vec<String> {
        if self.completed_actions.is_empty() {
            // まだ何も実行していない → 開始ノード
            return self.graph.start_actions();
        }

        // 完了したアクションの後続を収集
        let mut candidates = Vec::new();
        for completed in &self.completed_actions {
            for next in self.graph.valid_next_actions(completed) {
                if !self.completed_actions.contains(&next) && !candidates.contains(&next) {
                    candidates.push(next);
                }
            }
        }

        candidates
    }

    /// タスクが完了したか判定
    ///
    /// 終端アクションのいずれかが完了していれば true。
    pub fn is_task_complete(&self) -> bool {
        self.graph
            .terminal_actions()
            .iter()
            .any(|t| self.completed_actions.contains(t))
    }

    /// 進捗を取得（完了アクション数 / 全アクション数）
    pub fn progress(&self) -> f64 {
        if self.graph.available_actions.is_empty() {
            return 0.0;
        }
        self.completed_actions.len() as f64 / self.graph.available_actions.len() as f64
    }

    /// 依存グラフへの参照を取得
    pub fn graph(&self) -> &DependencyGraph {
        &self.graph
    }
}

// ============================================================================
// LearnedDependencyProvider - 学習データからグラフを生成
// ============================================================================

use crate::learn::offline::LearnedActionOrder;

/// 学習済みアクション順序から DependencyGraph を生成する Provider
///
/// `OfflineModel.action_order` を使用して、LLM を呼ばずに即座にグラフを構築する。
/// アクション集合が一致しない場合は `None` を返し、LLM フォールバックさせる。
///
/// # Example
///
/// ```ignore
/// use swarm_engine_core::learn::offline::LearnedActionOrder;
/// use swarm_engine_core::exploration::LearnedDependencyProvider;
///
/// let order = LearnedActionOrder::new(
///     vec!["Grep".to_string(), "Read".to_string()],  // discover
///     vec!["Restart".to_string()],                    // not_discover
///     &["Grep", "Read", "Restart"].map(String::from).to_vec(),
/// );
///
/// let provider = LearnedDependencyProvider::new(order);
/// let graph = provider.provide_graph("task", &actions);
/// ```
#[derive(Debug, Clone)]
pub struct LearnedDependencyProvider {
    action_order: LearnedActionOrder,
}

impl LearnedDependencyProvider {
    /// 新しい LearnedDependencyProvider を作成
    pub fn new(action_order: LearnedActionOrder) -> Self {
        Self { action_order }
    }

    /// 学習済み順序を取得
    pub fn action_order(&self) -> &LearnedActionOrder {
        &self.action_order
    }
}

impl DependencyGraphProvider for LearnedDependencyProvider {
    fn provide_graph(&self, task: &str, available_actions: &[String]) -> Option<DependencyGraph> {
        // アクション集合が一致するか確認
        if !self.action_order.matches_actions(available_actions) {
            tracing::debug!(
                expected_hash = self.action_order.action_set_hash,
                actual_hash = LearnedActionOrder::compute_hash(available_actions),
                "Action set mismatch, falling back to LLM"
            );
            return None;
        }

        // グラフを構築
        let sorted_discover = &self.action_order.discover;
        let sorted_not_discover = &self.action_order.not_discover;

        // 両方空の場合は有効なグラフを構築できない
        if sorted_discover.is_empty() && sorted_not_discover.is_empty() {
            tracing::warn!("Both discover and not_discover are empty, cannot build graph");
            return None;
        }

        tracing::info!(
            discover = ?self.action_order.discover,
            not_discover = ?self.action_order.not_discover,
            "Using learned action order (LLM skipped)"
        );

        let mut builder = DependencyGraphBuilder::new()
            .task(task)
            .available_actions(available_actions.iter().cloned());

        // 開始ノード設定
        if !sorted_discover.is_empty() {
            builder = builder.start_node(&sorted_discover[0]);
        } else if !sorted_not_discover.is_empty() {
            builder = builder.start_node(&sorted_not_discover[0]);
        }

        // 終端ノード設定
        if let Some(last) = sorted_not_discover.last() {
            builder = builder.terminal_node(last);
        } else if !sorted_discover.is_empty() {
            builder = builder.terminal_node(sorted_discover.last().unwrap());
        }

        // Discover 間のエッジ
        for window in sorted_discover.windows(2) {
            builder = builder.edge(&window[0], &window[1], 0.9);
        }

        // Discover → NotDiscover へのエッジ
        if !sorted_discover.is_empty() && !sorted_not_discover.is_empty() {
            builder = builder.edge(
                sorted_discover.last().unwrap(),
                &sorted_not_discover[0],
                0.9,
            );
        }

        // NotDiscover 間のエッジ
        for window in sorted_not_discover.windows(2) {
            builder = builder.edge(&window[0], &window[1], 0.9);
        }

        // バリデーション付きでビルド（失敗時は None）
        builder.build_validated().ok()
    }
}

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

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

    #[test]
    fn test_dependency_graph_builder() {
        let graph = DependencyGraph::builder()
            .task("Find auth function")
            .available_actions(["Grep", "List", "Read"])
            .edge("Grep", "Read", 0.95)
            .edge("List", "Grep", 0.60)
            .start_nodes(["Grep", "List"])
            .terminal_node("Read")
            .build();

        assert_eq!(graph.task(), "Find auth function");
        assert!(graph.is_start("Grep"));
        assert!(graph.is_start("List"));
        assert!(graph.is_terminal("Read"));
        assert!(graph.can_transition("Grep", "Read"));
        assert!(!graph.can_transition("Read", "Grep"));
    }

    #[test]
    fn test_valid_next_actions() {
        let graph = DependencyGraph::builder()
            .available_actions(["Grep", "List", "Read"])
            .edge("Grep", "Read", 0.95)
            .edge("List", "Grep", 0.60)
            .edge("List", "Read", 0.40)
            .start_nodes(["Grep", "List"])
            .terminal_node("Read")
            .build();

        // Grep の後は Read
        let next = graph.valid_next_actions("Grep");
        assert_eq!(next, vec!["Read"]);

        // List の後は Grep (0.60) と Read (0.40)、確信度順
        let next = graph.valid_next_actions("List");
        assert_eq!(next, vec!["Grep", "Read"]);

        // Read の後は何もない（終端）
        let next = graph.valid_next_actions("Read");
        assert!(next.is_empty());
    }

    #[test]
    fn test_static_planner_file_exploration() {
        let planner = StaticDependencyPlanner::new().with_file_exploration_pattern();

        let graph = planner
            .plan("Find auth.rs", &["Grep".to_string(), "Read".to_string()])
            .unwrap();

        assert!(graph.is_start("Grep"));
        assert!(graph.is_terminal("Read"));
    }

    #[test]
    fn test_graph_navigator() {
        let graph = DependencyGraph::builder()
            .available_actions(["Grep", "Read"])
            .edge("Grep", "Read", 0.95)
            .start_node("Grep")
            .terminal_node("Read")
            .build();

        let mut nav = GraphNavigator::new(graph);

        // 最初は Grep を推奨
        assert_eq!(nav.suggest_next(), vec!["Grep"]);
        assert!(!nav.is_task_complete());

        // Grep 完了
        nav.mark_completed("Grep");
        assert_eq!(nav.suggest_next(), vec!["Read"]);
        assert!(!nav.is_task_complete());

        // Read 完了 → タスク完了
        nav.mark_completed("Read");
        assert!(nav.is_task_complete());
        assert!(nav.suggest_next().is_empty());
    }

    #[test]
    fn test_llm_response_parsing() {
        let json = r#"{
            "edges": [
                {"from": "Grep", "to": "Read", "confidence": 0.95}
            ],
            "start": ["Grep"],
            "terminal": ["Read"],
            "reasoning": "Search first, then read"
        }"#;

        let response = LlmDependencyResponse::parse(json).unwrap();
        assert_eq!(response.edges.len(), 1);
        assert_eq!(response.start, vec!["Grep"]);
        assert_eq!(response.terminal, vec!["Read"]);
        assert!(response.reasoning.is_some());

        let graph = response.into_graph(
            "Find function",
            vec!["Grep".to_string(), "Read".to_string()],
        );
        assert!(graph.can_transition("Grep", "Read"));
    }

    #[test]
    fn test_mermaid_output() {
        let graph = DependencyGraph::builder()
            .available_actions(["Grep", "List", "Read"])
            .edge("Grep", "Read", 0.95)
            .edge("List", "Grep", 0.60)
            .start_nodes(["Grep", "List"])
            .terminal_node("Read")
            .build();

        let mermaid = graph.to_mermaid();
        assert!(mermaid.contains("graph LR"));
        assert!(mermaid.contains("Grep -->|95%| Read"));
        assert!(mermaid.contains("style Read fill:#f99"));
    }

    // =========================================================================
    // LearnedDependencyProvider Tests
    // =========================================================================

    #[test]
    fn test_learned_action_order_hash() {
        let actions = vec![
            "Grep".to_string(),
            "Read".to_string(),
            "Restart".to_string(),
        ];

        let order = LearnedActionOrder::new(
            vec!["Grep".to_string(), "Read".to_string()],
            vec!["Restart".to_string()],
            &actions,
        );

        // 同じアクション集合（順序違い）なら同じハッシュ
        let actions_reordered = vec![
            "Restart".to_string(),
            "Grep".to_string(),
            "Read".to_string(),
        ];
        assert!(order.matches_actions(&actions_reordered));

        // 異なるアクション集合なら不一致
        let actions_different = vec!["Grep".to_string(), "Read".to_string()];
        assert!(!order.matches_actions(&actions_different));
    }

    #[test]
    fn test_learned_dependency_provider_cache_hit() {
        let actions = vec![
            "Grep".to_string(),
            "Read".to_string(),
            "Restart".to_string(),
        ];

        let order = LearnedActionOrder::new(
            vec!["Grep".to_string(), "Read".to_string()],
            vec!["Restart".to_string()],
            &actions,
        );

        let provider = LearnedDependencyProvider::new(order);

        // キャッシュヒット
        let graph = provider.provide_graph("troubleshooting", &actions);
        assert!(graph.is_some());

        let graph = graph.unwrap();
        assert!(graph.is_start("Grep"));
        assert!(graph.is_terminal("Restart"));
        assert!(graph.can_transition("Grep", "Read"));
        assert!(graph.can_transition("Read", "Restart"));
    }

    #[test]
    fn test_learned_dependency_provider_cache_miss() {
        let original_actions = vec![
            "Grep".to_string(),
            "Read".to_string(),
            "Restart".to_string(),
        ];

        let order = LearnedActionOrder::new(
            vec!["Grep".to_string(), "Read".to_string()],
            vec!["Restart".to_string()],
            &original_actions,
        );

        let provider = LearnedDependencyProvider::new(order);

        // 異なるアクション集合 → キャッシュミス（None）
        let different_actions = vec!["Grep".to_string(), "Read".to_string()];
        let graph = provider.provide_graph("troubleshooting", &different_actions);
        assert!(graph.is_none());
    }

    #[test]
    fn test_learned_dependency_provider_discover_only() {
        let actions = vec!["Grep".to_string(), "Read".to_string()];

        let order = LearnedActionOrder::new(
            vec!["Grep".to_string(), "Read".to_string()],
            vec![], // NotDiscover なし
            &actions,
        );

        let provider = LearnedDependencyProvider::new(order);
        let graph = provider.provide_graph("search task", &actions);
        assert!(graph.is_some());

        let graph = graph.unwrap();
        assert!(graph.is_start("Grep"));
        assert!(graph.is_terminal("Read")); // Discover の最後が Terminal
        assert!(graph.can_transition("Grep", "Read"));
    }

    #[test]
    fn test_learned_dependency_provider_not_discover_only() {
        let actions = vec!["Restart".to_string(), "CheckStatus".to_string()];

        let order = LearnedActionOrder::new(
            vec![], // Discover なし
            vec!["Restart".to_string(), "CheckStatus".to_string()],
            &actions,
        );

        let provider = LearnedDependencyProvider::new(order);
        let graph = provider.provide_graph("ops task", &actions);
        assert!(graph.is_some());

        let graph = graph.unwrap();
        assert!(graph.is_start("Restart")); // NotDiscover の最初が Start
        assert!(graph.is_terminal("CheckStatus"));
        assert!(graph.can_transition("Restart", "CheckStatus"));
    }

    #[test]
    fn test_learned_dependency_provider_empty_lists() {
        let actions = vec!["Grep".to_string(), "Read".to_string()];

        let order = LearnedActionOrder::new(
            vec![], // Discover なし
            vec![], // NotDiscover なし
            &actions,
        );

        let provider = LearnedDependencyProvider::new(order);
        // 両方空の場合は None を返す
        let graph = provider.provide_graph("empty task", &actions);
        assert!(graph.is_none());
    }

    // =========================================================================
    // extract_json Tests
    // =========================================================================

    #[test]
    fn test_extract_json_simple() {
        let text = r#"Here is the result: {"edges": [], "start": ["A"], "terminal": ["B"]}"#;
        let json = LlmDependencyResponse::extract_json(text);
        assert!(json.is_some());
        let json = json.unwrap();
        assert!(json.starts_with('{'));
        assert!(json.ends_with('}'));
    }

    #[test]
    fn test_extract_json_nested() {
        let text = r#"Result: {"edges": [{"from": "A", "to": "B", "confidence": 0.9}], "start": ["A"], "terminal": ["B"]}"#;
        let json = LlmDependencyResponse::extract_json(text);
        assert!(json.is_some());

        // パースできることを確認
        let parsed: Result<LlmDependencyResponse, _> = serde_json::from_str(&json.unwrap());
        assert!(parsed.is_ok());
    }

    #[test]
    fn test_extract_json_with_string_braces() {
        // 文字列内の {} は無視される
        let text =
            r#"{"edges": [], "start": ["A"], "terminal": ["B"], "reasoning": "Use {pattern}"}"#;
        let json = LlmDependencyResponse::extract_json(text);
        assert!(json.is_some());
        assert_eq!(json.unwrap(), text);
    }

    #[test]
    fn test_extract_json_no_json() {
        let text = "This is just plain text without JSON";
        let json = LlmDependencyResponse::extract_json(text);
        assert!(json.is_none());
    }

    // =========================================================================
    // validate() Tests
    // =========================================================================

    #[test]
    fn test_validate_unknown_start_node() {
        let graph = DependencyGraph::builder()
            .available_actions(["Grep", "Read"])
            .start_node("Unknown") // available_actions に含まれない
            .terminal_node("Read")
            .build();

        let result = graph.validate();
        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            DependencyGraphError::UnknownAction(name) if name == "Unknown"
        ));
    }

    #[test]
    fn test_validate_unknown_terminal_node() {
        let graph = DependencyGraph::builder()
            .available_actions(["Grep", "Read"])
            .start_node("Grep")
            .terminal_node("Unknown") // available_actions に含まれない
            .build();

        let result = graph.validate();
        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            DependencyGraphError::UnknownAction(name) if name == "Unknown"
        ));
    }

    #[test]
    fn test_validate_valid_graph() {
        let graph = DependencyGraph::builder()
            .available_actions(["Grep", "Read"])
            .edge("Grep", "Read", 0.9)
            .start_node("Grep")
            .terminal_node("Read")
            .build();

        assert!(graph.validate().is_ok());
    }

    // =========================================================================
    // start_actions / terminal_actions Order Tests
    // =========================================================================

    #[test]
    fn test_start_actions_sorted() {
        let graph = DependencyGraph::builder()
            .available_actions(["Zebra", "Apple", "Mango"])
            .start_nodes(["Zebra", "Apple", "Mango"])
            .terminal_node("Zebra")
            .build();

        let actions = graph.start_actions();
        // アルファベット順にソートされる
        assert_eq!(actions, vec!["Apple", "Mango", "Zebra"]);
    }

    #[test]
    fn test_terminal_actions_sorted() {
        let graph = DependencyGraph::builder()
            .available_actions(["Zebra", "Apple", "Mango"])
            .start_node("Apple")
            .terminal_nodes(["Zebra", "Apple", "Mango"])
            .build();

        let actions = graph.terminal_actions();
        // アルファベット順にソートされる
        assert_eq!(actions, vec!["Apple", "Mango", "Zebra"]);
    }
}