swarm-engine-core 0.1.6

//! Context System - スコープ付きコンテキスト管理
//!
//! # 設計思想
//!
//! Manager/Worker が「見えるべき」情報を制御するためのスコープシステム。
//! データは ContextStore に正規化して格納し、ContextView で可視範囲を定義。
//! ContextResolver が View に応じて ResolvedContext を生成。
//!
//! ```text
//! ContextStore (正規化データ)
//!       │
//!       ├── ContextView::Global    → Manager用（全体が見える）
//!       ├── ContextView::Local     → Worker用（自分+Neighbor）
//!       └── ContextView::Custom    → カスタムフィルタ
//!       │
//!       ▼
//! ContextResolver.resolve(store, view) → ResolvedContext
//!       │
//!       ▼
//! LLM層でプロンプト生成
//! ```

use std::collections::HashMap;

use serde_json::Value;

use crate::actions::ActionsConfig;
use crate::state::Escalation;
use crate::types::WorkerId;

use crate::agent::{ManagerId, ManagerInstruction};

// ============================================================================
// ContextStore - 正規化データストア
// ============================================================================

/// 正規化されたコンテキストデータストア
///
/// 全ての情報を一箇所に集約し、View によって可視範囲を制御する。
#[derive(Debug, Clone)]
pub struct ContextStore {
    /// Global情報（tick, progress等）
    pub global: GlobalContext,
    /// 全Worker状態
    pub workers: HashMap<WorkerId, WorkerContext>,
    /// 全Manager情報
    pub managers: HashMap<ManagerId, ManagerContext>,
    /// Escalation一覧
    pub escalations: Vec<(WorkerId, Escalation)>,
    /// 利用可能Actions
    pub actions: Option<ActionsConfig>,
    /// 拡張メタデータ
    pub metadata: HashMap<String, Value>,
}

impl ContextStore {
    /// 新しい ContextStore を作成
    pub fn new(tick: u64) -> Self {
        Self {
            global: GlobalContext::new(tick),
            workers: HashMap::new(),
            managers: HashMap::new(),
            escalations: Vec::new(),
            actions: None,
            metadata: HashMap::new(),
        }
    }

    /// Worker を追加
    pub fn with_worker(mut self, ctx: WorkerContext) -> Self {
        self.workers.insert(ctx.id, ctx);
        self
    }

    /// Manager を追加
    pub fn with_manager(mut self, ctx: ManagerContext) -> Self {
        self.managers.insert(ctx.id, ctx);
        self
    }

    /// Escalation を追加
    pub fn with_escalation(mut self, worker_id: WorkerId, escalation: Escalation) -> Self {
        self.escalations.push((worker_id, escalation));
        self
    }

    /// Actions を設定
    pub fn with_actions(mut self, actions: ActionsConfig) -> Self {
        self.actions = Some(actions);
        self
    }

    /// メタデータを追加
    pub fn insert<V: Into<Value>>(mut self, key: impl Into<String>, value: V) -> Self {
        self.metadata.insert(key.into(), value.into());
        self
    }

    /// メタデータを取得
    pub fn get(&self, key: &str) -> Option<&Value> {
        self.metadata.get(key)
    }

    /// メタデータを文字列として取得
    pub fn get_str(&self, key: &str) -> Option<&str> {
        self.metadata.get(key).and_then(|v| v.as_str())
    }
}

// ============================================================================
// GlobalContext - 全体情報
// ============================================================================

/// Global情報（全リクエスト共通）
#[derive(Debug, Clone, Default)]
pub struct GlobalContext {
    /// 現在の tick
    pub tick: u64,
    /// 最大 tick
    pub max_ticks: u64,
    /// 進捗 (0.0 - 1.0)
    pub progress: f64,
    /// 成功率 (0.0 - 1.0)
    pub success_rate: f64,
    /// タスク説明
    pub task_description: Option<String>,
    /// ヒント
    pub hint: Option<String>,
}

impl GlobalContext {
    pub fn new(tick: u64) -> Self {
        Self {
            tick,
            ..Default::default()
        }
    }

