bzzz-core 0.1.0

//! State Model v1
//!
//! Core state objects for Bzzz orchestration:
//! - Run: A single execution of an Agent or Swarm
//! - Step: A single step within a Run
//! - Session: A logical grouping of related Runs
//! - Artifact: Output produced by a Run or Step

use std::path::PathBuf;
use std::time::{Duration, Instant, SystemTime};

// ============================================================================
// ID Types
// ============================================================================

/// Unique identifier for a Run
#[derive(Debug, Clone, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
pub struct RunId(pub String);

impl RunId {
    /// Generate a new unique RunId using UUID v4
    pub fn new() -> Self {
        RunId(uuid::Uuid::new_v4().to_string())
    }

    /// Create a RunId from a string
    pub fn from_string(s: String) -> Self {
        RunId(s)
    }

    /// Get the inner string
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

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

/// Unique identifier for a Step within a Run
#[derive(Debug, Clone, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
pub struct StepId(pub String);

impl StepId {
    /// Generate a new unique StepId
    pub fn new() -> Self {
        StepId(uuid::Uuid::new_v4().to_string())
    }

    /// Get the inner string
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

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

/// Unique identifier for a Session
#[derive(Debug, Clone, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
pub struct SessionId(pub String);

impl SessionId {
    /// Generate a new unique SessionId
    pub fn new() -> Self {
        SessionId(uuid::Uuid::new_v4().to_string())
    }

    /// Get the inner string
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

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

/// Unique identifier for an Artifact
#[derive(Debug, Clone, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
pub struct ArtifactId(pub String);

impl ArtifactId {
    /// Generate a new unique ArtifactId
    pub fn new() -> Self {
        ArtifactId(uuid::Uuid::new_v4().to_string())
    }

    /// Get the inner string
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

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

// ============================================================================
// Run Types
// ============================================================================

/// Status of a Run
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum RunStatus {
    /// Execution has not started
    Pending,
    /// Execution is in progress
    Running,
    /// Execution completed successfully
    Completed,
    /// Execution failed with error
    Failed,
    /// Execution was cancelled by user
    Cancelled,
    /// Execution exceeded time limit
    Timeout,
}

/// Target of a Run - what is being executed
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum RunTarget {
    /// Execute a single Agent
    Agent {
        /// Path to AgentSpec file
        spec_path: PathBuf,
    },
    /// Execute a Swarm
    Swarm {
        /// Path to SwarmFile
        swarmfile_path: PathBuf,
    },
    /// Execute via A2A protocol (remote HTTP agent)
    A2AAgent {
        /// A2A agent endpoint URL
        url: String,
    },
}

/// Kind of runtime environment
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
pub enum RuntimeKind {
    /// Local subprocess execution
    Local,
    /// Docker container execution
    Docker,
    /// HTTP remote execution
    Http,
    /// Cloud managed execution
    Cloud,
}

/// Error from a Run execution
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error, serde::Serialize, serde::Deserialize)]
pub enum RunError {
    /// Failed to start execution
    #[error("Startup failed: {message}")]
    StartupFailed {
        /// Error message
        message: String,
    },
    /// Execution failed with error
    #[error("Execution failed: {message}")]
    ExecutionFailed {
        /// Error message
        message: String,
    },
    /// Execution was cancelled
    #[error("Cancelled: {reason}")]
    Cancelled {
        /// Cancellation reason
        reason: String,
    },
    /// Execution exceeded time limit
    #[error("Timeout after {after:?}")]
    Timeout {
        /// Time limit that was exceeded
        after: Duration,
    },
    /// Individual worker exceeded its per-worker time limit
    #[error("Worker '{worker}' timed out after {after:?}")]
    WorkerTimeout {
        /// Worker name that timed out
        worker: String,
        /// Time limit that was exceeded
        after: Duration,
    },
    /// Invalid configuration
    #[error("Invalid config: {message}")]
    InvalidConfig {
        /// Error message
        message: String,
    },
    /// Runtime error
    #[error("Runtime error: {message}")]
    RuntimeError {
        /// Error message
        message: String,
    },
    /// Artifact error
    #[error("Artifact error: {message}")]
    ArtifactError {
        /// Error message
        message: String,
    },
    /// Resource not found (e.g., run, container, agent)
    #[error("Not found: {resource_type} '{id}'")]
    NotFound {
        /// Type of resource
        resource_type: String,
        /// Resource identifier
        id: String,
    },
    /// Pattern execution error
    #[error("Pattern '{pattern}' error in step '{step}': {message}")]
    PatternError {
        /// Pattern type
        pattern: String,
        /// Step identifier
        step: String,
        /// Error message
        message: String,
    },
    /// Agent spec validation error
    #[error("Agent spec validation failed: {message}")]
    ValidationError {
        /// Error message
        message: String,
    },
}

impl RunError {
    /// Create a StartupFailed error
    pub fn startup_failed(message: impl Into<String>) -> Self {
        RunError::StartupFailed {
            message: message.into(),
        }
    }

