swarm-engine-core 0.1.6

//! SwarmEngine State - 2層メモリモデル
//!
//! 共有領域（SharedState）と個別領域（WorkerStates）を分離し、
//! 並列Workerがそれぞれの個別Stateに高速アクセスできる構造。

use std::any::Any;
use std::collections::{HashMap, HashSet, VecDeque};
use std::time::Duration;

use rayon::prelude::*;

use crate::async_task::TaskStatus;
use crate::extensions::Extensions;
use crate::online_stats::SwarmStats;
use crate::types::{AgentId, TaskId, WorkerId};

// ============================================================================
// LlmStats - LLM 呼び出し統計
// ============================================================================

/// LLM 呼び出し統計
#[derive(Debug, Clone, Default)]
pub struct LlmStats {
    /// 総呼び出し回数
    pub invocations: u64,
    /// エラー回数
    pub errors: u64,
    /// 総実行時間
    pub total_duration: Duration,
}

impl LlmStats {
    /// 成功率
    pub fn success_rate(&self) -> f64 {
        if self.invocations == 0 {
            1.0
        } else {
            (self.invocations - self.errors) as f64 / self.invocations as f64
        }
    }

    /// LLM 呼び出しを記録
    pub fn record(&mut self, success: bool, duration: Duration) {
        self.invocations += 1;
        self.total_duration += duration;
        if !success {
            self.errors += 1;
        }
    }
}

/// State 全体
pub struct SwarmState {
    /// 共有領域 - 全Worker が ReadOnly 参照
    pub shared: SharedState,
    /// 個別領域 - Worker ごとの State（並列アクセス用）
    pub workers: WorkerStates,
}

impl SwarmState {
    pub fn new(worker_count: usize) -> Self {
        Self {
            shared: SharedState::default(),
            workers: WorkerStates::new(worker_count),
        }
    }

    /// Tick を進める
    pub fn advance_tick(&mut self) {
        self.shared.tick += 1;
    }
}

/// 共有領域
#[derive(Default)]
pub struct SharedState {
    /// 環境情報
    pub environment: Environment,
    /// 統合統計（ActionEvent ベース）
    pub stats: SwarmStats,
    /// Tick番号
    pub tick: u64,
    /// Agent間で共有するデータ（ReadOnly）
    pub shared_data: SharedData,
    /// 動的リソース（DB接続、HTTPクライアント等）
    pub extensions: Extensions,
    /// 平均Tick時間（ナノ秒）- EMA で計算
    pub avg_tick_duration_ns: u64,
    /// Goal Action（Terminal Action）を達成した Worker
    pub done_workers: HashSet<WorkerId>,
    /// Environment が Done(success=true) を返したかどうか
    pub environment_done: bool,
    /// LLM 呼び出し統計
    pub llm_stats: LlmStats,
}

impl SharedState {
    /// Worker を完了済みとしてマーク
    pub fn mark_worker_done(&mut self, worker_id: WorkerId) {
        self.done_workers.insert(worker_id);
        self.environment_done = true;
    }

    /// Worker が完了済みかどうか
    pub fn is_worker_done(&self, worker_id: WorkerId) -> bool {
        self.done_workers.contains(&worker_id)
    }

    /// 環境が完了したかどうか
    pub fn is_environment_done(&self) -> bool {
        self.environment_done
    }

    /// LLM 呼び出し回数
    pub fn llm_invocations(&self) -> u64 {
        self.llm_stats.invocations
    }

    /// LLM エラー回数
    pub fn llm_errors(&self) -> u64 {
        self.llm_stats.errors
    }
}

/// 環境情報
#[derive(Default)]
pub struct Environment {
    /// 環境変数
    pub variables: HashMap<String, String>,
    /// 設定フラグ
    pub flags: HashMap<String, bool>,
}

// ============================================================================
// Tick Snapshot - 完全履歴記録
// ============================================================================

/// Tick 毎のスナップショット（完全履歴）
///
/// 各Tickで何が起きたかを全て記録。Eval等での詳細分析に使用。
#[derive(Debug, Clone)]
pub struct TickSnapshot {
    /// Tick 番号
    pub tick: u64,
    /// 処理時間
    pub duration: std::time::Duration,
    /// Manager フェーズの記録（起動した場合のみ）
    pub manager_phase: Option<ManagerPhaseSnapshot>,
    /// 全 Worker の結果
    pub worker_results: Vec<WorkerResultSnapshot>,
}