    pub fn with_max_ticks(mut self, max: u64) -> Self {
        self.max_ticks = max;
        self
    }

    pub fn with_progress(mut self, progress: f64) -> Self {
        self.progress = progress;
        self
    }

    pub fn with_success_rate(mut self, rate: f64) -> Self {
        self.success_rate = rate;
        self
    }

    pub fn with_task(mut self, description: impl Into<String>) -> Self {
        self.task_description = Some(description.into());
        self
    }

    pub fn with_hint(mut self, hint: impl Into<String>) -> Self {
        self.hint = Some(hint.into());
        self
    }
}

// ============================================================================
// WorkerContext - Worker固有情報
// ============================================================================

/// Worker固有のコンテキスト
#[derive(Debug, Clone)]
pub struct WorkerContext {
    /// Worker ID
    pub id: WorkerId,
    /// 連続失敗数
    pub consecutive_failures: u32,
    /// 最新アクション名
    pub last_action: Option<String>,
    /// 最新アクションの成功/失敗
    pub last_success: Option<bool>,
    /// 履歴の長さ
    pub history_len: usize,
    /// Escalation 中かどうか
    pub has_escalation: bool,
    /// このWorkerに適用可能なAction候補
    pub candidates: Vec<String>,
    /// Worker固有メタデータ
    pub metadata: HashMap<String, Value>,
}

impl WorkerContext {
    pub fn new(id: WorkerId) -> Self {
        Self {
            id,
            consecutive_failures: 0,
            last_action: None,
            last_success: None,
            history_len: 0,
            has_escalation: false,
            candidates: Vec::new(),
            metadata: HashMap::new(),
        }
    }

    pub fn with_failures(mut self, count: u32) -> Self {
        self.consecutive_failures = count;
        self
    }

    pub fn with_last_action(mut self, action: impl Into<String>, success: bool) -> Self {
        self.last_action = Some(action.into());
        self.last_success = Some(success);
        self
    }

    pub fn with_history_len(mut self, len: usize) -> Self {
        self.history_len = len;
        self
    }

    pub fn with_escalation(mut self, has: bool) -> Self {
        self.has_escalation = has;
        self
    }

    pub fn with_candidates(mut self, candidates: Vec<String>) -> Self {
        self.candidates = candidates;
        self
    }
}

// ============================================================================
// ManagerContext - Manager固有情報
// ============================================================================

/// Manager固有のコンテキスト
#[derive(Debug, Clone)]
pub struct ManagerContext {
    /// Manager ID
    pub id: ManagerId,
    /// Manager名
    pub name: String,
    /// 最後に処理した tick
    pub last_tick: u64,
    /// Manager固有メタデータ
    pub metadata: HashMap<String, Value>,
}

impl ManagerContext {
    pub fn new(id: ManagerId) -> Self {
        Self {
            id,
            name: format!("Manager_{}", id.0),
            last_tick: 0,
            metadata: HashMap::new(),
        }
    }

    pub fn with_name(mut self, name: impl Into<String>) -> Self {
        self.name = name.into();
        self
    }

    pub fn with_last_tick(mut self, tick: u64) -> Self {
        self.last_tick = tick;
        self
    }
}

// ============================================================================
// ContextView - スコープ定義
// ============================================================================

/// コンテキストの可視範囲を定義
#[derive(Debug, Clone)]
pub enum ContextView {
    /// Manager用: 全体が見える
    Global { manager_id: ManagerId },
    /// Worker用: 自分 + Neighbor
    Local {
        worker_id: WorkerId,
        /// 可視範囲の Neighbor Worker IDs
        neighbor_ids: Vec<WorkerId>,
    },
    /// カスタムフィルタ（将来拡張）
    Custom {
        /// フィルタ名（デバッグ用）
        name: String,
        /// 可視 Worker IDs
        visible_worker_ids: Vec<WorkerId>,
        /// 可視 Manager IDs
        visible_manager_ids: Vec<ManagerId>,
    },
}

impl ContextView {
    /// Manager用の Global View を作成
    pub fn global(manager_id: ManagerId) -> Self {
        Self::Global { manager_id }
    }