    /// Create an ExecutionFailed error
    pub fn execution_failed(message: impl Into<String>) -> Self {
        RunError::ExecutionFailed {
            message: message.into(),
        }
    }

    /// Create a Cancelled error
    pub fn cancelled(reason: impl Into<String>) -> Self {
        RunError::Cancelled {
            reason: reason.into(),
        }
    }

    /// Create a Timeout error
    pub fn timeout(after: Duration) -> Self {
        RunError::Timeout { after }
    }

    /// Create a WorkerTimeout error
    pub fn worker_timeout(worker: impl Into<String>, after: Duration) -> Self {
        RunError::WorkerTimeout {
            worker: worker.into(),
            after,
        }
    }

    /// Create an InvalidConfig error
    pub fn invalid_config(message: impl Into<String>) -> Self {
        RunError::InvalidConfig {
            message: message.into(),
        }
    }

    /// Create a RuntimeError
    pub fn runtime_error(message: impl Into<String>) -> Self {
        RunError::RuntimeError {
            message: message.into(),
        }
    }

    /// Create a NotFound error
    pub fn not_found(resource_type: impl Into<String>, id: impl Into<String>) -> Self {
        RunError::NotFound {
            resource_type: resource_type.into(),
            id: id.into(),
        }
    }

    /// Create a PatternError
    pub fn pattern_error(
        pattern: impl Into<String>,
        step: impl Into<String>,
        message: impl Into<String>,
    ) -> Self {
        RunError::PatternError {
            pattern: pattern.into(),
            step: step.into(),
            message: message.into(),
        }
    }

    /// Create a ValidationError
    pub fn validation_error(message: impl Into<String>) -> Self {
        RunError::ValidationError {
            message: message.into(),
        }
    }

    /// Check if this error is retryable
    ///
    /// Retryable errors: Timeout, RuntimeError, PatternError
    /// Non-retryable: StartupFailed, InvalidConfig, ValidationError, NotFound, Cancelled
    pub fn is_retryable(&self) -> bool {
        matches!(
            self,
            RunError::Timeout { .. }
                | RunError::WorkerTimeout { .. }
                | RunError::RuntimeError { .. }
                | RunError::PatternError { .. }
        )
    }

    /// Check if this is a user-initiated cancellation
    pub fn is_cancellation(&self) -> bool {
        matches!(self, RunError::Cancelled { .. })
    }