/// Manager フェーズのスナップショット
#[derive(Debug, Clone)]
pub struct ManagerPhaseSnapshot {
    /// Batch リクエスト（Manager が何を聞いたか）
    pub batch_request: crate::agent::BatchDecisionRequest,
    /// LLM からの返答
    pub responses: Vec<(crate::types::WorkerId, crate::agent::DecisionResponse)>,
    /// 発行された Guidance（Worker毎）
    pub guidances: std::collections::HashMap<crate::types::WorkerId, crate::agent::Guidance>,
    /// LLM呼び出しエラー数（サーバー停止等）
    pub llm_errors: u64,
}

/// Worker の結果スナップショット
#[derive(Debug, Clone)]
pub struct WorkerResultSnapshot {
    pub worker_id: crate::types::WorkerId,
    /// 受け取った Guidance（あれば）
    pub guidance_received: Option<crate::agent::Guidance>,
    /// 実行結果
    pub result: WorkResultSnapshot,
}

/// WorkResult のスナップショット版（Clone可能）
#[derive(Debug, Clone)]
pub enum WorkResultSnapshot {
    /// 行動した
    Acted {
        action_result: ActionResultSnapshot,
        state_delta: Option<crate::agent::WorkerStateDelta>,
    },
    /// 継続中
    Continuing { progress: f32 },
    /// Guidance 要求
    NeedsGuidance {
        reason: String,
        context: crate::agent::GuidanceContext,
    },
    /// Escalation 要求
    Escalate {
        reason: crate::agent::EscalationReason,
        context: Option<String>,
    },
    /// 待機
    Idle,
    /// タスク完了
    Done {
        success: bool,
        message: Option<String>,
    },
}

/// ActionResult のスナップショット版（Clone可能）
#[derive(Debug, Clone)]
pub struct ActionResultSnapshot {
    pub success: bool,
    /// output のテキスト表現
    pub output_debug: Option<String>,
    pub duration: std::time::Duration,
    pub error: Option<String>,
}

impl ActionResultSnapshot {
    /// ActionResult から変換
    pub fn from_action_result(result: &crate::types::ActionResult) -> Self {
        Self {
            success: result.success,
            output_debug: result.output.as_ref().map(|o| o.as_text()),
            duration: result.duration,
            error: result.error.clone(),
        }
    }
}

/// SharedData の env エントリ保持数（デフォルト）
const DEFAULT_MAX_ENV_ENTRIES: usize = 500;

/// Agent間共有データ
pub struct SharedData {
    /// key-value ストア
    pub kv: HashMap<String, Vec<u8>>,
    /// 完了した非同期タスク一覧（Meta のみ）
    /// payload は AsyncTaskSystem に保持され、ReadPayload Action で取得
    pub completed_async_tasks: Vec<CompletedAsyncTask>,
    /// env エントリの最大保持数
    max_env_entries: usize,
}

impl Default for SharedData {
    fn default() -> Self {
        Self {
            kv: HashMap::new(),
            completed_async_tasks: Vec::new(),
            max_env_entries: DEFAULT_MAX_ENV_ENTRIES,
        }
    }
}

impl SharedData {
    /// env:{worker_id}:{tick} 形式の古いエントリをクリーンアップ
    ///
    /// kv 内の env:* エントリが max_env_entries を超えた場合、
    /// tick が古いものから削除する。
    pub fn cleanup_env_entries(&mut self) {
        // env:* エントリを収集
        let mut env_entries: Vec<(String, u64)> = self
            .kv
            .keys()
            .filter(|k| k.starts_with("env:"))
            .filter_map(|k| {
                // env:{worker_id}:{tick} から tick を抽出
                k.rsplit(':')
                    .next()?
                    .parse::<u64>()
                    .ok()
                    .map(|tick| (k.clone(), tick))
            })
            .collect();

        if env_entries.len() <= self.max_env_entries {
            return;
        }

        // tick でソート（古い順）
        env_entries.sort_by_key(|(_, tick)| *tick);

        // 超過分を削除
        let remove_count = env_entries.len() - self.max_env_entries;
        for (key, _) in env_entries.into_iter().take(remove_count) {
            self.kv.remove(&key);
        }
    }