    /// Worker用の Local View を作成（Neighborなし）
    pub fn local(worker_id: WorkerId) -> Self {
        Self::Local {
            worker_id,
            neighbor_ids: Vec::new(),
        }
    }

    /// Worker用の Local View を作成（Neighbor指定）
    pub fn local_with_neighbors(worker_id: WorkerId, neighbor_ids: Vec<WorkerId>) -> Self {
        Self::Local {
            worker_id,
            neighbor_ids,
        }
    }

    /// カスタム View を作成
    pub fn custom(
        name: impl Into<String>,
        visible_workers: Vec<WorkerId>,
        visible_managers: Vec<ManagerId>,
    ) -> Self {
        Self::Custom {
            name: name.into(),
            visible_worker_ids: visible_workers,
            visible_manager_ids: visible_managers,
        }
    }
}

// ============================================================================
// ActionCandidate - プロンプト構築用 Action 情報
// ============================================================================

/// プロンプト構築用の Action パラメータ情報
#[derive(Debug, Clone)]
pub struct ActionParam {
    /// パラメータ名
    pub name: String,
    /// 説明
    pub description: String,
    /// 必須かどうか
    pub required: bool,
}

/// プロンプト構築用の Action 情報
///
/// ActionsConfig から必要な情報だけを抽出したモデル。
/// LLM 層でプロンプト生成時に使用する。
#[derive(Debug, Clone)]
pub struct ActionCandidate {
    /// Action 名
    pub name: String,
    /// 説明
    pub description: String,
    /// パラメータ
    pub params: Vec<ActionParam>,
    /// 出力例（JSON 形式）
    pub example: Option<String>,
}

impl ActionCandidate {
    /// ActionsConfig から ActionCandidate のリストを生成
    pub fn from_config(config: &ActionsConfig) -> Vec<Self> {
        config
            .all_actions()
            .map(|def| ActionCandidate {
                name: def.name.clone(),
                description: def.description.clone(),
                params: def
                    .params
                    .iter()
                    .map(|p| ActionParam {
                        name: p.name.clone(),
                        description: p.description.clone(),
                        required: p.required,
                    })
                    .collect(),
                example: def.example.clone(),
            })
            .collect()
    }
}

// ============================================================================
// ResolvedContext - 解決済みコンテキスト
// ============================================================================

/// 解決済みコンテキスト（LLM層に渡る）
///
/// ContextResolver が ContextStore + ContextView から生成。
/// プロンプト生成に必要な情報のみを含む。
#[derive(Debug, Clone)]
pub struct ResolvedContext {
    /// Global情報
    pub global: GlobalContext,
    /// 可視範囲の Worker 情報
    pub visible_workers: Vec<WorkerContext>,
    /// 可視範囲の Escalation
    pub escalations: Vec<(WorkerId, Escalation)>,
    /// 利用可能な Action 候補（プロンプト構築用）
    pub candidates: Vec<ActionCandidate>,
    /// 追加メタデータ
    pub metadata: HashMap<String, Value>,
    /// このコンテキストの対象
    pub target: ContextTarget,
    /// 自分の last_output（Scope::Minimal 用）
    ///
    /// WorkerScope::Minimal の場合、visible_workers は空にし、
    /// 代わりにこのフィールドに自分の前回結果のみを格納する。
    pub self_last_output: Option<String>,
    /// Manager からの指示（前回 Guidance から抽出）
    ///
    /// 前回の Manager 判断を次回の Worker Prompt に埋め込むために使用。
    /// Manager.prepare() で前回の Guidance から生成して設定する。
    pub manager_instruction: Option<ManagerInstruction>,
}

/// コンテキストの対象
#[derive(Debug, Clone)]
pub enum ContextTarget {
    Manager(ManagerId),
    Worker(WorkerId),
}

impl ResolvedContext {
    /// 新しい ResolvedContext を作成
    pub fn new(global: GlobalContext, target: ContextTarget) -> Self {
        Self {
            global,
            visible_workers: Vec::new(),
            escalations: Vec::new(),
            candidates: Vec::new(),
            metadata: HashMap::new(),
            target,
            self_last_output: None,
            manager_instruction: None,
        }
    }