    /// Check if this is a configuration/validation error
    pub fn is_config_error(&self) -> bool {
        matches!(
            self,
            RunError::InvalidConfig { .. } | RunError::ValidationError { .. }
        )
    }
}

/// A single execution of an Agent or Swarm
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct Run {
    /// Unique identifier for this run
    pub id: RunId,
    /// Current status
    pub status: RunStatus,
    /// What is being executed
    pub target: RunTarget,
    /// Runtime environment
    pub runtime_kind: RuntimeKind,
    /// Optional session this run belongs to
    pub session_id: Option<SessionId>,
    /// When the run was created
    pub created_at: SystemTime,
    /// When execution started
    pub started_at: Option<SystemTime>,
    /// When execution finished
    pub finished_at: Option<SystemTime>,
    /// Optional timeout duration
    pub timeout: Option<Duration>,
    /// Error if failed
    pub error: Option<RunError>,
    /// Steps in this run
    pub steps: Vec<Step>,
    /// Output artifacts
    pub artifacts: Vec<ArtifactId>,
    /// Input data for the worker (resolved template expressions, JSON-serialized)
    #[serde(default)]
    pub input: Option<serde_json::Value>,
}

impl Run {
    /// Create a new run
    pub fn new(target: RunTarget, runtime_kind: RuntimeKind) -> Self {
        Run {
            id: RunId::new(),
            status: RunStatus::Pending,
            target,
            runtime_kind,
            session_id: None,
            created_at: SystemTime::now(),
            started_at: None,
            finished_at: None,
            timeout: None,
            error: None,
            steps: Vec::new(),
            artifacts: Vec::new(),
            input: None,
        }
    }

    /// Set input data for the worker
    pub fn with_input(mut self, input: serde_json::Value) -> Self {
        self.input = Some(input);
        self
    }

    /// Add session association
    pub fn with_session(mut self, session_id: SessionId) -> Self {
        self.session_id = Some(session_id);
        self
    }

    /// Set timeout duration
    pub fn with_timeout(mut self, timeout: Duration) -> Self {
        self.timeout = Some(timeout);
        self
    }

    /// Mark execution as started
    pub fn start(&mut self) {
        self.status = RunStatus::Running;
        self.started_at = Some(SystemTime::now());
    }

    /// Mark execution as completed
    pub fn complete(&mut self) {
        self.status = RunStatus::Completed;
        self.finished_at = Some(SystemTime::now());
    }

    /// Mark execution as failed
    pub fn fail(&mut self, error: RunError) {
        self.status = RunStatus::Failed;
        self.error = Some(error);
        self.finished_at = Some(SystemTime::now());
    }

    /// Mark execution as cancelled
    pub fn cancel(&mut self, reason: String) {
        self.status = RunStatus::Cancelled;
        self.error = Some(RunError::Cancelled { reason });
        self.finished_at = Some(SystemTime::now());
    }

    /// Mark execution as timed out
    pub fn timeout(&mut self) {
        self.status = RunStatus::Timeout;
        if let Some(timeout) = self.timeout {
            self.error = Some(RunError::Timeout { after: timeout });
        }
        self.finished_at = Some(SystemTime::now());
    }

    /// Get execution duration if completed
    pub fn duration(&self) -> Option<Duration> {
        match (self.started_at, self.finished_at) {
            (Some(start), Some(end)) => end.duration_since(start).ok(),
            _ => None,
        }
    }
}

// ============================================================================
// Step Types
// ============================================================================

/// Status of a Step
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum StepStatus {
    /// Step has not started
    Pending,
    /// Step is in progress
    Running,
    /// Step completed successfully
    Completed,
    /// Step failed with error
    Failed,
    /// Step was skipped
    Skipped,
    /// Step was cancelled
    Cancelled,
}

/// Error from a Step execution
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error, serde::Serialize, serde::Deserialize)]
pub enum StepError {
    /// Agent execution failed
    #[error("Agent failed: {message}")]
    AgentFailed {
        /// Error message
        message: String,
    },
    /// Invalid input provided
    #[error("Invalid input: {message}")]
    InvalidInput {
        /// Error message
        message: String,
    },
    /// Output generation failed
    #[error("Output failed: {message}")]
    OutputFailed {
        /// Error message
        message: String,
    },
    /// Step exceeded time limit
    #[error("Timeout after {after:?}")]
    Timeout {
        /// Time limit that was exceeded
        after: Duration,
    },
    /// Step was cancelled
    #[error("Cancelled: {reason}")]
    Cancelled {
        /// Cancellation reason
        reason: String,
    },
    /// Pattern execution error
    #[error("Pattern '{pattern}' error: {message}")]
    PatternError {
        /// Pattern name
        pattern: String,
        /// Error message
        message: String,
    },
}