    /// env エントリの最大保持数を設定
    pub fn set_max_env_entries(&mut self, max: usize) {
        self.max_env_entries = max;
    }
}

/// 完了した非同期タスクのメタ情報
///
/// payload は含まず、完了通知のみ。Worker は ReadPayload Action で payload を取得する。
#[derive(Debug, Clone)]
pub struct CompletedAsyncTask {
    /// タスクID
    pub task_id: TaskId,
    /// 発行した Worker（None = Manager 等から発行）
    pub worker_id: Option<WorkerId>,
    /// タスク種別（"web_search", "llm_call" など）
    pub task_type: String,
    /// 完了した Tick
    pub completed_at_tick: u64,
    /// ステータス
    pub status: TaskStatus,
    /// エラーメッセージ（失敗時）
    pub error: Option<String>,
}

/// Worker の永続状態コンテナ - GPU的なメモリ配置
pub struct WorkerStates {
    /// 連続メモリ領域に全 Worker の State を配置
    states: Vec<WorkerState>,
}

impl WorkerStates {
    pub fn new(count: usize) -> Self {
        let states = (0..count).map(|i| WorkerState::new(AgentId(i))).collect();
        Self { states }
    }

    /// Worker は自分の State のみ mutable アクセス可能
    pub fn get_mut(&mut self, id: AgentId) -> Option<&mut WorkerState> {
        self.states.get_mut(id.0)
    }

    /// 他 Worker の State は ReadOnly
    pub fn get(&self, id: AgentId) -> Option<&WorkerState> {
        self.states.get(id.0)
    }

    /// Worker 数を取得
    pub fn len(&self) -> usize {
        self.states.len()
    }

    /// 空かどうか
    pub fn is_empty(&self) -> bool {
        self.states.is_empty()
    }

    /// イテレーション
    pub fn iter(&self) -> impl Iterator<Item = &WorkerState> {
        self.states.iter()
    }

    /// 可変イテレーション
    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut WorkerState> {
        self.states.iter_mut()
    }

    /// 並列イテレーション（Rayon等で並列処理）
    pub fn par_iter_mut(&mut self) -> impl ParallelIterator<Item = &mut WorkerState> {
        self.states.par_iter_mut()
    }
}

/// Escalation 理由
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum EscalationReason {
    /// 連続失敗
    ConsecutiveFailures(u32),
    /// リソース不足
    ResourceExhausted,
    /// タイムアウト
    Timeout,
    /// Agent からの申告（Agent が自発的に介入を要求）
    AgentRequested(String),
    /// 不明なエラー
    Unknown(String),
}

/// Escalation 情報
#[derive(Debug, Clone)]
pub struct Escalation {
    /// 理由
    pub reason: EscalationReason,
    /// 発生 tick
    pub raised_at_tick: u64,
    /// 追加コンテキスト
    pub context: Option<String>,
}

impl Escalation {
    pub fn consecutive_failures(count: u32, tick: u64) -> Self {
        Self {
            reason: EscalationReason::ConsecutiveFailures(count),
            raised_at_tick: tick,
            context: None,
        }
    }

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

/// Worker の永続状態（Tick間で保持される）
///
/// WorkerStateDelta 経由で更新される。
pub struct WorkerState {
    /// Agent ID（内部識別用）
    pub id: AgentId,
    /// 内部状態（Worker 固有）
    internal_state: Option<Box<dyn Any + Send + Sync>>,
    /// 行動履歴
    pub history: ActionHistory,
    /// ローカルキャッシュ
    pub cache: LocalCache,
    /// 保留中の非同期Task（HashSet で O(1) 検索）
    pub pending_tasks: HashSet<TaskId>,
    /// Escalation 情報（Worker が設定、Manager が読み取り）
    pub escalation: Option<Escalation>,
    /// 連続失敗カウント
    pub consecutive_failures: u32,
    /// 最新アクションの出力（Environment からの結果）
    pub last_output: Option<String>,
}

impl WorkerState {
    pub fn new(id: AgentId) -> Self {
        Self {
            id,
            internal_state: None,
            history: ActionHistory::default(),
            cache: LocalCache::default(),
            pending_tasks: HashSet::new(),
            escalation: None,
            consecutive_failures: 0,
            last_output: None,
        }
    }

    /// Escalation を発生させる
    pub fn raise_escalation(&mut self, escalation: Escalation) {
        self.escalation = Some(escalation);
    }

