paladin-ai-core 0.5.1

//! Citadel State Persistence Domain Types
//!
//! This module defines the domain types for the Citadel state persistence system,
//! which enables automatic saving and restoration of Paladin agents and Battalion
//! orchestrations to the file system as JSON files.
//!
//! # Overview
//!
//! The Citadel acts as a safety mechanism ensuring that long-running agent processes,
//! complex multi-agent workflows, and valuable conversation contexts are never lost
//! due to system failures, restarts, or intentional shutdown.
//!
//! # Key Types
//!
//! - [`PaladinState`]: Complete serializable state of a Paladin agent
//! - [`BattalionState`]: Complete serializable state of a Battalion orchestration
//! - [`StateSummary`]: Summary information for listing saved states
//! - [`CheckpointData`]: Tracking data for Battalion resumption
//!
//! # Example
//!
//! ```rust,no_run
//! use paladin_core::platform::container::citadel::PaladinState;
//! use chrono::Utc;
//!
//! // Create a Paladin state for persistence
//! // Note: In practice, this would be constructed from an actual Paladin
//! // let state = PaladinState {
//! //     paladin: my_paladin,
//! //     garrison: garrison_entries,
//! //     execution_history: history,
//! //     created_at: Utc::now(),
//! //     updated_at: Utc::now(),
//! //     schema_version: "1.0.0".to_string(),
//! // };
//! ```

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use uuid::Uuid;

use super::garrison::GarrisonEntry;
use super::paladin::MaxLoops;
use crate::base::entity::node::Node;

/// Type alias for Paladin using the Node pattern
pub type Paladin = Node<PaladinData>;

/// Placeholder for PaladinData - should match Epic 1 implementation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PaladinData {
    pub system_prompt: String,
    pub name: String,
    pub user_name: String,
    pub model: String,
    pub temperature: f32,
    pub max_loops: MaxLoops,
    pub stop_words: Vec<String>,
    pub status: PaladinStatus,
    pub vision_enabled: bool,
    pub autonomous_planning: bool,
    pub autonomous_prompts: bool,
    pub agent_description: String,
    pub dynamic_temperature: bool,
}

impl Default for PaladinData {
    fn default() -> Self {
        Self {
            system_prompt: String::new(),
            name: String::new(),
            user_name: String::new(),
            model: String::new(),
            temperature: 0.7,
            max_loops: MaxLoops::Fixed(3),
            stop_words: Vec::new(),
            status: PaladinStatus::Idle,
            vision_enabled: false,
            autonomous_planning: false,
            autonomous_prompts: false,
            agent_description: String::new(),
            dynamic_temperature: false,
        }
    }
}

/// Status of a Paladin agent
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum PaladinStatus {
    Idle,
    Executing,
    Completed,
    Failed(String),
}

/// Serializable Paladin state for persistence
///
/// Contains the complete state of a Paladin agent including its configuration,
/// conversation history (Garrison), execution records, and metadata timestamps.
///
/// # Fields
///
/// - `paladin`: The Paladin entity with its configuration
/// - `garrison`: Complete conversation history as Garrison entries
/// - `execution_history`: Record of all executions for debugging/audit
/// - `created_at`: Timestamp when state was first created
/// - `updated_at`: Timestamp of last state update
/// - `schema_version`: Version identifier for state schema compatibility
///
/// # Example
///
/// ```rust,no_run
/// use paladin_core::platform::container::citadel::PaladinState;
/// use chrono::Utc;
///
/// // States are typically created by the persistence system
/// // let state = PaladinState {
/// //     paladin: my_paladin,
/// //     garrison: vec![],
/// //     execution_history: vec![],
/// //     created_at: Utc::now(),
/// //     updated_at: Utc::now(),
/// //     schema_version: "1.0.0".to_string(),
/// // };
/// ```
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PaladinState {
    /// The Paladin entity with its configuration
    pub paladin: Paladin,
    /// Complete conversation history
    pub garrison: Vec<GarrisonEntry>,
    /// Record of all executions
    pub execution_history: Vec<ExecutionRecord>,
    /// Timestamp when state was first created
    pub created_at: DateTime<Utc>,
    /// Timestamp of last state update
    pub updated_at: DateTime<Utc>,
    /// Schema version for compatibility tracking
    pub schema_version: String,
}