    /// Worker情報を追加
    pub fn with_workers(mut self, workers: Vec<WorkerContext>) -> Self {
        self.visible_workers = workers;
        self
    }

    /// Escalationを追加
    pub fn with_escalations(mut self, escalations: Vec<(WorkerId, Escalation)>) -> Self {
        self.escalations = escalations;
        self
    }

    /// 候補を設定
    pub fn with_candidates(mut self, candidates: Vec<ActionCandidate>) -> Self {
        self.candidates = candidates;
        self
    }

    /// ActionsConfig から候補を設定
    pub fn with_actions_config(mut self, config: &ActionsConfig) -> Self {
        self.candidates = ActionCandidate::from_config(config);
        self
    }

    /// メタデータを追加
    pub fn with_metadata(mut self, metadata: HashMap<String, Value>) -> Self {
        self.metadata = metadata;
        self
    }

    /// 自分の last_output を設定（Scope::Minimal 用）
    pub fn with_self_last_output(mut self, output: Option<String>) -> Self {
        self.self_last_output = output;
        self
    }

    /// Manager からの指示を設定
    pub fn with_manager_instruction(mut self, instruction: ManagerInstruction) -> Self {
        self.manager_instruction = Some(instruction);
        self
    }

    /// Escalation があるか
    pub fn has_escalations(&self) -> bool {
        !self.escalations.is_empty()
    }

    /// 対象がManagerか
    pub fn is_manager(&self) -> bool {
        matches!(self.target, ContextTarget::Manager(_))
    }

    /// 対象がWorkerか
    pub fn is_worker(&self) -> bool {
        matches!(self.target, ContextTarget::Worker(_))
    }
}

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

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

    #[test]
    fn test_context_store_builder() {
        let store = ContextStore::new(10)
            .with_worker(WorkerContext::new(WorkerId(0)))
            .with_worker(WorkerContext::new(WorkerId(1)))
            .with_manager(ManagerContext::new(ManagerId(0)))
            .insert("task", "Find the bug");

        assert_eq!(store.global.tick, 10);
        assert_eq!(store.workers.len(), 2);
        assert_eq!(store.managers.len(), 1);
        assert_eq!(store.get_str("task"), Some("Find the bug"));
    }

    #[test]
    fn test_context_view_creation() {
        let global = ContextView::global(ManagerId(0));
        assert!(matches!(global, ContextView::Global { .. }));

        let local = ContextView::local(WorkerId(0));
        assert!(matches!(local, ContextView::Local { .. }));

        let local_with_neighbors =
            ContextView::local_with_neighbors(WorkerId(0), vec![WorkerId(1), WorkerId(2)]);
        if let ContextView::Local { neighbor_ids, .. } = local_with_neighbors {
            assert_eq!(neighbor_ids.len(), 2);
        }
    }

    #[test]
    fn test_worker_context_builder() {
        let ctx = WorkerContext::new(WorkerId(0))
            .with_failures(2)
            .with_last_action("read:/path", true)
            .with_history_len(10)
            .with_escalation(true)
            .with_candidates(vec!["read".into(), "grep".into()]);

        assert_eq!(ctx.id, WorkerId(0));
        assert_eq!(ctx.consecutive_failures, 2);
        assert_eq!(ctx.last_action, Some("read:/path".to_string()));
        assert!(ctx.has_escalation);
        assert_eq!(ctx.candidates.len(), 2);
    }

    #[test]
    fn test_resolved_context() {
        let global = GlobalContext::new(5)
            .with_progress(0.5)
            .with_task("Test task");

        let candidates = vec![ActionCandidate {
            name: "action1".to_string(),
            description: "Test action".to_string(),
            params: vec![],
            example: None,
        }];

        let resolved = ResolvedContext::new(global, ContextTarget::Worker(WorkerId(0)))
            .with_workers(vec![WorkerContext::new(WorkerId(0))])
            .with_candidates(candidates);

        assert!(resolved.is_worker());
        assert!(!resolved.is_manager());
        assert_eq!(resolved.visible_workers.len(), 1);
        assert_eq!(resolved.candidates.len(), 1);
    }
}