    /// Escalation をクリア
    pub fn clear_escalation(&mut self) {
        self.escalation = None;
        self.consecutive_failures = 0;
    }

    /// 失敗を記録し、閾値を超えたら Escalation
    pub fn record_failure(&mut self, tick: u64, threshold: u32) -> bool {
        self.consecutive_failures += 1;
        if self.consecutive_failures >= threshold && self.escalation.is_none() {
            self.raise_escalation(Escalation::consecutive_failures(
                self.consecutive_failures,
                tick,
            ));
            true
        } else {
            false
        }
    }

    /// 成功を記録（連続失敗をリセット）
    pub fn record_success(&mut self) {
        self.consecutive_failures = 0;
    }

    /// 内部状態を設定
    pub fn set_state<T: Any + Send + Sync + 'static>(&mut self, state: T) {
        self.internal_state = Some(Box::new(state));
    }

    /// 内部状態を取得
    pub fn get_state<T: Any + Send + Sync + 'static>(&self) -> Option<&T> {
        self.internal_state.as_ref()?.downcast_ref()
    }

    /// 内部状態を可変で取得
    pub fn get_state_mut<T: Any + Send + Sync + 'static>(&mut self) -> Option<&mut T> {
        self.internal_state.as_mut()?.downcast_mut()
    }

    /// 非同期タスクを追加（O(1)）
    pub fn add_pending_task(&mut self, task_id: TaskId) {
        self.pending_tasks.insert(task_id);
    }

    /// 完了したタスクを削除（O(1)）
    pub fn complete_task(&mut self, task_id: TaskId) {
        self.pending_tasks.remove(&task_id);
    }
}

/// 行動履歴
///
/// VecDeque を使用したリングバッファ実装。
/// push() が O(1) で効率的に動作する。
pub struct ActionHistory {
    /// 最近のアクション（リングバッファ）
    entries: VecDeque<HistoryEntry>,
    /// 最大保持数
    max_entries: usize,
}

impl Default for ActionHistory {
    fn default() -> Self {
        Self::new(100) // デフォルト: 100エントリ
    }
}

impl ActionHistory {
    pub fn new(max_entries: usize) -> Self {
        Self {
            entries: VecDeque::with_capacity(max_entries),
            max_entries,
        }
    }

    /// エントリを追加（O(1)）
    pub fn push(&mut self, entry: HistoryEntry) {
        if self.max_entries > 0 && self.entries.len() >= self.max_entries {
            self.entries.pop_front(); // O(1) - VecDeque の利点
        }
        self.entries.push_back(entry);
    }

    /// 最新のエントリを取得
    pub fn latest(&self) -> Option<&HistoryEntry> {
        self.entries.back()
    }

    /// エントリ数を取得
    pub fn len(&self) -> usize {
        self.entries.len()
    }

    /// 空かどうか
    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    /// イテレータを取得
    pub fn iter(&self) -> impl Iterator<Item = &HistoryEntry> {
        self.entries.iter()
    }
}

/// 履歴エントリ
#[derive(Debug, Clone)]
pub struct HistoryEntry {
    pub tick: u64,
    pub action_name: String,
    pub success: bool,
}

/// ローカルキャッシュ
#[derive(Default)]
pub struct LocalCache {
    /// key-value キャッシュ
    data: HashMap<String, CacheEntry>,
}

impl LocalCache {
    /// キャッシュに設定
    pub fn set(&mut self, key: impl Into<String>, value: Vec<u8>, ttl_ticks: u64) {
        self.data.insert(
            key.into(),
            CacheEntry {
                value,
                expires_at_tick: ttl_ticks,
            },
        );
    }

    /// キャッシュから取得
    pub fn get(&self, key: &str, current_tick: u64) -> Option<&[u8]> {
        let entry = self.data.get(key)?;
        if entry.expires_at_tick > current_tick {
            Some(&entry.value)
        } else {
            None
        }
    }

    /// 期限切れエントリを削除
    pub fn cleanup(&mut self, current_tick: u64) {
        self.data.retain(|_, v| v.expires_at_tick > current_tick);
    }
}