impl PaladinState {
    /// Creates a new PaladinState with the current timestamp
    ///
    /// # Arguments
    ///
    /// * `paladin` - The Paladin entity to persist
    /// * `garrison` - Conversation history entries
    /// * `execution_history` - Record of past executions
    ///
    /// # Returns
    ///
    /// A new `PaladinState` with timestamps set to current time and schema version "1.0.0"
    pub fn new(
        paladin: Paladin,
        garrison: Vec<GarrisonEntry>,
        execution_history: Vec<ExecutionRecord>,
    ) -> Self {
        let now = Utc::now();
        Self {
            paladin,
            garrison,
            execution_history,
            created_at: now,
            updated_at: now,
            schema_version: "1.0.0".to_string(),
        }
    }

    /// Updates the state with new data and refreshes the updated_at timestamp
    pub fn update(
        &mut self,
        garrison: Vec<GarrisonEntry>,
        execution_history: Vec<ExecutionRecord>,
    ) {
        self.garrison = garrison;
        self.execution_history = execution_history;
        self.updated_at = Utc::now();
    }
}

/// Record of a single Paladin execution
///
/// Captures the input, output, status, and timing information for
/// debugging and audit purposes.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExecutionRecord {
    /// When the execution occurred
    pub timestamp: DateTime<Utc>,
    /// Input provided to the Paladin
    pub input: String,
    /// Output generated by the Paladin
    pub output: String,
    /// Execution status (Success, Failed, etc.)
    pub status: ExecutionStatus,
    /// Number of reasoning loops performed
    pub loops_used: u32,
}

/// Status of an execution
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum ExecutionStatus {
    Success,
    Failed(String),
    Timeout,
    StopWordDetected,
}

/// Serializable Battalion state for persistence
///
/// Contains the complete state of a Battalion orchestration including
/// the type of Battalion, configuration, states of all constituent Paladins,
/// and checkpoint data for resumption.
///
/// # Fields
///
/// - `id`: Unique identifier for the Battalion
/// - `battalion_type`: Type of Battalion (Formation, Phalanx, Campaign, ChainOfCommand)
/// - `config`: Battalion configuration parameters
/// - `paladin_states`: Complete state of all constituent Paladins
/// - `checkpoint`: Current execution checkpoint for resumption
/// - `created_at`: Timestamp when state was first created
/// - `updated_at`: Timestamp of last state update
/// - `schema_version`: Version identifier for state schema compatibility
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BattalionState {
    /// Unique identifier for this Battalion
    pub id: Uuid,
    /// Type of Battalion orchestration
    pub battalion_type: String,
    /// Battalion configuration parameters
    pub config: BattalionConfig,
    /// States of all constituent Paladins
    pub paladin_states: Vec<PaladinState>,
    /// Current checkpoint for resumption
    pub checkpoint: Option<CheckpointData>,
    /// Timestamp when state was first created
    pub created_at: DateTime<Utc>,
    /// Timestamp of last state update
    pub updated_at: DateTime<Utc>,
    /// Schema version for compatibility tracking
    pub schema_version: String,
}

impl BattalionState {
    /// Creates a new BattalionState
    ///
    /// # Arguments
    ///
    /// * `battalion_type` - Type of Battalion (e.g., "Formation", "Phalanx")
    /// * `config` - Battalion configuration
    /// * `paladin_states` - States of constituent Paladins
    /// * `checkpoint` - Optional checkpoint data
    pub fn new(
        battalion_type: impl Into<String>,
        config: BattalionConfig,
        paladin_states: Vec<PaladinState>,
        checkpoint: Option<CheckpointData>,
    ) -> Self {
        let now = Utc::now();
        Self {
            id: Uuid::new_v4(),
            battalion_type: battalion_type.into(),
            config,
            paladin_states,
            checkpoint,
            created_at: now,
            updated_at: now,
            schema_version: "1.0.0".to_string(),
        }
    }
}