impl StepError {
    /// Create an AgentFailed error
    pub fn agent_failed(message: impl Into<String>) -> Self {
        StepError::AgentFailed {
            message: message.into(),
        }
    }

    /// Create an InvalidInput error
    pub fn invalid_input(message: impl Into<String>) -> Self {
        StepError::InvalidInput {
            message: message.into(),
        }
    }

    /// Create an OutputFailed error
    pub fn output_failed(message: impl Into<String>) -> Self {
        StepError::OutputFailed {
            message: message.into(),
        }
    }

    /// Create a Timeout error
    pub fn timeout(after: Duration) -> Self {
        StepError::Timeout { after }
    }

    /// Create a Cancelled error
    pub fn cancelled(reason: impl Into<String>) -> Self {
        StepError::Cancelled {
            reason: reason.into(),
        }
    }

    /// Create a PatternError
    pub fn pattern_error(pattern: impl Into<String>, message: impl Into<String>) -> Self {
        StepError::PatternError {
            pattern: pattern.into(),
            message: message.into(),
        }
    }

    /// Check if this error is retryable
    pub fn is_retryable(&self) -> bool {
        matches!(
            self,
            StepError::Timeout { .. } | StepError::PatternError { .. }
        )
    }

    /// Check if this is a cancellation
    pub fn is_cancellation(&self) -> bool {
        matches!(self, StepError::Cancelled { .. })
    }
}

/// A single step within a Run
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct Step {
    /// Unique identifier for this step
    pub id: StepId,
    /// Run this step belongs to
    pub run_id: RunId,
    /// Current status
    pub status: StepStatus,
    /// Step name
    pub name: String,
    /// Path to agent spec
    pub agent_spec_path: PathBuf,
    /// When step was created
    pub created_at: SystemTime,
    /// When step started
    pub started_at: Option<SystemTime>,
    /// When step finished
    pub finished_at: Option<SystemTime>,
    /// Error if failed
    pub error: Option<StepError>,
    /// Output artifacts
    pub output_artifacts: Vec<ArtifactId>,
}

impl Step {
    /// Create a new step
    pub fn new(run_id: RunId, name: String, agent_spec_path: PathBuf) -> Self {
        Step {
            id: StepId::new(),
            run_id,
            status: StepStatus::Pending,
            name,
            agent_spec_path,
            created_at: SystemTime::now(),
            started_at: None,
            finished_at: None,
            error: None,
            output_artifacts: Vec::new(),
        }
    }

    /// Mark step as started
    pub fn start(&mut self) {
        self.status = StepStatus::Running;
        self.started_at = Some(SystemTime::now());
    }

    /// Mark step as completed with artifacts
    pub fn complete(&mut self, artifacts: Vec<ArtifactId>) {
        self.status = StepStatus::Completed;
        self.output_artifacts = artifacts;
        self.finished_at = Some(SystemTime::now());
    }

    /// Mark step as failed
    pub fn fail(&mut self, error: StepError) {
        self.status = StepStatus::Failed;
        self.error = Some(error);
        self.finished_at = Some(SystemTime::now());
    }

    /// Mark step as skipped
    pub fn skip(&mut self) {
        self.status = StepStatus::Skipped;
        self.finished_at = Some(SystemTime::now());
    }