/// キャッシュエントリ
struct CacheEntry {
    value: Vec<u8>,
    expires_at_tick: u64,
}

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

    #[test]
    fn test_swarm_state_creation() {
        let state = SwarmState::new(3);
        assert_eq!(state.workers.len(), 3);
        assert_eq!(state.shared.tick, 0);
    }

    #[test]
    fn test_swarm_state_advance_tick() {
        let mut state = SwarmState::new(1);
        assert_eq!(state.shared.tick, 0);

        state.advance_tick();
        assert_eq!(state.shared.tick, 1);

        state.advance_tick();
        assert_eq!(state.shared.tick, 2);
    }

    #[test]
    fn test_worker_states_access() {
        let mut states = WorkerStates::new(3);
        assert_eq!(states.len(), 3);
        assert!(!states.is_empty());

        // get_mut でアクセス
        let ws = states.get_mut(AgentId(1)).unwrap();
        assert_eq!(ws.id.0, 1);

        // 存在しない ID
        assert!(states.get(AgentId(10)).is_none());
    }

    #[test]
    fn test_worker_state_internal() {
        let mut ws = WorkerState::new(AgentId(0));

        // 初期状態は None
        assert!(ws.get_state::<i32>().is_none());

        // 状態を設定
        ws.set_state(42i32);
        assert_eq!(ws.get_state::<i32>(), Some(&42));

        // 可変アクセス
        if let Some(state) = ws.get_state_mut::<i32>() {
            *state = 100;
        }
        assert_eq!(ws.get_state::<i32>(), Some(&100));

        // 型が違う場合は None
        assert!(ws.get_state::<String>().is_none());
    }

    #[test]
    fn test_worker_state_pending_tasks() {
        let mut ws = WorkerState::new(AgentId(0));
        assert!(ws.pending_tasks.is_empty());

        ws.add_pending_task(TaskId(1));
        ws.add_pending_task(TaskId(2));
        assert_eq!(ws.pending_tasks.len(), 2);
        assert!(ws.pending_tasks.contains(&TaskId(1)));
        assert!(ws.pending_tasks.contains(&TaskId(2)));

        ws.complete_task(TaskId(1));
        assert_eq!(ws.pending_tasks.len(), 1);
        assert!(!ws.pending_tasks.contains(&TaskId(1)));
        assert!(ws.pending_tasks.contains(&TaskId(2)));
    }

    #[test]
    fn test_action_history() {
        let mut history = ActionHistory::new(3);

        history.push(HistoryEntry {
            tick: 0,
            action_name: "action1".to_string(),
            success: true,
        });
        history.push(HistoryEntry {
            tick: 1,
            action_name: "action2".to_string(),
            success: false,
        });

        assert_eq!(history.len(), 2);
        assert_eq!(history.latest().unwrap().action_name, "action2");

        // 最大数を超えると古いものが削除される
        history.push(HistoryEntry {
            tick: 2,
            action_name: "action3".to_string(),
            success: true,
        });
        history.push(HistoryEntry {
            tick: 3,
            action_name: "action4".to_string(),
            success: true,
        });

        assert_eq!(history.len(), 3);
        // action1 は削除され、action2 が先頭になる
        let entries: Vec<_> = history.iter().collect();
        assert_eq!(entries[0].action_name, "action2");
    }

    #[test]
    fn test_local_cache() {
        let mut cache = LocalCache::default();

        cache.set("key1", vec![1, 2, 3], 10);
        cache.set("key2", vec![4, 5, 6], 5);

        // 有効期限内
        assert_eq!(cache.get("key1", 0), Some([1u8, 2, 3].as_slice()));
        assert_eq!(cache.get("key2", 4), Some([4u8, 5, 6].as_slice()));

        // 有効期限切れ
        assert!(cache.get("key2", 5).is_none());
        assert!(cache.get("key2", 10).is_none());

        // key1 はまだ有効
        assert_eq!(cache.get("key1", 9), Some([1u8, 2, 3].as_slice()));

        // cleanup
        cache.cleanup(6);
        assert!(cache.get("key1", 0).is_some()); // まだ残っている
        cache.cleanup(11);
        assert!(cache.get("key1", 0).is_none()); // 削除された
    }

    #[test]
    fn test_environment() {
        let mut env = Environment::default();
        env.variables
            .insert("PATH".to_string(), "/usr/bin".to_string());
        env.flags.insert("debug".to_string(), true);

        assert_eq!(env.variables.get("PATH"), Some(&"/usr/bin".to_string()));
        assert_eq!(env.flags.get("debug"), Some(&true));
    }
}