/// Configuration parameters for Battalion orchestration
///
/// Contains settings that control how a Battalion executes its Paladins.
/// This is a placeholder and will be expanded in Epic 4.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct BattalionConfig {
    /// Maximum concurrent Paladins (for Phalanx)
    #[serde(default)]
    pub max_concurrency: Option<usize>,
    /// Timeout for entire Battalion execution
    #[serde(default)]
    pub timeout_seconds: Option<u64>,
    /// Whether to continue on Paladin failure
    #[serde(default)]
    pub continue_on_error: bool,
}

/// Checkpoint data for Battalion execution resumption
///
/// Tracks the current progress through a Battalion workflow, enabling
/// resumption from the last successfully completed Paladin.
///
/// # Fields
///
/// - `last_completed_index`: Index of last successfully completed Paladin
/// - `completed_paladins`: IDs of all completed Paladins
/// - `failed_paladins`: IDs of any failed Paladins
/// - `checkpoint_time`: When this checkpoint was created
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CheckpointData {
    /// Index of last successfully completed Paladin in sequence
    pub last_completed_index: Option<usize>,
    /// IDs of all successfully completed Paladins
    pub completed_paladins: Vec<Uuid>,
    /// IDs of any failed Paladins
    pub failed_paladins: Vec<Uuid>,
    /// When this checkpoint was created
    pub checkpoint_time: DateTime<Utc>,
}

impl CheckpointData {
    /// Creates a new empty checkpoint
    pub fn new() -> Self {
        Self {
            last_completed_index: None,
            completed_paladins: Vec::new(),
            failed_paladins: Vec::new(),
            checkpoint_time: Utc::now(),
        }
    }

    /// Records a completed Paladin in the checkpoint
    pub fn mark_completed(&mut self, paladin_id: Uuid, index: usize) {
        self.completed_paladins.push(paladin_id);
        self.last_completed_index = Some(index);
        self.checkpoint_time = Utc::now();
    }

    /// Records a failed Paladin in the checkpoint
    pub fn mark_failed(&mut self, paladin_id: Uuid) {
        self.failed_paladins.push(paladin_id);
        self.checkpoint_time = Utc::now();
    }
}

impl Default for CheckpointData {
    fn default() -> Self {
        Self::new()
    }
}

/// Summary information for a saved state
///
/// Provides metadata about a saved state file without loading
/// the complete state data.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StateSummary {
    /// Unique identifier for the state
    pub id: Uuid,
    /// Type of state (Paladin or Battalion)
    pub state_type: StateType,
    /// When the state was created
    pub created_at: DateTime<Utc>,
    /// When the state was last updated
    pub updated_at: DateTime<Utc>,
    /// File path to the saved state
    pub file_path: PathBuf,
}