    /// Mark step as cancelled
    pub fn cancel(&mut self, reason: String) {
        self.status = StepStatus::Cancelled;
        self.error = Some(StepError::Cancelled { reason });
        self.finished_at = Some(SystemTime::now());
    }
}

// ============================================================================
// Session Types
// ============================================================================

/// Status of a Session
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum SessionStatus {
    /// Session is active
    Active,
    /// Session has ended
    Ended,
    /// Session was cancelled
    Cancelled,
}

/// A logical grouping of related Runs
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct Session {
    /// Unique identifier
    pub id: SessionId,
    /// Current status
    pub status: SessionStatus,
    /// Session name
    pub name: String,
    /// When session was created
    pub created_at: SystemTime,
    /// When session ended
    pub ended_at: Option<SystemTime>,
    /// Runs in this session
    pub run_ids: Vec<RunId>,
}

impl Session {
    /// Create a new session
    pub fn new(name: String) -> Self {
        Session {
            id: SessionId::new(),
            status: SessionStatus::Active,
            name,
            created_at: SystemTime::now(),
            ended_at: None,
            run_ids: Vec::new(),
        }
    }

    /// Add a run to this session
    pub fn add_run(&mut self, run_id: RunId) {
        self.run_ids.push(run_id);
    }

    /// End the session
    pub fn end(&mut self) {
        self.status = SessionStatus::Ended;
        self.ended_at = Some(SystemTime::now());
    }

    /// Cancel the session
    pub fn cancel(&mut self) {
        self.status = SessionStatus::Cancelled;
        self.ended_at = Some(SystemTime::now());
    }
}

// ============================================================================
// Artifact Types
// ============================================================================

/// Location where an artifact is stored
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum ArtifactLocation {
    /// Local filesystem path
    Local {
        /// File path
        path: PathBuf,
    },
    /// HTTP URL
    Http {
        /// URL to fetch artifact
        url: String,
    },
    /// Cloud storage URI
    CloudStorage {
        /// Cloud storage URI
        uri: String,
    },
    /// Inline content
    Inline {
        /// Raw content bytes
        content: Vec<u8>,
    },
}

/// Checksum for artifact integrity
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct Checksum {
    /// Algorithm used (e.g., "sha256")
    pub algorithm: String,
    /// Checksum value
    pub value: String,
}

impl Checksum {
    /// Create a new checksum
    pub fn new(algorithm: impl Into<String>, value: impl Into<String>) -> Self {
        Checksum {
            algorithm: algorithm.into(),
            value: value.into(),
        }
    }

    /// Create a SHA-256 checksum
    pub fn sha256(value: impl Into<String>) -> Self {
        Checksum::new("sha256", value)
    }
}

/// Output produced by a Run or Step
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct Artifact {
    /// Unique identifier
    pub id: ArtifactId,
    /// Artifact name
    pub name: String,
    /// MIME content type
    pub content_type: String,
    /// Size in bytes
    pub size: u64,
    /// Where the artifact is stored
    pub location: ArtifactLocation,
    /// Optional integrity checksum
    pub checksum: Option<Checksum>,
    /// When artifact was created
    pub created_at: SystemTime,
    /// Run that produced this artifact
    pub source_run: Option<RunId>,
    /// Step that produced this artifact
    pub source_step: Option<StepId>,
}

impl Artifact {
    /// Create a new artifact
    pub fn new(name: String, content_type: String, location: ArtifactLocation) -> Self {
        Artifact {
            id: ArtifactId::new(),
            name,
            content_type,
            size: 0,
            location,
            checksum: None,
            created_at: SystemTime::now(),
            source_run: None,
            source_step: None,
        }
    }

    /// Set artifact size
    pub fn with_size(mut self, size: u64) -> Self {
        self.size = size;
        self
    }

    /// Set artifact checksum
    pub fn with_checksum(mut self, checksum: Checksum) -> Self {
        self.checksum = Some(checksum);
        self
    }