/// Type of saved state
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum StateType {
    /// Paladin agent state
    Paladin,
    /// Battalion orchestration state
    Battalion,
}

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

    fn create_test_paladin_data() -> PaladinData {
        PaladinData {
            system_prompt: "You are a helpful assistant".to_string(),
            name: "TestPaladin".to_string(),
            user_name: "TestUser".to_string(),
            model: "gpt-4".to_string(),
            temperature: 0.7,
            max_loops: MaxLoops::Fixed(3),
            stop_words: vec!["STOP".to_string()],
            status: PaladinStatus::Idle,
            vision_enabled: false,
            ..Default::default()
        }
    }

    fn create_test_paladin() -> Paladin {
        Node::new(create_test_paladin_data(), Some("TestPaladin".to_string()))
    }

    #[test]
    fn test_paladin_state_creation() {
        let paladin = create_test_paladin();
        let garrison = vec![];
        let history = vec![];

        let state = PaladinState::new(paladin, garrison, history);

        assert_eq!(state.schema_version, "1.0.0");
        assert_eq!(state.garrison.len(), 0);
        assert_eq!(state.execution_history.len(), 0);
        assert!(state.created_at <= Utc::now());
        assert!(state.updated_at <= Utc::now());
    }

    #[test]
    fn test_paladin_state_serialization_roundtrip() {
        let paladin = create_test_paladin();
        let garrison = vec![];
        let history = vec![ExecutionRecord {
            timestamp: Utc::now(),
            input: "test input".to_string(),
            output: "test output".to_string(),
            status: ExecutionStatus::Success,
            loops_used: 1,
        }];

        let state = PaladinState::new(paladin, garrison, history);

        // Serialize to JSON
        let json = serde_json::to_string(&state).expect("Failed to serialize");

        // Deserialize back
        let deserialized: PaladinState =
            serde_json::from_str(&json).expect("Failed to deserialize");

        // Verify critical fields match
        assert_eq!(deserialized.schema_version, state.schema_version);
        assert_eq!(deserialized.execution_history.len(), 1);
        assert_eq!(deserialized.execution_history[0].input, "test input");
    }

    #[test]
    fn test_battalion_state_creation() {
        let paladin = create_test_paladin();
        let paladin_state = PaladinState::new(paladin, vec![], vec![]);
        let config = BattalionConfig::default();

        let battalion_state = BattalionState::new("Formation", config, vec![paladin_state], None);

        assert_eq!(battalion_state.battalion_type, "Formation");
        assert_eq!(battalion_state.schema_version, "1.0.0");
        assert_eq!(battalion_state.paladin_states.len(), 1);
        assert!(battalion_state.checkpoint.is_none());
    }

    #[test]
    fn test_battalion_state_serialization_roundtrip() {
        let paladin = create_test_paladin();
        let paladin_state = PaladinState::new(paladin, vec![], vec![]);
        let config = BattalionConfig {
            max_concurrency: Some(4),
            timeout_seconds: Some(300),
            continue_on_error: true,
        };
        let checkpoint = Some(CheckpointData::new());

        let battalion_state =
            BattalionState::new("Phalanx", config, vec![paladin_state], checkpoint);

        // Serialize to JSON
        let json = serde_json::to_string(&battalion_state).expect("Failed to serialize");

        // Deserialize back
        let deserialized: BattalionState =
            serde_json::from_str(&json).expect("Failed to deserialize");

        // Verify critical fields match
        assert_eq!(deserialized.battalion_type, "Phalanx");
        assert_eq!(deserialized.schema_version, "1.0.0");
        assert!(deserialized.checkpoint.is_some());
    }

    #[test]
    fn test_state_summary_creation() {
        let summary = StateSummary {
            id: Uuid::new_v4(),
            state_type: StateType::Paladin,
            created_at: Utc::now(),
            updated_at: Utc::now(),
            file_path: PathBuf::from("/tmp/test-state.json"),
        };

        assert_eq!(summary.state_type, StateType::Paladin);
        assert_eq!(summary.file_path, PathBuf::from("/tmp/test-state.json"));
    }

    #[test]
    fn test_checkpoint_data_serialization() {
        let mut checkpoint = CheckpointData::new();
        let paladin_id = Uuid::new_v4();
        checkpoint.mark_completed(paladin_id, 0);

        // Serialize to JSON
        let json = serde_json::to_string(&checkpoint).expect("Failed to serialize");

        // Deserialize back
        let deserialized: CheckpointData =
            serde_json::from_str(&json).expect("Failed to deserialize");

        assert_eq!(deserialized.last_completed_index, Some(0));
        assert_eq!(deserialized.completed_paladins.len(), 1);
        assert_eq!(deserialized.completed_paladins[0], paladin_id);
    }

    #[test]
    fn test_schema_version_field_present() {
        let paladin = create_test_paladin();
        let state = PaladinState::new(paladin, vec![], vec![]);

        // Serialize and check JSON contains schema_version
        let json = serde_json::to_string(&state).expect("Failed to serialize");
        assert!(json.contains("schema_version"));
        assert!(json.contains("1.0.0"));
    }

    #[test]
    fn test_json_output_human_readable() {
        let paladin = create_test_paladin();
        let state = PaladinState::new(paladin, vec![], vec![]);

        // Serialize with pretty printing
        let json = serde_json::to_string_pretty(&state).expect("Failed to serialize");

        // Should be multi-line (pretty printed)
        assert!(json.contains('\n'));
        // Should contain readable field names
        assert!(json.contains("\"schema_version\""));
        assert!(json.contains("\"created_at\""));
        assert!(json.contains("\"updated_at\""));
    }

    #[test]
    fn test_checkpoint_mark_completed() {
        let mut checkpoint = CheckpointData::new();
        let id1 = Uuid::new_v4();
        let id2 = Uuid::new_v4();

        checkpoint.mark_completed(id1, 0);
        checkpoint.mark_completed(id2, 1);

        assert_eq!(checkpoint.last_completed_index, Some(1));
        assert_eq!(checkpoint.completed_paladins.len(), 2);
        assert!(checkpoint.failed_paladins.is_empty());
    }

    #[test]
    fn test_checkpoint_mark_failed() {
        let mut checkpoint = CheckpointData::new();
        let id = Uuid::new_v4();

        checkpoint.mark_failed(id);

        assert_eq!(checkpoint.failed_paladins.len(), 1);
        assert_eq!(checkpoint.failed_paladins[0], id);
        assert!(checkpoint.completed_paladins.is_empty());
    }

    #[test]
    fn test_paladin_state_update() {
        let paladin = create_test_paladin();
        let mut state = PaladinState::new(paladin, vec![], vec![]);

        let original_updated = state.updated_at;

        // Small delay to ensure timestamp difference
        std::thread::sleep(std::time::Duration::from_millis(10));

        let new_history = vec![ExecutionRecord {
            timestamp: Utc::now(),
            input: "new input".to_string(),
            output: "new output".to_string(),
            status: ExecutionStatus::Success,
            loops_used: 2,
        }];

        state.update(vec![], new_history);

        assert!(state.updated_at > original_updated);
        assert_eq!(state.execution_history.len(), 1);
    }

    #[test]
    fn test_execution_status_variants() {
        let success = ExecutionStatus::Success;
        let failed = ExecutionStatus::Failed("error".to_string());
        let timeout = ExecutionStatus::Timeout;
        let stop = ExecutionStatus::StopWordDetected;

        // Test serialization of all variants
        assert!(serde_json::to_string(&success).is_ok());
        assert!(serde_json::to_string(&failed).is_ok());
        assert!(serde_json::to_string(&timeout).is_ok());
        assert!(serde_json::to_string(&stop).is_ok());
    }

    #[test]
    fn test_state_summary_serialization() {
        let summary = StateSummary {
            id: Uuid::new_v4(),
            state_type: StateType::Paladin,
            created_at: Utc::now(),
            updated_at: Utc::now(),
            file_path: PathBuf::from("/test/path.json"),
        };

        // Serialize to JSON
        let json = serde_json::to_string(&summary).expect("Failed to serialize StateSummary");

        // Deserialize back
        let deserialized: StateSummary =
            serde_json::from_str(&json).expect("Failed to deserialize StateSummary");

        assert_eq!(deserialized.id, summary.id);
        assert_eq!(deserialized.state_type, summary.state_type);
        assert_eq!(deserialized.file_path, summary.file_path);
    }

    #[test]
    fn test_state_type_serialization() {
        let paladin_type = StateType::Paladin;
        let battalion_type = StateType::Battalion;

        // Serialize both types
        let paladin_json = serde_json::to_string(&paladin_type).expect("Failed to serialize");
        let battalion_json = serde_json::to_string(&battalion_type).expect("Failed to serialize");

        // Deserialize back
        let deser_paladin: StateType =
            serde_json::from_str(&paladin_json).expect("Failed to deserialize");
        let deser_battalion: StateType =
            serde_json::from_str(&battalion_json).expect("Failed to deserialize");

        assert_eq!(deser_paladin, StateType::Paladin);
        assert_eq!(deser_battalion, StateType::Battalion);
    }

    #[test]
    fn test_all_domain_types_are_serializable() {
        // This test ensures all types used in state persistence are serializable
        let paladin = create_test_paladin();
        let garrison_entry = GarrisonEntry::new(
            crate::platform::container::garrison::ConversationRole::User,
            "test".to_string(),
        );
        let exec_record = ExecutionRecord {
            timestamp: Utc::now(),
            input: "test".to_string(),
            output: "result".to_string(),
            status: ExecutionStatus::Success,
            loops_used: 1,
        };
        let checkpoint = CheckpointData::new();
        let config = BattalionConfig::default();

        // All these should serialize successfully
        assert!(serde_json::to_string(&paladin).is_ok());
        assert!(serde_json::to_string(&garrison_entry).is_ok());
        assert!(serde_json::to_string(&exec_record).is_ok());
        assert!(serde_json::to_string(&checkpoint).is_ok());
        assert!(serde_json::to_string(&config).is_ok());
    }
}