    /// Set source run
    pub fn from_run(mut self, run_id: RunId) -> Self {
        self.source_run = Some(run_id);
        self
    }

    /// Set source step
    pub fn from_step(mut self, step_id: StepId) -> Self {
        self.source_step = Some(step_id);
        self
    }
}

// ============================================================================
// Execution Handle
// ============================================================================

/// A handle to an active execution
#[derive(Debug, Clone)]
pub struct ExecutionHandle {
    /// Run ID for this execution
    pub run_id: RunId,
    /// Runtime environment
    pub runtime_kind: RuntimeKind,
    /// When handle was created
    pub created_at: Instant,
    /// Runtime-specific handle string
    pub runtime_handle: String,
}

impl ExecutionHandle {
    /// Create a new execution handle
    pub fn new(run_id: RunId, runtime_kind: RuntimeKind, runtime_handle: String) -> Self {
        ExecutionHandle {
            run_id,
            runtime_kind,
            created_at: Instant::now(),
            runtime_handle,
        }
    }

    /// Get elapsed time since handle creation
    pub fn elapsed(&self) -> Duration {
        self.created_at.elapsed()
    }
}

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

    #[test]
    fn test_run_input_default_none() {
        let run = Run::new(
            RunTarget::Agent {
                spec_path: PathBuf::from("agent.yaml"),
            },
            RuntimeKind::Local,
        );
        assert!(run.input.is_none());
    }

    #[test]
    fn test_run_with_input() {
        use serde_json::json;
        let value = json!({"key": "value"});
        let run = Run::new(
            RunTarget::Agent {
                spec_path: PathBuf::from("agent.yaml"),
            },
            RuntimeKind::Local,
        )
        .with_input(value.clone());
        assert_eq!(run.input, Some(value));
    }

    #[test]
    fn test_run_id_generation() {
        let id1 = RunId::new();
        let id2 = RunId::new();
        assert_ne!(id1, id2);
    }

    #[test]
    fn test_run_lifecycle() {
        let mut run = Run::new(
            RunTarget::Agent {
                spec_path: PathBuf::from("agent.yaml"),
            },
            RuntimeKind::Local,
        );
        assert_eq!(run.status, RunStatus::Pending);
        assert!(run.started_at.is_none());

        run.start();
        assert_eq!(run.status, RunStatus::Running);

        run.complete();
        assert_eq!(run.status, RunStatus::Completed);
        assert!(run.duration().is_some());
    }

    #[test]
    fn test_run_failure() {
        let mut run = Run::new(
            RunTarget::Agent {
                spec_path: PathBuf::from("agent.yaml"),
            },
            RuntimeKind::Local,
        );
        run.start();
        run.fail(RunError::ExecutionFailed {
            message: "test".into(),
        });
        assert_eq!(run.status, RunStatus::Failed);
    }

    #[test]
    fn test_step_lifecycle() {
        let run_id = RunId::new();
        let mut step = Step::new(run_id, "test-step".into(), PathBuf::from("agent.yaml"));
        assert_eq!(step.status, StepStatus::Pending);

        step.start();
        assert_eq!(step.status, StepStatus::Running);

        step.complete(vec![ArtifactId::new()]);
        assert_eq!(step.status, StepStatus::Completed);
        assert!(!step.output_artifacts.is_empty());
    }

    #[test]
    fn test_session() {
        let mut session = Session::new("test".into());
        session.add_run(RunId::new());
        assert_eq!(session.run_ids.len(), 1);
        session.end();
        assert_eq!(session.status, SessionStatus::Ended);
    }

    #[test]
    fn test_artifact() {
        let artifact = Artifact::new(
            "output.json".into(),
            "application/json".into(),
            ArtifactLocation::Local {
                path: PathBuf::from("/tmp/out.json"),
            },
        )
        .with_size(1024);
        assert_eq!(artifact.name, "output.json");
        assert_eq!(artifact.size, 1024);
    }
}