durable_execution_sdk/

context.rs

//! DurableContext and operation identifier types for the AWS Durable Execution SDK.
//!
//! This module provides the main `DurableContext` interface that user code interacts with,
//! as well as the `OperationIdentifier` type for deterministic operation ID generation.

use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, RwLock};

use blake2::{Blake2b512, Digest};

use crate::error::DurableResult;
use crate::sealed::Sealed;
use crate::state::ExecutionState;
use crate::traits::{DurableValue, StepFn};
use crate::types::OperationId;

/// Identifies an operation within a durable execution.
///
/// Each operation has a unique ID that is deterministically generated
/// based on the parent context and step counter. This ensures that
/// the same operation gets the same ID across replays.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct OperationIdentifier {
    /// The unique operation ID (deterministically generated)
    pub operation_id: String,
    /// The parent operation ID (None for root operations)
    pub parent_id: Option<String>,
    /// Optional human-readable name for the operation
    pub name: Option<String>,
}

impl OperationIdentifier {
    /// Creates a new OperationIdentifier with the given values.
    pub fn new(
        operation_id: impl Into<String>,
        parent_id: Option<String>,
        name: Option<String>,
    ) -> Self {
        Self {
            operation_id: operation_id.into(),
            parent_id,
            name,
        }
    }

    /// Creates a root operation identifier (no parent).
    pub fn root(operation_id: impl Into<String>) -> Self {
        Self {
            operation_id: operation_id.into(),
            parent_id: None,
            name: None,
        }
    }

    /// Creates an operation identifier with a parent.
    pub fn with_parent(operation_id: impl Into<String>, parent_id: impl Into<String>) -> Self {
        Self {
            operation_id: operation_id.into(),
            parent_id: Some(parent_id.into()),
            name: None,
        }
    }

    /// Sets the name for this operation identifier.
    pub fn with_name(mut self, name: impl Into<String>) -> Self {
        self.name = Some(name.into());
        self
    }

    /// Returns the operation ID as an `OperationId` newtype.
    #[inline]
    pub fn operation_id_typed(&self) -> OperationId {
        OperationId::from(self.operation_id.clone())
    }

    /// Returns the parent ID as an `OperationId` newtype, if present.
    #[inline]
    pub fn parent_id_typed(&self) -> Option<OperationId> {
        self.parent_id
            .as_ref()
            .map(|id| OperationId::from(id.clone()))
    }

    /// Applies this identifier's parent_id and name to an `OperationUpdate`.
    ///
    /// This is a convenience method to avoid repeated boilerplate when building
    /// operation updates across handlers.
    pub fn apply_to(
        &self,
        mut update: crate::operation::OperationUpdate,
    ) -> crate::operation::OperationUpdate {
        if let Some(ref parent_id) = self.parent_id {
            update = update.with_parent_id(parent_id);
        }
        if let Some(ref name) = self.name {
            update = update.with_name(name);
        }
        update
    }
}

impl std::fmt::Display for OperationIdentifier {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if let Some(ref name) = self.name {
            write!(f, "{}({})", name, self.operation_id)
        } else {
            write!(f, "{}", self.operation_id)
        }
    }
}

/// Generates deterministic operation IDs using blake2b hashing.
///
/// The ID is generated by hashing the parent ID (or execution ARN for root)
/// combined with the step counter value. This ensures:
/// - Same inputs always produce the same ID
/// - Different step counts produce different IDs
/// - IDs are unique within an execution
#[derive(Debug)]
pub struct OperationIdGenerator {
    /// The base identifier (execution ARN or parent operation ID)
    base_id: String,
    /// Thread-safe step counter
    step_counter: AtomicU64,
}

impl OperationIdGenerator {
    /// Creates a new OperationIdGenerator with the given base ID.
    ///
    /// # Arguments
    ///
    /// * `base_id` - The base identifier to use for hashing (typically execution ARN or parent ID)
    pub fn new(base_id: impl Into<String>) -> Self {
        Self {
            base_id: base_id.into(),
            step_counter: AtomicU64::new(0),
        }
    }

    /// Creates a new OperationIdGenerator with a specific starting counter value.
    ///
    /// # Arguments
    ///
    /// * `base_id` - The base identifier to use for hashing
    /// * `initial_counter` - The initial value for the step counter
    pub fn with_counter(base_id: impl Into<String>, initial_counter: u64) -> Self {
        Self {
            base_id: base_id.into(),
            step_counter: AtomicU64::new(initial_counter),
        }
    }

    /// Generates the next operation ID.
    ///
    /// This method atomically increments the step counter and generates
    /// a deterministic ID based on the base ID and counter value.
    ///
    /// # Returns
    ///
    /// A unique, deterministic operation ID string.
    ///
    /// # Memory Ordering
    ///
    /// Uses `Ordering::Relaxed` because the step counter only needs to provide
    /// unique values - there's no synchronization requirement with other data.
    /// Each call gets a unique counter value, and the hash function ensures
    /// deterministic ID generation regardless of ordering between threads.
    ///
    /// Requirements: 4.1, 4.6
    pub fn next_id(&self) -> String {
        // Relaxed ordering is sufficient: we only need uniqueness, not synchronization.
        // The counter value is combined with base_id in a hash, so the only requirement
        // is that each call gets a different counter value, which fetch_add guarantees.
        let counter = self.step_counter.fetch_add(1, Ordering::Relaxed);
        generate_operation_id(&self.base_id, counter)
    }

    /// Generates an operation ID for a specific counter value without incrementing.
    ///
    /// This is useful for testing or when you need to generate an ID
    /// for a known counter value.
    ///
    /// # Arguments
    ///
    /// * `counter` - The counter value to use for ID generation
    ///
    /// # Returns
    ///
    /// A deterministic operation ID string.
    pub fn id_for_counter(&self, counter: u64) -> String {
        generate_operation_id(&self.base_id, counter)
    }

    /// Returns the current counter value without incrementing.
    ///
    /// # Memory Ordering
    ///
    /// Uses `Ordering::Relaxed` because this is an informational read that doesn't
    /// need to synchronize with other operations. The value may be slightly stale
    /// in concurrent scenarios, but this is acceptable for debugging/monitoring use.
    ///
    /// Requirements: 4.1, 4.6
    pub fn current_counter(&self) -> u64 {
        // Relaxed ordering is sufficient: this is an informational read with no
        // synchronization requirements. Callers should not rely on this value
        // for correctness in concurrent scenarios.
        self.step_counter.load(Ordering::Relaxed)
    }

    /// Returns the base ID used for generation.
    pub fn base_id(&self) -> &str {
        &self.base_id
    }

    /// Creates a child generator with a new base ID.
    ///
    /// The child generator starts with counter 0 and uses the provided
    /// parent operation ID as its base.
    ///
    /// # Arguments
    ///
    /// * `parent_operation_id` - The operation ID to use as the base for the child
    pub fn create_child(&self, parent_operation_id: impl Into<String>) -> Self {
        Self::new(parent_operation_id)
    }
}

impl Clone for OperationIdGenerator {
    fn clone(&self) -> Self {
        Self {
            base_id: self.base_id.clone(),
            // Relaxed ordering is sufficient: we're creating a snapshot of the counter
            // for a new generator. The clone doesn't need to synchronize with the
            // original - it just needs a reasonable starting point.
            step_counter: AtomicU64::new(self.step_counter.load(Ordering::Relaxed)),
        }
    }
}

/// Generates a deterministic operation ID using blake2b hashing.
///
/// # Arguments
///
/// * `base_id` - The base identifier (execution ARN or parent operation ID)
/// * `counter` - The step counter value
///
/// # Returns
///
/// A hex-encoded operation ID string (first 32 characters of the hash).
pub fn generate_operation_id(base_id: &str, counter: u64) -> String {
    let mut hasher = Blake2b512::new();
    hasher.update(base_id.as_bytes());
    hasher.update(counter.to_le_bytes());
    let result = hasher.finalize();

    // Take first 16 bytes (32 hex chars) for a reasonably short but unique ID
    hex::encode(&result[..16])
}

/// Logger trait for structured logging in durable executions.
///
/// This trait is compatible with the `tracing` crate and allows
/// custom logging implementations.
///
/// # Sealed Trait
///
/// This trait is sealed and cannot be implemented outside of this crate.
/// This allows the SDK maintainers to evolve the logging interface without
/// breaking external code. If you need custom logging behavior, use the
/// provided factory functions or wrap one of the existing implementations.
#[allow(private_bounds)]
pub trait Logger: Sealed + Send + Sync {
    /// Logs a debug message.
    fn debug(&self, message: &str, info: &LogInfo);
    /// Logs an info message.
    fn info(&self, message: &str, info: &LogInfo);
    /// Logs a warning message.
    fn warn(&self, message: &str, info: &LogInfo);
    /// Logs an error message.
    fn error(&self, message: &str, info: &LogInfo);
}

/// Context information for log messages.
///
/// This struct provides context for log messages including execution ARN,
/// operation IDs, and replay status. The `is_replay` flag indicates whether
/// the current operation is being replayed from a checkpoint.
#[derive(Debug, Clone, Default)]
pub struct LogInfo {
    /// The durable execution ARN
    pub durable_execution_arn: Option<String>,
    /// The current operation ID
    pub operation_id: Option<String>,
    /// The parent operation ID
    pub parent_id: Option<String>,
    /// Whether the current operation is being replayed from a checkpoint
    ///
    /// When `true`, the operation is returning a previously checkpointed result
    /// without re-executing the operation's logic. Loggers can use this flag
    /// to suppress or annotate logs during replay.
    pub is_replay: bool,
    /// Additional key-value pairs
    pub extra: Vec<(String, String)>,
}

impl LogInfo {
    /// Creates a new LogInfo with the given execution ARN.
    pub fn new(durable_execution_arn: impl Into<String>) -> Self {
        Self {
            durable_execution_arn: Some(durable_execution_arn.into()),
            ..Default::default()
        }
    }

    /// Sets the operation ID.
    pub fn with_operation_id(mut self, operation_id: impl Into<String>) -> Self {
        self.operation_id = Some(operation_id.into());
        self
    }

    /// Sets the parent ID.
    pub fn with_parent_id(mut self, parent_id: impl Into<String>) -> Self {
        self.parent_id = Some(parent_id.into());
        self
    }

    /// Sets the replay flag.
    ///
    /// When set to `true`, indicates that the current operation is being
    /// replayed from a checkpoint rather than executing fresh.
    ///
    /// # Arguments
    ///
    /// * `is_replay` - Whether the operation is being replayed
    pub fn with_replay(mut self, is_replay: bool) -> Self {
        self.is_replay = is_replay;
        self
    }

    /// Adds an extra key-value pair.
    pub fn with_extra(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
        self.extra.push((key.into(), value.into()));
        self
    }
}

/// Creates a tracing span for a durable operation.
///
/// This function creates a structured tracing span that includes all relevant
/// context for observability and debugging. The span can be used to trace
/// execution flow and correlate logs across operations.
///
/// # Arguments
///
/// * `operation_type` - The type of operation (e.g., "step", "wait", "callback", "invoke", "map", "parallel")
/// * `op_id` - The operation identifier containing operation_id, parent_id, and name
/// * `durable_execution_arn` - The ARN of the durable execution
///
/// # Returns
///
/// A `tracing::Span` that can be entered during operation execution.
///
/// # Example
///
/// ```rust,ignore
/// let span = create_operation_span("step", &op_id, state.durable_execution_arn());
/// let _guard = span.enter();
/// // ... operation execution ...
/// ```
pub fn create_operation_span(
    operation_type: &str,
    op_id: &OperationIdentifier,
    durable_execution_arn: &str,
) -> tracing::Span {
    tracing::info_span!(
        "durable_operation",
        operation_type = %operation_type,
        operation_id = %op_id.operation_id,
        parent_id = ?op_id.parent_id,
        name = ?op_id.name,
        durable_execution_arn = %durable_execution_arn,
        status = tracing::field::Empty,
    )
}

/// Default logger implementation using the `tracing` crate.
///
/// This logger includes the `is_replay` flag in all log messages to help
/// distinguish between fresh executions and replayed operations.
///
/// Extra fields from `LogInfo` are included in the tracing output as a formatted
/// string of key-value pairs, making them queryable in log aggregation systems.
#[derive(Debug, Clone, Default)]
pub struct TracingLogger;

// Implement Sealed for TracingLogger to allow it to implement Logger
impl Sealed for TracingLogger {}

impl TracingLogger {
    /// Formats extra fields as a string of key-value pairs.
    ///
    /// Returns an empty string if there are no extra fields, otherwise returns
    /// a comma-separated list of "key=value" pairs.
    fn format_extra_fields(extra: &[(String, String)]) -> String {
        if extra.is_empty() {
            String::new()
        } else {
            extra
                .iter()
                .map(|(k, v)| format!("{}={}", k, v))
                .collect::<Vec<_>>()
                .join(", ")
        }
    }
}

impl Logger for TracingLogger {
    fn debug(&self, message: &str, info: &LogInfo) {
        let extra_fields = Self::format_extra_fields(&info.extra);
        tracing::debug!(
            durable_execution_arn = ?info.durable_execution_arn,
            operation_id = ?info.operation_id,
            parent_id = ?info.parent_id,
            is_replay = info.is_replay,
            extra = %extra_fields,
            "{}",
            message
        );
    }

    fn info(&self, message: &str, info: &LogInfo) {
        let extra_fields = Self::format_extra_fields(&info.extra);
        tracing::info!(
            durable_execution_arn = ?info.durable_execution_arn,
            operation_id = ?info.operation_id,
            parent_id = ?info.parent_id,
            is_replay = info.is_replay,
            extra = %extra_fields,
            "{}",
            message
        );
    }

    fn warn(&self, message: &str, info: &LogInfo) {
        let extra_fields = Self::format_extra_fields(&info.extra);
        tracing::warn!(
            durable_execution_arn = ?info.durable_execution_arn,
            operation_id = ?info.operation_id,
            parent_id = ?info.parent_id,
            is_replay = info.is_replay,
            extra = %extra_fields,
            "{}",
            message
        );
    }

    fn error(&self, message: &str, info: &LogInfo) {
        let extra_fields = Self::format_extra_fields(&info.extra);
        tracing::error!(
            durable_execution_arn = ?info.durable_execution_arn,
            operation_id = ?info.operation_id,
            parent_id = ?info.parent_id,
            is_replay = info.is_replay,
            extra = %extra_fields,
            "{}",
            message
        );
    }
}

/// Configuration for replay-aware logging behavior.
///
/// This configuration controls how the `ReplayAwareLogger` handles log messages
/// during replay. Users can choose to suppress all logs during replay, allow
/// only certain log levels, or log all messages regardless of replay status.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum ReplayLoggingConfig {
    /// Suppress all logs during replay (default).
    ///
    /// When in replay mode, no log messages will be emitted. This is useful
    /// to reduce noise in logs when replaying previously executed operations.
    #[default]
    SuppressAll,

    /// Allow all logs during replay.
    ///
    /// Log messages will be emitted regardless of replay status. The `is_replay`
    /// flag in `LogInfo` can still be used to distinguish replay logs.
    AllowAll,

    /// Allow only error logs during replay.
    ///
    /// Only error-level messages will be emitted during replay. This is useful
    /// when you want to see errors but suppress informational messages.
    ErrorsOnly,

    /// Allow error and warning logs during replay.
    ///
    /// Error and warning-level messages will be emitted during replay.
    /// Debug and info messages will be suppressed.
    WarningsAndErrors,
}

/// A logger wrapper that can suppress logs during replay based on configuration.
///
/// `ReplayAwareLogger` wraps any `Logger` implementation and adds replay-aware
/// behavior. When the `is_replay` flag in `LogInfo` is `true`, the logger will
/// suppress or allow logs based on the configured `ReplayLoggingConfig`.
///
/// # Example
///
/// ```rust
/// use durable_execution_sdk::{TracingLogger, ReplayAwareLogger, ReplayLoggingConfig};
/// use std::sync::Arc;
///
/// // Create a replay-aware logger that suppresses all logs during replay
/// let logger = ReplayAwareLogger::new(
///     Arc::new(TracingLogger),
///     ReplayLoggingConfig::SuppressAll,
/// );
///
/// // Create a replay-aware logger that allows errors during replay
/// let logger_with_errors = ReplayAwareLogger::new(
///     Arc::new(TracingLogger),
///     ReplayLoggingConfig::ErrorsOnly,
/// );
/// ```
pub struct ReplayAwareLogger {
    /// The underlying logger to delegate to
    inner: Arc<dyn Logger>,
    /// Configuration for replay logging behavior
    config: ReplayLoggingConfig,
}

// Implement Sealed for ReplayAwareLogger to allow it to implement Logger
impl Sealed for ReplayAwareLogger {}

impl ReplayAwareLogger {
    /// Creates a new `ReplayAwareLogger` with the specified configuration.
    ///
    /// # Arguments
    ///
    /// * `inner` - The underlying logger to delegate to
    /// * `config` - Configuration for replay logging behavior
    ///
    /// # Example
    ///
    /// ```rust
    /// use durable_execution_sdk::{TracingLogger, ReplayAwareLogger, ReplayLoggingConfig};
    /// use std::sync::Arc;
    ///
    /// let logger = ReplayAwareLogger::new(
    ///     Arc::new(TracingLogger),
    ///     ReplayLoggingConfig::SuppressAll,
    /// );
    /// ```
    pub fn new(inner: Arc<dyn Logger>, config: ReplayLoggingConfig) -> Self {
        Self { inner, config }
    }

    /// Creates a new `ReplayAwareLogger` that suppresses all logs during replay.
    ///
    /// This is a convenience constructor for the most common use case.
    ///
    /// # Arguments
    ///
    /// * `inner` - The underlying logger to delegate to
    pub fn suppress_replay(inner: Arc<dyn Logger>) -> Self {
        Self::new(inner, ReplayLoggingConfig::SuppressAll)
    }

    /// Creates a new `ReplayAwareLogger` that allows all logs during replay.
    ///
    /// # Arguments
    ///
    /// * `inner` - The underlying logger to delegate to
    pub fn allow_all(inner: Arc<dyn Logger>) -> Self {
        Self::new(inner, ReplayLoggingConfig::AllowAll)
    }

    /// Returns the current replay logging configuration.
    pub fn config(&self) -> ReplayLoggingConfig {
        self.config
    }

    /// Returns a reference to the underlying logger.
    pub fn inner(&self) -> &Arc<dyn Logger> {
        &self.inner
    }

    /// Checks if a log at the given level should be suppressed during replay.
    fn should_suppress(&self, info: &LogInfo, level: LogLevel) -> bool {
        if !info.is_replay {
            return false;
        }

        match self.config {
            ReplayLoggingConfig::SuppressAll => true,
            ReplayLoggingConfig::AllowAll => false,
            ReplayLoggingConfig::ErrorsOnly => level != LogLevel::Error,
            ReplayLoggingConfig::WarningsAndErrors => {
                level != LogLevel::Error && level != LogLevel::Warn
            }
        }
    }
}

/// Internal enum for log levels used by ReplayAwareLogger.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum LogLevel {
    Debug,
    Info,
    Warn,
    Error,
}

impl std::fmt::Debug for ReplayAwareLogger {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ReplayAwareLogger")
            .field("config", &self.config)
            .finish()
    }
}

impl Logger for ReplayAwareLogger {
    fn debug(&self, message: &str, info: &LogInfo) {
        if !self.should_suppress(info, LogLevel::Debug) {
            self.inner.debug(message, info);
        }
    }

    fn info(&self, message: &str, info: &LogInfo) {
        if !self.should_suppress(info, LogLevel::Info) {
            self.inner.info(message, info);
        }
    }

    fn warn(&self, message: &str, info: &LogInfo) {
        if !self.should_suppress(info, LogLevel::Warn) {
            self.inner.warn(message, info);
        }
    }

    fn error(&self, message: &str, info: &LogInfo) {
        if !self.should_suppress(info, LogLevel::Error) {
            self.inner.error(message, info);
        }
    }
}

/// A custom logger that delegates to user-provided closures.
///
/// This struct allows users to provide custom logging behavior without
/// implementing the sealed `Logger` trait directly. Each log level can
/// have its own closure for handling log messages.
///
/// # Example
///
/// ```rust
/// use durable_execution_sdk::context::{custom_logger, LogInfo};
/// use std::sync::Arc;
///
/// // Create a custom logger that prints to stdout
/// let logger = custom_logger(
///     |msg, info| println!("[DEBUG] {}: {:?}", msg, info),
///     |msg, info| println!("[INFO] {}: {:?}", msg, info),
///     |msg, info| println!("[WARN] {}: {:?}", msg, info),
///     |msg, info| println!("[ERROR] {}: {:?}", msg, info),
/// );
/// ```
pub struct CustomLogger<D, I, W, E>
where
    D: Fn(&str, &LogInfo) + Send + Sync,
    I: Fn(&str, &LogInfo) + Send + Sync,
    W: Fn(&str, &LogInfo) + Send + Sync,
    E: Fn(&str, &LogInfo) + Send + Sync,
{
    debug_fn: D,
    info_fn: I,
    warn_fn: W,
    error_fn: E,
}

// Implement Sealed for CustomLogger to allow it to implement Logger
impl<D, I, W, E> Sealed for CustomLogger<D, I, W, E>
where
    D: Fn(&str, &LogInfo) + Send + Sync,
    I: Fn(&str, &LogInfo) + Send + Sync,
    W: Fn(&str, &LogInfo) + Send + Sync,
    E: Fn(&str, &LogInfo) + Send + Sync,
{
}

impl<D, I, W, E> Logger for CustomLogger<D, I, W, E>
where
    D: Fn(&str, &LogInfo) + Send + Sync,
    I: Fn(&str, &LogInfo) + Send + Sync,
    W: Fn(&str, &LogInfo) + Send + Sync,
    E: Fn(&str, &LogInfo) + Send + Sync,
{
    fn debug(&self, message: &str, info: &LogInfo) {
        (self.debug_fn)(message, info);
    }

    fn info(&self, message: &str, info: &LogInfo) {
        (self.info_fn)(message, info);
    }

    fn warn(&self, message: &str, info: &LogInfo) {
        (self.warn_fn)(message, info);
    }

    fn error(&self, message: &str, info: &LogInfo) {
        (self.error_fn)(message, info);
    }
}

impl<D, I, W, E> std::fmt::Debug for CustomLogger<D, I, W, E>
where
    D: Fn(&str, &LogInfo) + Send + Sync,
    I: Fn(&str, &LogInfo) + Send + Sync,
    W: Fn(&str, &LogInfo) + Send + Sync,
    E: Fn(&str, &LogInfo) + Send + Sync,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CustomLogger").finish()
    }
}

/// Creates a custom logger with user-provided closures for each log level.
///
/// This factory function allows users to create custom logging behavior without
/// implementing the sealed `Logger` trait directly. Each log level has its own
/// closure that receives the log message and context information.
///
/// # Arguments
///
/// * `debug_fn` - Closure called for debug-level messages
/// * `info_fn` - Closure called for info-level messages
/// * `warn_fn` - Closure called for warning-level messages
/// * `error_fn` - Closure called for error-level messages
///
/// # Example
///
/// ```rust
/// use durable_execution_sdk::context::{custom_logger, LogInfo};
/// use std::sync::Arc;
///
/// // Create a custom logger that prints to stdout
/// let logger = custom_logger(
///     |msg, info| println!("[DEBUG] {}: {:?}", msg, info),
///     |msg, info| println!("[INFO] {}: {:?}", msg, info),
///     |msg, info| println!("[WARN] {}: {:?}", msg, info),
///     |msg, info| println!("[ERROR] {}: {:?}", msg, info),
/// );
///
/// // Use with Arc for sharing across contexts
/// let shared_logger = Arc::new(logger);
/// ```
pub fn custom_logger<D, I, W, E>(
    debug_fn: D,
    info_fn: I,
    warn_fn: W,
    error_fn: E,
) -> CustomLogger<D, I, W, E>
where
    D: Fn(&str, &LogInfo) + Send + Sync,
    I: Fn(&str, &LogInfo) + Send + Sync,
    W: Fn(&str, &LogInfo) + Send + Sync,
    E: Fn(&str, &LogInfo) + Send + Sync,
{
    CustomLogger {
        debug_fn,
        info_fn,
        warn_fn,
        error_fn,
    }
}

/// Creates a simple custom logger that uses a single closure for all log levels.
///
/// This is a convenience function for cases where you want the same handling
/// for all log levels. The closure receives the log level as a string, the
/// message, and the log info.
///
/// # Arguments
///
/// * `log_fn` - Closure called for all log messages. Receives (level, message, info).
///
/// # Example
///
/// ```rust
/// use durable_execution_sdk::context::{simple_custom_logger, LogInfo};
/// use std::sync::Arc;
///
/// // Create a simple logger that prints all messages with their level
/// let logger = simple_custom_logger(|level, msg, info| {
///     println!("[{}] {}: {:?}", level, msg, info);
/// });
///
/// // Use with Arc for sharing across contexts
/// let shared_logger = Arc::new(logger);
/// ```
pub fn simple_custom_logger<F>(log_fn: F) -> impl Logger
where
    F: Fn(&str, &str, &LogInfo) + Send + Sync + Clone + 'static,
{
    let debug_fn = log_fn.clone();
    let info_fn = log_fn.clone();
    let warn_fn = log_fn.clone();
    let error_fn = log_fn;

    custom_logger(
        move |msg, info| debug_fn("DEBUG", msg, info),
        move |msg, info| info_fn("INFO", msg, info),
        move |msg, info| warn_fn("WARN", msg, info),
        move |msg, info| error_fn("ERROR", msg, info),
    )
}

/// The main context for durable execution operations.
///
/// `DurableContext` provides all the durable operations that user code
/// can use to build reliable workflows. It handles checkpointing, replay,
/// and state management automatically.
///
/// # Thread Safety
///
/// `DurableContext` is `Send + Sync` and can be safely shared across
/// async tasks. The internal state uses appropriate synchronization
/// primitives for concurrent access.
///
/// # Example
///
/// ```rust,ignore
/// async fn my_workflow(ctx: DurableContext) -> Result<String, DurableError> {
///     // Execute a step (automatically checkpointed)
///     let result = ctx.step(|_| Ok("hello".to_string()), None).await?;
///     
///     // Wait for 5 seconds (suspends Lambda, resumes later)
///     ctx.wait(Duration::from_seconds(5), None).await?;
///     
///     Ok(result)
/// }
/// ```
pub struct DurableContext {
    /// The execution state manager
    state: Arc<ExecutionState>,
    /// The Lambda context (if running in Lambda)
    lambda_context: Option<lambda_runtime::Context>,
    /// The parent operation ID (None for root context)
    parent_id: Option<String>,
    /// Operation ID generator for this context
    id_generator: Arc<OperationIdGenerator>,
    /// Logger for structured logging (wrapped in RwLock for runtime reconfiguration)
    logger: Arc<RwLock<Arc<dyn Logger>>>,
}

// Ensure DurableContext is Send + Sync
static_assertions::assert_impl_all!(DurableContext: Send, Sync);

impl DurableContext {
    /// Creates a new DurableContext with the given state.
    ///
    /// # Arguments
    ///
    /// * `state` - The execution state manager
    pub fn new(state: Arc<ExecutionState>) -> Self {
        let base_id = state.durable_execution_arn().to_string();
        Self {
            state,
            lambda_context: None,
            parent_id: None,
            id_generator: Arc::new(OperationIdGenerator::new(base_id)),
            logger: Arc::new(RwLock::new(Arc::new(TracingLogger))),
        }
    }

    /// Creates a DurableContext from a Lambda context.
    ///
    /// This is the primary factory method used when running in AWS Lambda.
    ///
    /// # Arguments
    ///
    /// * `state` - The execution state manager
    /// * `lambda_context` - The Lambda runtime context
    pub fn from_lambda_context(
        state: Arc<ExecutionState>,
        lambda_context: lambda_runtime::Context,
    ) -> Self {
        let base_id = state.durable_execution_arn().to_string();
        Self {
            state,
            lambda_context: Some(lambda_context),
            parent_id: None,
            id_generator: Arc::new(OperationIdGenerator::new(base_id)),
            logger: Arc::new(RwLock::new(Arc::new(TracingLogger))),
        }
    }

    /// Creates a child context for nested operations.
    ///
    /// Child contexts have their own step counter but share the same
    /// execution state. Operations in a child context are tracked
    /// with the parent's operation ID.
    ///
    /// # Arguments
    ///
    /// * `parent_operation_id` - The operation ID of the parent operation
    pub fn create_child_context(&self, parent_operation_id: impl Into<String>) -> Self {
        let parent_id = parent_operation_id.into();
        Self {
            state: self.state.clone(),
            lambda_context: self.lambda_context.clone(),
            parent_id: Some(parent_id.clone()),
            id_generator: Arc::new(OperationIdGenerator::new(parent_id)),
            logger: self.logger.clone(),
        }
    }

    /// Sets a custom logger for this context.
    ///
    /// # Arguments
    ///
    /// * `logger` - The logger implementation to use
    pub fn set_logger(&mut self, logger: Arc<dyn Logger>) {
        *self.logger.write().unwrap() = logger;
    }

    /// Returns a new context with the specified logger.
    ///
    /// # Arguments
    ///
    /// * `logger` - The logger implementation to use
    pub fn with_logger(self, logger: Arc<dyn Logger>) -> Self {
        *self.logger.write().unwrap() = logger;
        self
    }

    /// Reconfigures the logger for this context at runtime.
    ///
    /// All subsequent log calls on this context (and any clones sharing the
    /// same underlying `RwLock`) will use the new logger.
    ///
    /// # Arguments
    ///
    /// * `logger` - The new logger implementation to use
    pub fn configure_logger(&self, logger: Arc<dyn Logger>) {
        *self.logger.write().unwrap() = logger;
    }

    /// Returns a reference to the execution state.
    pub fn state(&self) -> &Arc<ExecutionState> {
        &self.state
    }

    /// Returns the durable execution ARN.
    pub fn durable_execution_arn(&self) -> &str {
        self.state.durable_execution_arn()
    }

    /// Returns the parent operation ID, if any.
    pub fn parent_id(&self) -> Option<&str> {
        self.parent_id.as_deref()
    }

    /// Returns a reference to the Lambda context, if available.
    pub fn lambda_context(&self) -> Option<&lambda_runtime::Context> {
        self.lambda_context.as_ref()
    }

    /// Returns a reference to the logger.
    pub fn logger(&self) -> Arc<dyn Logger> {
        self.logger.read().unwrap().clone()
    }

    /// Generates the next operation ID for this context.
    ///
    /// This method is thread-safe and will generate unique, deterministic
    /// IDs based on the context's base ID and step counter.
    pub fn next_operation_id(&self) -> String {
        self.id_generator.next_id()
    }

    /// Creates an OperationIdentifier for the next operation.
    ///
    /// # Arguments
    ///
    /// * `name` - Optional human-readable name for the operation
    pub fn next_operation_identifier(&self, name: Option<String>) -> OperationIdentifier {
        OperationIdentifier::new(self.next_operation_id(), self.parent_id.clone(), name)
    }

    /// Returns the current step counter value without incrementing.
    pub fn current_step_counter(&self) -> u64 {
        self.id_generator.current_counter()
    }

    /// Creates log info for the current context.
    ///
    /// The returned `LogInfo` includes the current replay status from the
    /// execution state, allowing loggers to distinguish between fresh
    /// executions and replayed operations.
    pub fn create_log_info(&self) -> LogInfo {
        let mut info = LogInfo::new(self.durable_execution_arn());
        if let Some(ref parent_id) = self.parent_id {
            info = info.with_parent_id(parent_id);
        }
        // Include replay status from execution state
        info = info.with_replay(self.state.is_replay());
        info
    }

    /// Creates log info with an operation ID.
    ///
    /// The returned `LogInfo` includes the current replay status from the
    /// execution state, allowing loggers to distinguish between fresh
    /// executions and replayed operations.
    pub fn create_log_info_with_operation(&self, operation_id: &str) -> LogInfo {
        self.create_log_info().with_operation_id(operation_id)
    }

    /// Creates log info with explicit replay status.
    ///
    /// This method allows callers to explicitly set the replay status,
    /// which is useful when the operation-specific replay status differs
    /// from the global execution state replay status.
    ///
    /// # Arguments
    ///
    /// * `operation_id` - The operation ID to include in the log info
    /// * `is_replay` - Whether this specific operation is being replayed
    pub fn create_log_info_with_replay(&self, operation_id: &str, is_replay: bool) -> LogInfo {
        let mut info = LogInfo::new(self.durable_execution_arn());
        if let Some(ref parent_id) = self.parent_id {
            info = info.with_parent_id(parent_id);
        }
        info.with_operation_id(operation_id).with_replay(is_replay)
    }

    // ========================================================================
    // Simplified Logging API
    // ========================================================================

    /// Logs a message at INFO level with automatic context.
    ///
    /// This method automatically includes the durable_execution_arn and parent_id
    /// in the log output without requiring the caller to specify them.
    ///
    /// # Arguments
    ///
    /// * `message` - The message to log
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// ctx.log_info("Processing order started");
    /// ```
    pub fn log_info(&self, message: &str) {
        self.log_with_level(LogLevel::Info, message, &[]);
    }

    /// Logs a message at INFO level with extra fields.
    ///
    /// This method automatically includes the durable_execution_arn and parent_id
    /// in the log output, plus any additional fields specified.
    ///
    /// # Arguments
    ///
    /// * `message` - The message to log
    /// * `fields` - Additional key-value pairs to include in the log
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// ctx.log_info_with("Processing order", &[("order_id", "123"), ("amount", "99.99")]);
    /// ```
    pub fn log_info_with(&self, message: &str, fields: &[(&str, &str)]) {
        self.log_with_level(LogLevel::Info, message, fields);
    }

    /// Logs a message at DEBUG level with automatic context.
    ///
    /// This method automatically includes the durable_execution_arn and parent_id
    /// in the log output without requiring the caller to specify them.
    ///
    /// # Arguments
    ///
    /// * `message` - The message to log
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// ctx.log_debug("Entering validation step");
    /// ```
    pub fn log_debug(&self, message: &str) {
        self.log_with_level(LogLevel::Debug, message, &[]);
    }

    /// Logs a message at DEBUG level with extra fields.
    ///
    /// This method automatically includes the durable_execution_arn and parent_id
    /// in the log output, plus any additional fields specified.
    ///
    /// # Arguments
    ///
    /// * `message` - The message to log
    /// * `fields` - Additional key-value pairs to include in the log
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// ctx.log_debug_with("Variable state", &[("x", "42"), ("y", "100")]);
    /// ```
    pub fn log_debug_with(&self, message: &str, fields: &[(&str, &str)]) {
        self.log_with_level(LogLevel::Debug, message, fields);
    }

    /// Logs a message at WARN level with automatic context.
    ///
    /// This method automatically includes the durable_execution_arn and parent_id
    /// in the log output without requiring the caller to specify them.
    ///
    /// # Arguments
    ///
    /// * `message` - The message to log
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// ctx.log_warn("Retry attempt 3 of 5");
    /// ```
    pub fn log_warn(&self, message: &str) {
        self.log_with_level(LogLevel::Warn, message, &[]);
    }

    /// Logs a message at WARN level with extra fields.
    ///
    /// This method automatically includes the durable_execution_arn and parent_id
    /// in the log output, plus any additional fields specified.
    ///
    /// # Arguments
    ///
    /// * `message` - The message to log
    /// * `fields` - Additional key-value pairs to include in the log
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// ctx.log_warn_with("Rate limit approaching", &[("current", "95"), ("limit", "100")]);
    /// ```
    pub fn log_warn_with(&self, message: &str, fields: &[(&str, &str)]) {
        self.log_with_level(LogLevel::Warn, message, fields);
    }

    /// Logs a message at ERROR level with automatic context.
    ///
    /// This method automatically includes the durable_execution_arn and parent_id
    /// in the log output without requiring the caller to specify them.
    ///
    /// # Arguments
    ///
    /// * `message` - The message to log
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// ctx.log_error("Failed to process payment");
    /// ```
    pub fn log_error(&self, message: &str) {
        self.log_with_level(LogLevel::Error, message, &[]);
    }

    /// Logs a message at ERROR level with extra fields.
    ///
    /// This method automatically includes the durable_execution_arn and parent_id
    /// in the log output, plus any additional fields specified.
    ///
    /// # Arguments
    ///
    /// * `message` - The message to log
    /// * `fields` - Additional key-value pairs to include in the log
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// ctx.log_error_with("Payment failed", &[("error_code", "INSUFFICIENT_FUNDS"), ("amount", "150.00")]);
    /// ```
    pub fn log_error_with(&self, message: &str, fields: &[(&str, &str)]) {
        self.log_with_level(LogLevel::Error, message, fields);
    }

    /// Internal helper method to log at a specific level with optional extra fields.
    ///
    /// This method creates a LogInfo with automatic context (durable_execution_arn, parent_id)
    /// and any additional fields, then delegates to the configured logger.
    ///
    /// # Arguments
    ///
    /// * `level` - The log level
    /// * `message` - The message to log
    /// * `extra` - Additional key-value pairs to include
    fn log_with_level(&self, level: LogLevel, message: &str, extra: &[(&str, &str)]) {
        let mut log_info = self.create_log_info();

        for (key, value) in extra {
            log_info = log_info.with_extra(*key, *value);
        }

        let logger = self.logger.read().unwrap();
        match level {
            LogLevel::Debug => logger.debug(message, &log_info),
            LogLevel::Info => logger.info(message, &log_info),
            LogLevel::Warn => logger.warn(message, &log_info),
            LogLevel::Error => logger.error(message, &log_info),
        }
    }

    /// Returns the original user input from the EXECUTION operation.
    ///
    /// This method deserializes the input payload from the EXECUTION operation's
    /// ExecutionDetails.InputPayload field into the requested type.
    ///
    /// # Type Parameters
    ///
    /// * `T` - The type to deserialize the input into. Must implement `DeserializeOwned`.
    ///
    /// # Returns
    ///
    /// `Ok(T)` if the input exists and can be deserialized, or a `DurableError` if:
    /// - No EXECUTION operation exists
    /// - No input payload is available
    /// - Deserialization fails
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// #[derive(Deserialize)]
    /// struct OrderEvent {
    ///     order_id: String,
    ///     amount: f64,
    /// }
    ///
    /// async fn my_workflow(ctx: DurableContext) -> Result<(), DurableError> {
    ///     // Get the original input that started this execution
    ///     let event: OrderEvent = ctx.get_original_input()?;
    ///     println!("Processing order: {}", event.order_id);
    ///     Ok(())
    /// }
    /// ```
    pub fn get_original_input<T>(&self) -> DurableResult<T>
    where
        T: serde::de::DeserializeOwned,
    {
        // Get the raw input payload from the execution state
        let input_payload = self.state.get_original_input_raw().ok_or_else(|| {
            crate::error::DurableError::Validation {
                message: "No original input available. The EXECUTION operation may not exist or has no input payload.".to_string(),
            }
        })?;

        // Deserialize the input payload to the requested type
        serde_json::from_str(input_payload).map_err(|e| crate::error::DurableError::SerDes {
            message: format!("Failed to deserialize original input: {}", e),
        })
    }

    /// Returns the raw original user input as a string, if available.
    ///
    /// This method returns the raw JSON string from the EXECUTION operation's
    /// ExecutionDetails.InputPayload field without deserializing it.
    ///
    /// # Returns
    ///
    /// `Some(&str)` if the input exists, `None` otherwise.
    pub fn get_original_input_raw(&self) -> Option<&str> {
        self.state.get_original_input_raw()
    }

    /// Completes the execution with a successful result via checkpointing.
    ///
    /// This method checkpoints a SUCCEED action on the EXECUTION operation,
    /// which is useful for large results that exceed the Lambda response size limit (6MB).
    /// After calling this method, the Lambda function should return an empty result.
    ///
    /// # Arguments
    ///
    /// * `result` - The result to checkpoint. Must implement `Serialize`.
    ///
    /// # Returns
    ///
    /// `Ok(())` on success, or a `DurableError` if:
    /// - No EXECUTION operation exists
    /// - Serialization fails
    /// - The checkpoint fails
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// async fn my_workflow(ctx: DurableContext) -> Result<(), DurableError> {
    ///     let large_result = compute_large_result().await?;
    ///     
    ///     // Check if result would exceed Lambda response limit
    ///     if DurableExecutionInvocationOutput::would_exceed_max_size(&large_result) {
    ///         // Checkpoint the result instead of returning it
    ///         ctx.complete_execution_success(&large_result).await?;
    ///         // Return empty result - the actual result is checkpointed
    ///         return Ok(());
    ///     }
    ///     
    ///     Ok(())
    /// }
    /// ```
    pub async fn complete_execution_success<T>(&self, result: &T) -> DurableResult<()>
    where
        T: serde::Serialize,
    {
        let serialized =
            serde_json::to_string(result).map_err(|e| crate::error::DurableError::SerDes {
                message: format!("Failed to serialize execution result: {}", e),
            })?;

        self.state
            .complete_execution_success(Some(serialized))
            .await
    }

    /// Completes the execution with a failure via checkpointing.
    ///
    /// This method checkpoints a FAIL action on the EXECUTION operation.
    /// After calling this method, the Lambda function should return a FAILED status.
    ///
    /// # Arguments
    ///
    /// * `error` - The error details to checkpoint
    ///
    /// # Returns
    ///
    /// `Ok(())` on success, or a `DurableError` if:
    /// - No EXECUTION operation exists
    /// - The checkpoint fails
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// async fn my_workflow(ctx: DurableContext) -> Result<(), DurableError> {
    ///     if let Err(e) = process_order().await {
    ///         // Checkpoint the failure
    ///         ctx.complete_execution_failure(ErrorObject::new("ProcessingError", &e.to_string())).await?;
    ///         return Err(DurableError::execution(&e.to_string()));
    ///     }
    ///     Ok(())
    /// }
    /// ```
    pub async fn complete_execution_failure(
        &self,
        error: crate::error::ErrorObject,
    ) -> DurableResult<()> {
        self.state.complete_execution_failure(error).await
    }

    /// Completes the execution with a successful result, automatically handling large results.
    ///
    /// This method checks if the result would exceed the Lambda response size limit (6MB).
    /// If so, it checkpoints the result via the EXECUTION operation and returns `true`.
    /// If the result fits within the limit, it returns `false` and the caller should
    /// return the result normally.
    ///
    /// # Arguments
    ///
    /// * `result` - The result to potentially checkpoint. Must implement `Serialize`.
    ///
    /// # Returns
    ///
    /// `Ok(true)` if the result was checkpointed (caller should return empty result),
    /// `Ok(false)` if the result fits within limits (caller should return it normally),
    /// or a `DurableError` if checkpointing fails.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// async fn my_workflow(ctx: DurableContext) -> Result<LargeResult, DurableError> {
    ///     let result = compute_result().await?;
    ///     
    ///     // Automatically handle large results
    ///     if ctx.complete_execution_if_large(&result).await? {
    ///         // Result was checkpointed, return a placeholder
    ///         // The actual result is stored in the EXECUTION operation
    ///         return Ok(LargeResult::default());
    ///     }
    ///     
    ///     // Result fits within limits, return normally
    ///     Ok(result)
    /// }
    /// ```
    pub async fn complete_execution_if_large<T>(&self, result: &T) -> DurableResult<bool>
    where
        T: serde::Serialize,
    {
        if crate::lambda::DurableExecutionInvocationOutput::would_exceed_max_size(result) {
            self.complete_execution_success(result).await?;
            Ok(true)
        } else {
            Ok(false)
        }
    }

    // ========================================================================
    // Durable Operations
    // ========================================================================

    /// Executes a step operation with automatic checkpointing.
    ///
    /// Steps are the fundamental unit of work in durable executions. Each step
    /// is checkpointed, allowing the workflow to resume from the last completed
    /// step after interruptions.
    ///
    /// # Arguments
    ///
    /// * `func` - The function to execute
    /// * `config` - Optional step configuration (retry strategy, semantics, serdes)
    ///
    /// # Returns
    ///
    /// The result of the step function, or an error if execution fails.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// let result: i32 = ctx.step(|_step_ctx| {
    ///     Ok(42)
    /// }, None).await?;
    /// ```
    pub async fn step<T, F>(
        &self,
        func: F,
        config: Option<crate::config::StepConfig>,
    ) -> DurableResult<T>
    where
        T: DurableValue,
        F: StepFn<T>,
    {
        let op_id = self.next_operation_identifier(None);
        let config = config.unwrap_or_default();

        let logger = self.logger.read().unwrap().clone();
        let result =
            crate::handlers::step_handler(func, &self.state, &op_id, &config, &logger).await;

        // Track replay after completion
        if result.is_ok() {
            self.state.track_replay(&op_id.operation_id).await;
        }

        result
    }

    /// Executes a named step operation with automatic checkpointing.
    ///
    /// Same as `step`, but allows specifying a human-readable name for the operation.
    ///
    /// # Arguments
    ///
    /// * `name` - Human-readable name for the step
    /// * `func` - The function to execute
    /// * `config` - Optional step configuration
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// let result: i32 = ctx.step_named("validate_input", |_step_ctx| {
    ///     Ok(42)
    /// }, None).await?;
    /// ```
    pub async fn step_named<T, F>(
        &self,
        name: &str,
        func: F,
        config: Option<crate::config::StepConfig>,
    ) -> DurableResult<T>
    where
        T: DurableValue,
        F: StepFn<T>,
    {
        let op_id = self.next_operation_identifier(Some(name.to_string()));
        let config = config.unwrap_or_default();

        let logger = self.logger.read().unwrap().clone();
        let result =
            crate::handlers::step_handler(func, &self.state, &op_id, &config, &logger).await;

        // Track replay after completion
        if result.is_ok() {
            self.state.track_replay(&op_id.operation_id).await;
        }

        result
    }

    /// Pauses execution for a specified duration.
    ///
    /// Wait operations suspend the Lambda execution and resume after the
    /// specified duration has elapsed. This is efficient because it doesn't
    /// block Lambda resources during the wait.
    ///
    /// # Arguments
    ///
    /// * `duration` - The duration to wait (must be at least 1 second)
    /// * `name` - Optional human-readable name for the operation
    ///
    /// # Returns
    ///
    /// Ok(()) when the wait has elapsed, or an error if validation fails.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// // Wait for 5 seconds
    /// ctx.wait(Duration::from_seconds(5), None).await?;
    ///
    /// // Wait with a name
    /// ctx.wait(Duration::from_minutes(1), Some("wait_for_processing")).await?;
    /// ```
    pub async fn wait(
        &self,
        duration: crate::duration::Duration,
        name: Option<&str>,
    ) -> DurableResult<()> {
        let op_id = self.next_operation_identifier(name.map(|s| s.to_string()));

        let logger = self.logger.read().unwrap().clone();
        let result = crate::handlers::wait_handler(duration, &self.state, &op_id, &logger).await;

        // Track replay after completion (only if not suspended)
        if result.is_ok() {
            self.state.track_replay(&op_id.operation_id).await;
        }

        result
    }

    /// Cancels an active wait operation.
    ///
    /// This method allows cancelling a wait operation that was previously started.
    /// If the wait has already completed (succeeded, failed, or timed out), this
    /// method will return Ok(()) without making any changes.
    ///
    /// # Arguments
    ///
    /// * `operation_id` - The operation ID of the wait to cancel
    ///
    /// # Returns
    ///
    /// Ok(()) if the wait was cancelled or was already completed, or an error if:
    /// - The operation doesn't exist
    /// - The operation is not a WAIT operation
    /// - The checkpoint fails
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// // Start a wait in a parallel branch
    /// let wait_op_id = ctx.next_operation_id();
    ///
    /// // Later, cancel the wait from another branch
    /// ctx.cancel_wait(&wait_op_id).await?;
    /// ```
    pub async fn cancel_wait(&self, operation_id: &str) -> DurableResult<()> {
        let logger = self.logger.read().unwrap().clone();
        crate::handlers::wait_cancel_handler(&self.state, operation_id, &logger).await
    }

    /// Creates a callback and returns a handle to wait for the result.
    ///
    /// Callbacks allow external systems to signal completion of asynchronous
    /// operations. The callback ID can be shared with external systems, which
    /// can then call the Lambda durable execution callback API.
    ///
    /// # Arguments
    ///
    /// * `config` - Optional callback configuration (timeout, heartbeat)
    ///
    /// # Returns
    ///
    /// A `Callback<T>` handle that can be used to wait for the result.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// let callback = ctx.create_callback::<ApprovalResponse>(None).await?;
    ///
    /// // Share callback.callback_id with external system
    /// notify_approver(&callback.callback_id).await?;
    ///
    /// // Wait for the callback result
    /// let approval = callback.result().await?;
    /// ```
    pub async fn create_callback<T>(
        &self,
        config: Option<crate::config::CallbackConfig>,
    ) -> DurableResult<crate::handlers::Callback<T>>
    where
        T: serde::Serialize + serde::de::DeserializeOwned,
    {
        let op_id = self.next_operation_identifier(None);
        let config = config.unwrap_or_default();

        let logger = self.logger.read().unwrap().clone();
        let result = crate::handlers::callback_handler(&self.state, &op_id, &config, &logger).await;

        // Track replay after completion
        if result.is_ok() {
            self.state.track_replay(&op_id.operation_id).await;
        }

        result
    }

    /// Creates a named callback and returns a handle to wait for the result.
    ///
    /// Same as `create_callback`, but allows specifying a human-readable name.
    pub async fn create_callback_named<T>(
        &self,
        name: &str,
        config: Option<crate::config::CallbackConfig>,
    ) -> DurableResult<crate::handlers::Callback<T>>
    where
        T: serde::Serialize + serde::de::DeserializeOwned,
    {
        let op_id = self.next_operation_identifier(Some(name.to_string()));
        let config = config.unwrap_or_default();

        let logger = self.logger.read().unwrap().clone();
        let result = crate::handlers::callback_handler(&self.state, &op_id, &config, &logger).await;

        // Track replay after completion
        if result.is_ok() {
            self.state.track_replay(&op_id.operation_id).await;
        }

        result
    }

    /// Invokes another durable Lambda function.
    ///
    /// This method calls another Lambda function and waits for its result.
    /// The invocation is checkpointed, so if the workflow is interrupted,
    /// it will resume with the result of the invocation.
    ///
    /// # Arguments
    ///
    /// * `function_name` - The name or ARN of the Lambda function to invoke
    /// * `payload` - The payload to send to the function
    /// * `config` - Optional invoke configuration (timeout, serdes)
    ///
    /// # Returns
    ///
    /// The result from the invoked function, or an error if invocation fails.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// let result: ProcessingResult = ctx.invoke(
    ///     "process-order-function",
    ///     OrderPayload { order_id: "123".to_string() },
    ///     None,
    /// ).await?;
    /// ```
    pub async fn invoke<P, R>(
        &self,
        function_name: &str,
        payload: P,
        config: Option<crate::config::InvokeConfig<P, R>>,
    ) -> DurableResult<R>
    where
        P: serde::Serialize + serde::de::DeserializeOwned + Send,
        R: serde::Serialize + serde::de::DeserializeOwned + Send,
    {
        let op_id = self.next_operation_identifier(Some(format!("invoke:{}", function_name)));
        let config = config.unwrap_or_default();

        let logger = self.logger.read().unwrap().clone();
        let result = crate::handlers::invoke_handler(
            function_name,
            payload,
            &self.state,
            &op_id,
            &config,
            &logger,
        )
        .await;

        // Track replay after completion
        if result.is_ok() {
            self.state.track_replay(&op_id.operation_id).await;
        }

        result
    }

    /// Processes a collection in parallel with configurable concurrency.
    ///
    /// Map operations execute a function for each item in the collection,
    /// with configurable concurrency limits and failure tolerance.
    ///
    /// # Arguments
    ///
    /// * `items` - The collection of items to process
    /// * `func` - The function to apply to each item
    /// * `config` - Optional map configuration (concurrency, completion criteria)
    ///
    /// # Returns
    ///
    /// A `BatchResult<U>` containing results for all items.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// let results = ctx.map(
    ///     vec![1, 2, 3, 4, 5],
    ///     |child_ctx, item, index| async move {
    ///         Ok(item * 2)
    ///     },
    ///     Some(MapConfig {
    ///         max_concurrency: Some(3),
    ///         ..Default::default()
    ///     }),
    /// ).await?;
    /// ```
    pub async fn map<T, U, F, Fut>(
        &self,
        items: Vec<T>,
        func: F,
        config: Option<crate::config::MapConfig>,
    ) -> DurableResult<crate::concurrency::BatchResult<U>>
    where
        T: serde::Serialize + serde::de::DeserializeOwned + Send + Sync + Clone + 'static,
        U: serde::Serialize + serde::de::DeserializeOwned + Send + 'static,
        F: Fn(DurableContext, T, usize) -> Fut + Send + Sync + Clone + 'static,
        Fut: std::future::Future<Output = DurableResult<U>> + Send + 'static,
    {
        let op_id = self.next_operation_identifier(Some("map".to_string()));
        let config = config.unwrap_or_default();

        let logger = self.logger.read().unwrap().clone();
        let result =
            crate::handlers::map_handler(items, func, &self.state, &op_id, self, &config, &logger)
                .await;

        // Track replay after completion
        if result.is_ok() {
            self.state.track_replay(&op_id.operation_id).await;
        }

        result
    }

    /// Executes multiple operations in parallel.
    ///
    /// Parallel operations execute multiple independent functions concurrently,
    /// with configurable concurrency limits and completion criteria.
    ///
    /// # Arguments
    ///
    /// * `branches` - The list of functions to execute in parallel
    /// * `config` - Optional parallel configuration (concurrency, completion criteria)
    ///
    /// # Returns
    ///
    /// A `BatchResult<T>` containing results for all branches.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// let results = ctx.parallel(
    ///     vec![
    ///         |ctx| async move { Ok(fetch_data_a(&ctx).await?) },
    ///         |ctx| async move { Ok(fetch_data_b(&ctx).await?) },
    ///         |ctx| async move { Ok(fetch_data_c(&ctx).await?) },
    ///     ],
    ///     None,
    /// ).await?;
    /// ```
    pub async fn parallel<T, F, Fut>(
        &self,
        branches: Vec<F>,
        config: Option<crate::config::ParallelConfig>,
    ) -> DurableResult<crate::concurrency::BatchResult<T>>
    where
        T: serde::Serialize + serde::de::DeserializeOwned + Send + 'static,
        F: FnOnce(DurableContext) -> Fut + Send + 'static,
        Fut: std::future::Future<Output = DurableResult<T>> + Send + 'static,
    {
        let op_id = self.next_operation_identifier(Some("parallel".to_string()));
        let config = config.unwrap_or_default();

        let logger = self.logger.read().unwrap().clone();
        let result = crate::handlers::parallel_handler(
            branches,
            &self.state,
            &op_id,
            self,
            &config,
            &logger,
        )
        .await;

        // Track replay after completion
        if result.is_ok() {
            self.state.track_replay(&op_id.operation_id).await;
        }

        result
    }

    /// Executes a function in a child context.
    ///
    /// Child contexts provide isolation for nested workflows. Operations in
    /// a child context are tracked separately and can be checkpointed as a unit.
    ///
    /// # Arguments
    ///
    /// * `func` - The function to execute in the child context
    /// * `config` - Optional child context configuration
    ///
    /// # Returns
    ///
    /// The result of the child function, or an error if execution fails.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// let result = ctx.run_in_child_context(|child_ctx| async move {
    ///     let step1 = child_ctx.step(|_| Ok(1), None).await?;
    ///     let step2 = child_ctx.step(|_| Ok(2), None).await?;
    ///     Ok(step1 + step2)
    /// }, None).await?;
    /// ```
    pub async fn run_in_child_context<T, F, Fut>(
        &self,
        func: F,
        config: Option<crate::config::ChildConfig>,
    ) -> DurableResult<T>
    where
        T: serde::Serialize + serde::de::DeserializeOwned + Send,
        F: FnOnce(DurableContext) -> Fut + Send,
        Fut: std::future::Future<Output = DurableResult<T>> + Send,
    {
        let op_id = self.next_operation_identifier(Some("child_context".to_string()));
        let config = config.unwrap_or_default();

        let logger = self.logger.read().unwrap().clone();
        let result =
            crate::handlers::child_handler(func, &self.state, &op_id, self, &config, &logger).await;

        // Track replay after completion
        if result.is_ok() {
            self.state.track_replay(&op_id.operation_id).await;
        }

        result
    }

    /// Executes a named function in a child context.
    ///
    /// Same as `run_in_child_context`, but allows specifying a human-readable name.
    pub async fn run_in_child_context_named<T, F, Fut>(
        &self,
        name: &str,
        func: F,
        config: Option<crate::config::ChildConfig>,
    ) -> DurableResult<T>
    where
        T: serde::Serialize + serde::de::DeserializeOwned + Send,
        F: FnOnce(DurableContext) -> Fut + Send,
        Fut: std::future::Future<Output = DurableResult<T>> + Send,
    {
        let op_id = self.next_operation_identifier(Some(name.to_string()));
        let config = config.unwrap_or_default();

        let logger = self.logger.read().unwrap().clone();
        let result =
            crate::handlers::child_handler(func, &self.state, &op_id, self, &config, &logger).await;

        // Track replay after completion
        if result.is_ok() {
            self.state.track_replay(&op_id.operation_id).await;
        }

        result
    }

    /// Polls until a condition is met.
    ///
    /// This method repeatedly checks a condition until it returns a successful
    /// result. Between checks, it waits for a configurable duration using the
    /// RETRY mechanism with NextAttemptDelaySeconds.
    ///
    /// # Implementation
    ///
    /// This method is implemented as a single STEP operation with RETRY mechanism,
    /// which is more efficient than using multiple steps and waits. The state is
    /// passed as Payload on retry (not Error), and the attempt number is tracked
    /// in StepDetails.Attempt.
    ///
    /// # Arguments
    ///
    /// * `check` - The function to check the condition
    /// * `config` - Configuration for the wait (interval, max attempts, timeout)
    ///
    /// # Returns
    ///
    /// The result when the condition is met, or an error if timeout/max attempts exceeded.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// let result = ctx.wait_for_condition(
    ///     |state, ctx| {
    ///         // Check if order is ready
    ///         let status = check_order_status(&state.order_id)?;
    ///         if status == "ready" {
    ///             Ok(OrderReady { order_id: state.order_id.clone() })
    ///         } else {
    ///             Err("Order not ready yet".into())
    ///         }
    ///     },
    ///     WaitForConditionConfig::from_interval(
    ///         OrderState { order_id: "123".to_string() },
    ///         Duration::from_seconds(5),
    ///         Some(10),
    ///     ),
    /// ).await?;
    /// ```
    pub async fn wait_for_condition<T, S, F>(
        &self,
        check: F,
        config: WaitForConditionConfig<S>,
    ) -> DurableResult<T>
    where
        T: serde::Serialize + serde::de::DeserializeOwned + Send,
        S: serde::Serialize + serde::de::DeserializeOwned + Clone + Send + Sync,
        F: Fn(&S, &WaitForConditionContext) -> Result<T, Box<dyn std::error::Error + Send + Sync>>
            + Send
            + Sync,
    {
        let op_id = self.next_operation_identifier(Some("wait_for_condition".to_string()));

        let logger = self.logger.read().unwrap().clone();
        // Use the new wait_for_condition_handler which implements the single STEP with RETRY pattern
        // This is more efficient than the previous child context + multiple steps approach
        let result = crate::handlers::wait_for_condition_handler(
            check,
            config,
            &self.state,
            &op_id,
            &logger,
        )
        .await;

        // Track replay after completion (only if not suspended)
        if result.is_ok() {
            self.state.track_replay(&op_id.operation_id).await;
        }

        result
    }

    /// Creates a callback and waits for the result with a submitter function.
    ///
    /// This is a convenience method that combines callback creation with
    /// a submitter function that sends the callback ID to an external system.
    /// The submitter execution is checkpointed within a child context to ensure
    /// replay safety - the submitter will not be re-executed during replay.
    ///
    /// # Arguments
    ///
    /// * `submitter` - Function that receives the callback ID and submits it to external system
    /// * `config` - Optional callback configuration (timeout, heartbeat)
    ///
    /// # Returns
    ///
    /// The callback result from the external system.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// let approval = ctx.wait_for_callback(
    ///     |callback_id| async move {
    ///         // Send callback ID to approval system
    ///         send_approval_request(&callback_id, &request).await
    ///     },
    ///     Some(CallbackConfig {
    ///         timeout: Duration::from_hours(24),
    ///         ..Default::default()
    ///     }),
    /// ).await?;
    /// ```
    pub async fn wait_for_callback<T, F, Fut>(
        &self,
        submitter: F,
        config: Option<crate::config::CallbackConfig>,
    ) -> DurableResult<T>
    where
        T: serde::Serialize + serde::de::DeserializeOwned + Send + Sync,
        F: FnOnce(String) -> Fut + Send + 'static,
        Fut: std::future::Future<Output = Result<(), Box<dyn std::error::Error + Send + Sync>>>
            + Send
            + 'static,
    {
        // Create the callback first with the provided configuration (Requirement 1.4)
        let callback: crate::handlers::Callback<T> = self.create_callback(config).await?;
        let callback_id = callback.callback_id.clone();

        // Generate operation ID for the child context that will execute the submitter
        let op_id = self.next_operation_identifier(Some("wait_for_callback_submitter".to_string()));

        // Execute the submitter within a child context for proper checkpointing (Requirements 1.1, 1.2)
        // The child context ensures the submitter execution is tracked and not re-executed during replay
        let child_config = crate::config::ChildConfig::default();

        let logger = self.logger.read().unwrap().clone();
        crate::handlers::child_handler(
            |child_ctx| {
                let callback_id = callback_id.clone();
                async move {
                    // Use a step to checkpoint that we're about to execute the submitter
                    // This step returns () on success, indicating submission completed
                    child_ctx
                        .step_named(
                            "execute_submitter",
                            move |_| {
                                // The step just marks that we're executing the submitter
                                // The actual async submitter call happens after this checkpoint
                                Ok(())
                            },
                            None,
                        )
                        .await?;

                    // Now execute the actual submitter
                    // If we're replaying and the step above succeeded, we won't reach here
                    // because the child context will return the checkpointed result
                    submitter(callback_id).await.map_err(|e| {
                        crate::error::DurableError::UserCode {
                            message: e.to_string(),
                            error_type: "SubmitterError".to_string(),
                            stack_trace: None,
                        }
                    })?;

                    Ok(())
                }
            },
            &self.state,
            &op_id,
            self,
            &child_config,
            &logger,
        )
        .await?;

        // Track replay after child context completion
        self.state.track_replay(&op_id.operation_id).await;

        // Wait for the callback result
        callback.result().await
    }

    // ========================================================================
    // Promise Combinators
    // ========================================================================

    /// Waits for all futures to complete successfully.
    ///
    /// Returns all results if all futures succeed, or returns the first error encountered.
    /// This is implemented within a STEP operation for durability.
    ///
    /// # Arguments
    ///
    /// * `futures` - Vector of futures to execute
    ///
    /// # Returns
    ///
    /// `Ok(Vec<T>)` if all futures succeed, or `Err` with the first error.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// let results = ctx.all(vec![
    ///     ctx.step(|_| Ok(1), None),
    ///     ctx.step(|_| Ok(2), None),
    ///     ctx.step(|_| Ok(3), None),
    /// ]).await?;
    /// assert_eq!(results, vec![1, 2, 3]);
    /// ```
    pub async fn all<T, Fut>(&self, futures: Vec<Fut>) -> DurableResult<Vec<T>>
    where
        T: serde::Serialize + serde::de::DeserializeOwned + Send + Clone + 'static,
        Fut: std::future::Future<Output = DurableResult<T>> + Send + 'static,
    {
        let op_id = self.next_operation_identifier(Some("all".to_string()));

        let logger = self.logger.read().unwrap().clone();
        let result = crate::handlers::all_handler(futures, &self.state, &op_id, &logger).await;

        // Track replay after completion
        if result.is_ok() {
            self.state.track_replay(&op_id.operation_id).await;
        }

        result
    }

    /// Waits for all futures to settle (success or failure).
    ///
    /// Returns a BatchResult containing outcomes for all futures, regardless of success or failure.
    /// This is implemented within a STEP operation for durability.
    ///
    /// # Arguments
    ///
    /// * `futures` - Vector of futures to execute
    ///
    /// # Returns
    ///
    /// `BatchResult<T>` containing results for all futures.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// let results = ctx.all_settled(vec![
    ///     ctx.step(|_| Ok(1), None),
    ///     ctx.step(|_| Err("error".into()), None),
    ///     ctx.step(|_| Ok(3), None),
    /// ]).await?;
    /// assert_eq!(results.success_count(), 2);
    /// assert_eq!(results.failure_count(), 1);
    /// ```
    pub async fn all_settled<T, Fut>(
        &self,
        futures: Vec<Fut>,
    ) -> DurableResult<crate::concurrency::BatchResult<T>>
    where
        T: serde::Serialize + serde::de::DeserializeOwned + Send + Clone + 'static,
        Fut: std::future::Future<Output = DurableResult<T>> + Send + 'static,
    {
        let op_id = self.next_operation_identifier(Some("all_settled".to_string()));

        let logger = self.logger.read().unwrap().clone();
        let result =
            crate::handlers::all_settled_handler(futures, &self.state, &op_id, &logger).await;

        // Track replay after completion
        if result.is_ok() {
            self.state.track_replay(&op_id.operation_id).await;
        }

        result
    }

    /// Returns the result of the first future to settle.
    ///
    /// Returns the result (success or failure) of whichever future completes first.
    /// This is implemented within a STEP operation for durability.
    ///
    /// # Arguments
    ///
    /// * `futures` - Vector of futures to execute
    ///
    /// # Returns
    ///
    /// The result of the first future to settle.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// let result = ctx.race(vec![
    ///     ctx.step(|_| Ok(1), None),
    ///     ctx.step(|_| Ok(2), None),
    /// ]).await?;
    /// // result is either 1 or 2, whichever completed first
    /// ```
    pub async fn race<T, Fut>(&self, futures: Vec<Fut>) -> DurableResult<T>
    where
        T: serde::Serialize + serde::de::DeserializeOwned + Send + Clone + 'static,
        Fut: std::future::Future<Output = DurableResult<T>> + Send + 'static,
    {
        let op_id = self.next_operation_identifier(Some("race".to_string()));

        let logger = self.logger.read().unwrap().clone();
        let result = crate::handlers::race_handler(futures, &self.state, &op_id, &logger).await;

        // Track replay after completion
        if result.is_ok() {
            self.state.track_replay(&op_id.operation_id).await;
        }

        result
    }

    /// Returns the result of the first future to succeed.
    ///
    /// Returns the result of the first future to succeed. If all futures fail,
    /// returns an error containing all the failures.
    /// This is implemented within a STEP operation for durability.
    ///
    /// # Arguments
    ///
    /// * `futures` - Vector of futures to execute
    ///
    /// # Returns
    ///
    /// The result of the first future to succeed, or an error if all fail.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// let result = ctx.any(vec![
    ///     ctx.step(|_| Err("error".into()), None),
    ///     ctx.step(|_| Ok(2), None),
    ///     ctx.step(|_| Ok(3), None),
    /// ]).await?;
    /// // result is either 2 or 3, whichever succeeded first
    /// ```
    pub async fn any<T, Fut>(&self, futures: Vec<Fut>) -> DurableResult<T>
    where
        T: serde::Serialize + serde::de::DeserializeOwned + Send + Clone + 'static,
        Fut: std::future::Future<Output = DurableResult<T>> + Send + 'static,
    {
        let op_id = self.next_operation_identifier(Some("any".to_string()));

        let logger = self.logger.read().unwrap().clone();
        let result = crate::handlers::any_handler(futures, &self.state, &op_id, &logger).await;

        // Track replay after completion
        if result.is_ok() {
            self.state.track_replay(&op_id.operation_id).await;
        }

        result
    }
}

/// Configuration for wait_for_condition operations.
#[allow(clippy::type_complexity)]
pub struct WaitForConditionConfig<S> {
    /// Initial state to pass to the check function.
    pub initial_state: S,
    /// Wait strategy function that determines polling behavior.
    ///
    /// Takes a reference to the current state and the attempt number (1-indexed),
    /// and returns a [`WaitDecision`](crate::config::WaitDecision).
    pub wait_strategy: Box<dyn Fn(&S, usize) -> crate::config::WaitDecision + Send + Sync>,
    /// Overall timeout for the operation.
    pub timeout: Option<crate::duration::Duration>,
    /// Optional custom serializer/deserializer.
    pub serdes: Option<std::sync::Arc<dyn crate::config::SerDesAny>>,
}

impl<S> std::fmt::Debug for WaitForConditionConfig<S>
where
    S: std::fmt::Debug,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("WaitForConditionConfig")
            .field("initial_state", &self.initial_state)
            .field("wait_strategy", &"<fn>")
            .field("timeout", &self.timeout)
            .field("serdes", &self.serdes.is_some())
            .finish()
    }
}

impl<S> WaitForConditionConfig<S> {
    /// Creates a backward-compatible `WaitForConditionConfig` from interval and max_attempts.
    ///
    /// This constructor converts the old-style `interval` + `max_attempts` parameters
    /// into a wait strategy that uses a fixed delay (no backoff, no jitter).
    ///
    /// # Arguments
    ///
    /// * `initial_state` - Initial state to pass to the check function.
    /// * `interval` - Fixed interval between condition checks.
    /// * `max_attempts` - Maximum number of attempts (`None` for unlimited).
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// use durable_execution_sdk::{WaitForConditionConfig, Duration};
    ///
    /// let config = WaitForConditionConfig::from_interval(
    ///     my_initial_state,
    ///     Duration::from_seconds(5),
    ///     Some(10),
    /// );
    /// ```
    pub fn from_interval(
        initial_state: S,
        interval: crate::duration::Duration,
        max_attempts: Option<usize>,
    ) -> Self
    where
        S: Send + Sync + 'static,
    {
        let interval_secs = interval.to_seconds();
        let max = max_attempts.unwrap_or(usize::MAX);

        Self {
            initial_state,
            wait_strategy: Box::new(move |_state: &S, attempts_made: usize| {
                // In the backward-compatible mode, the check function's Ok/Err
                // determines whether polling is done. The strategy only provides
                // the delay and enforces max_attempts.
                // Return Done when max_attempts exceeded (handler will checkpoint failure
                // if check returned Err, or success if check returned Ok).
                if attempts_made >= max {
                    return crate::config::WaitDecision::Done;
                }
                crate::config::WaitDecision::Continue {
                    delay: crate::duration::Duration::from_seconds(interval_secs),
                }
            }),
            timeout: None,
            serdes: None,
        }
    }
}

impl<S: Default + Send + Sync + 'static> Default for WaitForConditionConfig<S> {
    fn default() -> Self {
        Self::from_interval(
            S::default(),
            crate::duration::Duration::from_seconds(5),
            None,
        )
    }
}

/// Context provided to wait_for_condition check functions.
#[derive(Debug, Clone)]
pub struct WaitForConditionContext {
    /// Current attempt number (1-indexed).
    pub attempt: usize,
    /// Maximum number of attempts (None for unlimited).
    pub max_attempts: Option<usize>,
}

impl Clone for DurableContext {
    fn clone(&self) -> Self {
        Self {
            state: self.state.clone(),
            lambda_context: self.lambda_context.clone(),
            parent_id: self.parent_id.clone(),
            id_generator: self.id_generator.clone(),
            logger: self.logger.clone(),
        }
    }
}

impl std::fmt::Debug for DurableContext {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("DurableContext")
            .field("durable_execution_arn", &self.durable_execution_arn())
            .field("parent_id", &self.parent_id)
            .field("step_counter", &self.current_step_counter())
            .finish_non_exhaustive()
    }
}

// Add hex encoding dependency
mod hex {
    const HEX_CHARS: &[u8; 16] = b"0123456789abcdef";

    pub fn encode(bytes: &[u8]) -> String {
        let mut result = String::with_capacity(bytes.len() * 2);
        for &byte in bytes {
            result.push(HEX_CHARS[(byte >> 4) as usize] as char);
            result.push(HEX_CHARS[(byte & 0x0f) as usize] as char);
        }
        result
    }
}

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

    #[test]
    fn test_operation_identifier_new() {
        let id = OperationIdentifier::new(
            "op-123",
            Some("parent-456".to_string()),
            Some("my-step".to_string()),
        );
        assert_eq!(id.operation_id, "op-123");
        assert_eq!(id.parent_id, Some("parent-456".to_string()));
        assert_eq!(id.name, Some("my-step".to_string()));
    }

    #[test]
    fn test_operation_identifier_root() {
        let id = OperationIdentifier::root("op-123");
        assert_eq!(id.operation_id, "op-123");
        assert!(id.parent_id.is_none());
        assert!(id.name.is_none());
    }

    #[test]
    fn test_operation_identifier_with_parent() {
        let id = OperationIdentifier::with_parent("op-123", "parent-456");
        assert_eq!(id.operation_id, "op-123");
        assert_eq!(id.parent_id, Some("parent-456".to_string()));
        assert!(id.name.is_none());
    }

    #[test]
    fn test_operation_identifier_with_name() {
        let id = OperationIdentifier::root("op-123").with_name("my-step");
        assert_eq!(id.name, Some("my-step".to_string()));
    }

    #[test]
    fn test_operation_identifier_display() {
        let id = OperationIdentifier::root("op-123");
        assert_eq!(format!("{}", id), "op-123");

        let id_with_name = OperationIdentifier::root("op-123").with_name("my-step");
        assert_eq!(format!("{}", id_with_name), "my-step(op-123)");
    }

    #[test]
    fn test_generate_operation_id_deterministic() {
        let id1 = generate_operation_id("base-123", 0);
        let id2 = generate_operation_id("base-123", 0);
        assert_eq!(id1, id2);
    }

    #[test]
    fn test_generate_operation_id_different_counters() {
        let id1 = generate_operation_id("base-123", 0);
        let id2 = generate_operation_id("base-123", 1);
        assert_ne!(id1, id2);
    }

    #[test]
    fn test_generate_operation_id_different_bases() {
        let id1 = generate_operation_id("base-123", 0);
        let id2 = generate_operation_id("base-456", 0);
        assert_ne!(id1, id2);
    }

    #[test]
    fn test_generate_operation_id_length() {
        let id = generate_operation_id("base-123", 0);
        assert_eq!(id.len(), 32); // 16 bytes = 32 hex chars
    }

    #[test]
    fn test_operation_id_generator_new() {
        let gen = OperationIdGenerator::new("base-123");
        assert_eq!(gen.base_id(), "base-123");
        assert_eq!(gen.current_counter(), 0);
    }

    #[test]
    fn test_operation_id_generator_with_counter() {
        let gen = OperationIdGenerator::with_counter("base-123", 10);
        assert_eq!(gen.current_counter(), 10);
    }

    #[test]
    fn test_operation_id_generator_next_id() {
        let gen = OperationIdGenerator::new("base-123");

        let id1 = gen.next_id();
        assert_eq!(gen.current_counter(), 1);

        let id2 = gen.next_id();
        assert_eq!(gen.current_counter(), 2);

        assert_ne!(id1, id2);
    }

    #[test]
    fn test_operation_id_generator_id_for_counter() {
        let gen = OperationIdGenerator::new("base-123");

        let id_0 = gen.id_for_counter(0);
        let id_1 = gen.id_for_counter(1);

        // id_for_counter should not increment the counter
        assert_eq!(gen.current_counter(), 0);

        // Should match what next_id would produce
        let next = gen.next_id();
        assert_eq!(next, id_0);

        let next = gen.next_id();
        assert_eq!(next, id_1);
    }

    #[test]
    fn test_operation_id_generator_create_child() {
        let gen = OperationIdGenerator::new("base-123");
        gen.next_id(); // Increment parent counter

        let child = gen.create_child("child-op-id");
        assert_eq!(child.base_id(), "child-op-id");
        assert_eq!(child.current_counter(), 0);
    }

    #[test]
    fn test_operation_id_generator_clone() {
        let gen = OperationIdGenerator::new("base-123");
        gen.next_id();
        gen.next_id();

        let cloned = gen.clone();
        assert_eq!(cloned.base_id(), gen.base_id());
        assert_eq!(cloned.current_counter(), gen.current_counter());
    }

    #[test]
    fn test_operation_id_generator_thread_safety() {
        use std::thread;

        let gen = Arc::new(OperationIdGenerator::new("base-123"));
        let mut handles = vec![];

        for _ in 0..10 {
            let gen_clone = gen.clone();
            handles.push(thread::spawn(move || {
                let mut ids = vec![];
                for _ in 0..100 {
                    ids.push(gen_clone.next_id());
                }
                ids
            }));
        }

        let mut all_ids = vec![];
        for handle in handles {
            all_ids.extend(handle.join().unwrap());
        }

        // All IDs should be unique
        let unique_count = {
            let mut set = std::collections::HashSet::new();
            for id in &all_ids {
                set.insert(id.clone());
            }
            set.len()
        };

        assert_eq!(unique_count, 1000);
        assert_eq!(gen.current_counter(), 1000);
    }

    #[test]
    fn test_log_info_new() {
        let info = LogInfo::new("arn:test");
        assert_eq!(info.durable_execution_arn, Some("arn:test".to_string()));
        assert!(info.operation_id.is_none());
        assert!(info.parent_id.is_none());
    }

    #[test]
    fn test_log_info_with_operation_id() {
        let info = LogInfo::new("arn:test").with_operation_id("op-123");
        assert_eq!(info.operation_id, Some("op-123".to_string()));
    }

    #[test]
    fn test_log_info_with_parent_id() {
        let info = LogInfo::new("arn:test").with_parent_id("parent-456");
        assert_eq!(info.parent_id, Some("parent-456".to_string()));
    }

    #[test]
    fn test_log_info_with_extra() {
        let info = LogInfo::new("arn:test")
            .with_extra("key1", "value1")
            .with_extra("key2", "value2");
        assert_eq!(info.extra.len(), 2);
        assert_eq!(info.extra[0], ("key1".to_string(), "value1".to_string()));
    }

    #[test]
    fn test_hex_encode() {
        assert_eq!(hex::encode(&[0x00]), "00");
        assert_eq!(hex::encode(&[0xff]), "ff");
        assert_eq!(hex::encode(&[0x12, 0x34, 0xab, 0xcd]), "1234abcd");
    }
}

#[cfg(test)]
mod durable_context_tests {
    use super::*;
    use crate::client::SharedDurableServiceClient;
    use crate::lambda::InitialExecutionState;
    use std::sync::Arc;

    #[cfg(test)]
    fn create_mock_client() -> SharedDurableServiceClient {
        use crate::client::MockDurableServiceClient;
        Arc::new(MockDurableServiceClient::new())
    }

    fn create_test_state() -> Arc<ExecutionState> {
        let client = create_mock_client();
        Arc::new(ExecutionState::new(
            "arn:aws:lambda:us-east-1:123456789012:function:test:durable:abc123",
            "token-123",
            InitialExecutionState::new(),
            client,
        ))
    }

    #[test]
    fn test_durable_context_new() {
        let state = create_test_state();
        let ctx = DurableContext::new(state.clone());

        assert_eq!(
            ctx.durable_execution_arn(),
            "arn:aws:lambda:us-east-1:123456789012:function:test:durable:abc123"
        );
        assert!(ctx.parent_id().is_none());
        assert!(ctx.lambda_context().is_none());
        assert_eq!(ctx.current_step_counter(), 0);
    }

    #[test]
    fn test_durable_context_next_operation_id() {
        let state = create_test_state();
        let ctx = DurableContext::new(state);

        let id1 = ctx.next_operation_id();
        let id2 = ctx.next_operation_id();

        assert_ne!(id1, id2);
        assert_eq!(ctx.current_step_counter(), 2);
    }

    #[test]
    fn test_durable_context_next_operation_identifier() {
        let state = create_test_state();
        let ctx = DurableContext::new(state);

        let op_id = ctx.next_operation_identifier(Some("my-step".to_string()));

        assert!(op_id.parent_id.is_none());
        assert_eq!(op_id.name, Some("my-step".to_string()));
        assert!(!op_id.operation_id.is_empty());
    }

    #[test]
    fn test_durable_context_create_child_context() {
        let state = create_test_state();
        let ctx = DurableContext::new(state);

        // Generate a parent operation ID
        let parent_op_id = ctx.next_operation_id();

        // Create child context
        let child_ctx = ctx.create_child_context(&parent_op_id);

        assert_eq!(child_ctx.parent_id(), Some(parent_op_id.as_str()));
        assert_eq!(child_ctx.current_step_counter(), 0);
        assert_eq!(
            child_ctx.durable_execution_arn(),
            ctx.durable_execution_arn()
        );
    }

    #[test]
    fn test_durable_context_child_generates_different_ids() {
        let state = create_test_state();
        let ctx = DurableContext::new(state);

        let parent_op_id = ctx.next_operation_id();
        let child_ctx = ctx.create_child_context(&parent_op_id);

        // Child should generate different IDs than parent
        let child_id = child_ctx.next_operation_id();
        let parent_id = ctx.next_operation_id();

        assert_ne!(child_id, parent_id);
    }

    #[test]
    fn test_durable_context_child_operation_identifier_has_parent() {
        let state = create_test_state();
        let ctx = DurableContext::new(state);

        let parent_op_id = ctx.next_operation_id();
        let child_ctx = ctx.create_child_context(&parent_op_id);

        let child_op_id = child_ctx.next_operation_identifier(None);

        assert_eq!(child_op_id.parent_id, Some(parent_op_id));
    }

    #[test]
    fn test_durable_context_set_logger() {
        let state = create_test_state();
        let mut ctx = DurableContext::new(state);

        // Create a custom logger
        let custom_logger: Arc<dyn Logger> = Arc::new(TracingLogger);
        ctx.set_logger(custom_logger);

        // Just verify it doesn't panic - the logger is set
        let _ = ctx.logger();
    }

    #[test]
    fn test_durable_context_with_logger() {
        let state = create_test_state();
        let ctx = DurableContext::new(state);

        let custom_logger: Arc<dyn Logger> = Arc::new(TracingLogger);
        let ctx_with_logger = ctx.with_logger(custom_logger);

        // Just verify it doesn't panic - the logger is set
        let _ = ctx_with_logger.logger();
    }

    #[test]
    fn test_durable_context_create_log_info() {
        let state = create_test_state();
        let ctx = DurableContext::new(state);

        let info = ctx.create_log_info();

        assert_eq!(
            info.durable_execution_arn,
            Some("arn:aws:lambda:us-east-1:123456789012:function:test:durable:abc123".to_string())
        );
        assert!(info.parent_id.is_none());
    }

    #[test]
    fn test_durable_context_create_log_info_with_parent() {
        let state = create_test_state();
        let ctx = DurableContext::new(state);

        let parent_op_id = ctx.next_operation_id();
        let child_ctx = ctx.create_child_context(&parent_op_id);

        let info = child_ctx.create_log_info();

        assert_eq!(info.parent_id, Some(parent_op_id));
    }

    #[test]
    fn test_durable_context_create_log_info_with_operation() {
        let state = create_test_state();
        let ctx = DurableContext::new(state);

        let info = ctx.create_log_info_with_operation("op-123");

        assert_eq!(info.operation_id, Some("op-123".to_string()));
    }

    #[test]
    fn test_log_info_with_replay() {
        let info = LogInfo::new("arn:test")
            .with_operation_id("op-123")
            .with_replay(true);

        assert!(info.is_replay);
        assert_eq!(info.operation_id, Some("op-123".to_string()));
    }

    #[test]
    fn test_log_info_default_not_replay() {
        let info = LogInfo::default();
        assert!(!info.is_replay);
    }

    #[test]
    fn test_replay_logging_config_default() {
        let config = ReplayLoggingConfig::default();
        assert_eq!(config, ReplayLoggingConfig::SuppressAll);
    }

    #[test]
    fn test_replay_aware_logger_suppress_all() {
        use std::sync::atomic::{AtomicUsize, Ordering};

        // Create counters for each log level
        let debug_count = Arc::new(AtomicUsize::new(0));
        let info_count = Arc::new(AtomicUsize::new(0));
        let warn_count = Arc::new(AtomicUsize::new(0));
        let error_count = Arc::new(AtomicUsize::new(0));

        let inner = Arc::new(custom_logger(
            {
                let count = debug_count.clone();
                move |_, _| {
                    count.fetch_add(1, Ordering::SeqCst);
                }
            },
            {
                let count = info_count.clone();
                move |_, _| {
                    count.fetch_add(1, Ordering::SeqCst);
                }
            },
            {
                let count = warn_count.clone();
                move |_, _| {
                    count.fetch_add(1, Ordering::SeqCst);
                }
            },
            {
                let count = error_count.clone();
                move |_, _| {
                    count.fetch_add(1, Ordering::SeqCst);
                }
            },
        ));

        let logger = ReplayAwareLogger::new(inner, ReplayLoggingConfig::SuppressAll);

        // Non-replay logs should pass through
        let non_replay_info = LogInfo::new("arn:test").with_replay(false);
        logger.debug("test", &non_replay_info);
        logger.info("test", &non_replay_info);
        logger.warn("test", &non_replay_info);
        logger.error("test", &non_replay_info);

        assert_eq!(debug_count.load(Ordering::SeqCst), 1);
        assert_eq!(info_count.load(Ordering::SeqCst), 1);
        assert_eq!(warn_count.load(Ordering::SeqCst), 1);
        assert_eq!(error_count.load(Ordering::SeqCst), 1);

        // Replay logs should be suppressed
        let replay_info = LogInfo::new("arn:test").with_replay(true);
        logger.debug("test", &replay_info);
        logger.info("test", &replay_info);
        logger.warn("test", &replay_info);
        logger.error("test", &replay_info);

        // Counts should not have increased
        assert_eq!(debug_count.load(Ordering::SeqCst), 1);
        assert_eq!(info_count.load(Ordering::SeqCst), 1);
        assert_eq!(warn_count.load(Ordering::SeqCst), 1);
        assert_eq!(error_count.load(Ordering::SeqCst), 1);
    }

    #[test]
    fn test_replay_aware_logger_allow_all() {
        use std::sync::atomic::{AtomicUsize, Ordering};

        let call_count = Arc::new(AtomicUsize::new(0));

        let inner = Arc::new(custom_logger(
            {
                let count = call_count.clone();
                move |_, _| {
                    count.fetch_add(1, Ordering::SeqCst);
                }
            },
            {
                let count = call_count.clone();
                move |_, _| {
                    count.fetch_add(1, Ordering::SeqCst);
                }
            },
            {
                let count = call_count.clone();
                move |_, _| {
                    count.fetch_add(1, Ordering::SeqCst);
                }
            },
            {
                let count = call_count.clone();
                move |_, _| {
                    count.fetch_add(1, Ordering::SeqCst);
                }
            },
        ));

        let logger = ReplayAwareLogger::allow_all(inner);

        // All logs should pass through even during replay
        let replay_info = LogInfo::new("arn:test").with_replay(true);
        logger.debug("test", &replay_info);
        logger.info("test", &replay_info);
        logger.warn("test", &replay_info);
        logger.error("test", &replay_info);

        assert_eq!(call_count.load(Ordering::SeqCst), 4);
    }

    #[test]
    fn test_replay_aware_logger_errors_only() {
        use std::sync::atomic::{AtomicUsize, Ordering};

        let debug_count = Arc::new(AtomicUsize::new(0));
        let info_count = Arc::new(AtomicUsize::new(0));
        let warn_count = Arc::new(AtomicUsize::new(0));
        let error_count = Arc::new(AtomicUsize::new(0));

        let inner = Arc::new(custom_logger(
            {
                let count = debug_count.clone();
                move |_, _| {
                    count.fetch_add(1, Ordering::SeqCst);
                }
            },
            {
                let count = info_count.clone();
                move |_, _| {
                    count.fetch_add(1, Ordering::SeqCst);
                }
            },
            {
                let count = warn_count.clone();
                move |_, _| {
                    count.fetch_add(1, Ordering::SeqCst);
                }
            },
            {
                let count = error_count.clone();
                move |_, _| {
                    count.fetch_add(1, Ordering::SeqCst);
                }
            },
        ));

        let logger = ReplayAwareLogger::new(inner, ReplayLoggingConfig::ErrorsOnly);

        // During replay, only errors should pass through
        let replay_info = LogInfo::new("arn:test").with_replay(true);
        logger.debug("test", &replay_info);
        logger.info("test", &replay_info);
        logger.warn("test", &replay_info);
        logger.error("test", &replay_info);

        assert_eq!(debug_count.load(Ordering::SeqCst), 0);
        assert_eq!(info_count.load(Ordering::SeqCst), 0);
        assert_eq!(warn_count.load(Ordering::SeqCst), 0);
        assert_eq!(error_count.load(Ordering::SeqCst), 1);
    }

    #[test]
    fn test_replay_aware_logger_warnings_and_errors() {
        use std::sync::atomic::{AtomicUsize, Ordering};

        let debug_count = Arc::new(AtomicUsize::new(0));
        let info_count = Arc::new(AtomicUsize::new(0));
        let warn_count = Arc::new(AtomicUsize::new(0));
        let error_count = Arc::new(AtomicUsize::new(0));

        let inner = Arc::new(custom_logger(
            {
                let count = debug_count.clone();
                move |_, _| {
                    count.fetch_add(1, Ordering::SeqCst);
                }
            },
            {
                let count = info_count.clone();
                move |_, _| {
                    count.fetch_add(1, Ordering::SeqCst);
                }
            },
            {
                let count = warn_count.clone();
                move |_, _| {
                    count.fetch_add(1, Ordering::SeqCst);
                }
            },
            {
                let count = error_count.clone();
                move |_, _| {
                    count.fetch_add(1, Ordering::SeqCst);
                }
            },
        ));

        let logger = ReplayAwareLogger::new(inner, ReplayLoggingConfig::WarningsAndErrors);

        // During replay, only warnings and errors should pass through
        let replay_info = LogInfo::new("arn:test").with_replay(true);
        logger.debug("test", &replay_info);
        logger.info("test", &replay_info);
        logger.warn("test", &replay_info);
        logger.error("test", &replay_info);

        assert_eq!(debug_count.load(Ordering::SeqCst), 0);
        assert_eq!(info_count.load(Ordering::SeqCst), 0);
        assert_eq!(warn_count.load(Ordering::SeqCst), 1);
        assert_eq!(error_count.load(Ordering::SeqCst), 1);
    }

    #[test]
    fn test_replay_aware_logger_suppress_replay_constructor() {
        let inner: Arc<dyn Logger> = Arc::new(TracingLogger);
        let logger = ReplayAwareLogger::suppress_replay(inner);

        assert_eq!(logger.config(), ReplayLoggingConfig::SuppressAll);
    }

    #[test]
    fn test_replay_aware_logger_debug() {
        let inner: Arc<dyn Logger> = Arc::new(TracingLogger);
        let logger = ReplayAwareLogger::new(inner, ReplayLoggingConfig::SuppressAll);

        let debug_str = format!("{:?}", logger);
        assert!(debug_str.contains("ReplayAwareLogger"));
        assert!(debug_str.contains("SuppressAll"));
    }

    #[test]
    fn test_durable_context_create_log_info_with_replay_method() {
        let state = create_test_state();
        let ctx = DurableContext::new(state);

        let info = ctx.create_log_info_with_replay("op-123", true);

        assert_eq!(info.operation_id, Some("op-123".to_string()));
        assert!(info.is_replay);
    }

    #[test]
    fn test_durable_context_clone() {
        let state = create_test_state();
        let ctx = DurableContext::new(state);

        ctx.next_operation_id();
        ctx.next_operation_id();

        let cloned = ctx.clone();

        assert_eq!(cloned.durable_execution_arn(), ctx.durable_execution_arn());
        assert_eq!(cloned.current_step_counter(), ctx.current_step_counter());
    }

    #[test]
    fn test_durable_context_debug() {
        let state = create_test_state();
        let ctx = DurableContext::new(state);

        let debug_str = format!("{:?}", ctx);

        assert!(debug_str.contains("DurableContext"));
        assert!(debug_str.contains("durable_execution_arn"));
    }

    #[test]
    fn test_durable_context_state_access() {
        let state = create_test_state();
        let ctx = DurableContext::new(state.clone());

        // Verify we can access the state
        assert!(Arc::ptr_eq(&ctx.state(), &state));
    }

    #[test]
    fn test_durable_context_send_sync() {
        // This test verifies at compile time that DurableContext is Send + Sync
        fn assert_send_sync<T: Send + Sync>() {}
        assert_send_sync::<DurableContext>();
    }

    // ========================================================================
    // Simplified Logging API Tests
    // ========================================================================

    #[test]
    fn test_log_info_method() {
        use std::sync::atomic::{AtomicUsize, Ordering};
        use std::sync::Mutex;

        let info_count = Arc::new(AtomicUsize::new(0));
        let captured_info = Arc::new(Mutex::new(None::<LogInfo>));

        let captured_info_clone = captured_info.clone();
        let inner = Arc::new(custom_logger(
            |_, _| {},
            {
                let count = info_count.clone();
                move |_, info: &LogInfo| {
                    count.fetch_add(1, Ordering::SeqCst);
                    *captured_info_clone.lock().unwrap() = Some(info.clone());
                }
            },
            |_, _| {},
            |_, _| {},
        ));

        let state = create_test_state();
        let ctx = DurableContext::new(state).with_logger(inner);

        ctx.log_info("Test message");

        assert_eq!(info_count.load(Ordering::SeqCst), 1);

        let captured = captured_info.lock().unwrap();
        let info = captured.as_ref().unwrap();
        assert_eq!(
            info.durable_execution_arn,
            Some("arn:aws:lambda:us-east-1:123456789012:function:test:durable:abc123".to_string())
        );
    }

    #[test]
    fn test_log_debug_method() {
        use std::sync::atomic::{AtomicUsize, Ordering};

        let debug_count = Arc::new(AtomicUsize::new(0));

        let inner = Arc::new(custom_logger(
            {
                let count = debug_count.clone();
                move |_, _| {
                    count.fetch_add(1, Ordering::SeqCst);
                }
            },
            |_, _| {},
            |_, _| {},
            |_, _| {},
        ));

        let state = create_test_state();
        let ctx = DurableContext::new(state).with_logger(inner);

        ctx.log_debug("Debug message");

        assert_eq!(debug_count.load(Ordering::SeqCst), 1);
    }

    #[test]
    fn test_log_warn_method() {
        use std::sync::atomic::{AtomicUsize, Ordering};

        let warn_count = Arc::new(AtomicUsize::new(0));

        let inner = Arc::new(custom_logger(
            |_, _| {},
            |_, _| {},
            {
                let count = warn_count.clone();
                move |_, _| {
                    count.fetch_add(1, Ordering::SeqCst);
                }
            },
            |_, _| {},
        ));

        let state = create_test_state();
        let ctx = DurableContext::new(state).with_logger(inner);

        ctx.log_warn("Warning message");

        assert_eq!(warn_count.load(Ordering::SeqCst), 1);
    }

    #[test]
    fn test_log_error_method() {
        use std::sync::atomic::{AtomicUsize, Ordering};

        let error_count = Arc::new(AtomicUsize::new(0));

        let inner = Arc::new(custom_logger(|_, _| {}, |_, _| {}, |_, _| {}, {
            let count = error_count.clone();
            move |_, _| {
                count.fetch_add(1, Ordering::SeqCst);
            }
        }));

        let state = create_test_state();
        let ctx = DurableContext::new(state).with_logger(inner);

        ctx.log_error("Error message");

        assert_eq!(error_count.load(Ordering::SeqCst), 1);
    }

    #[test]
    fn test_log_info_with_extra_fields() {
        use std::sync::Mutex;

        let captured_info = Arc::new(Mutex::new(None::<LogInfo>));

        let captured_info_clone = captured_info.clone();
        let inner = Arc::new(custom_logger(
            |_, _| {},
            move |_, info: &LogInfo| {
                *captured_info_clone.lock().unwrap() = Some(info.clone());
            },
            |_, _| {},
            |_, _| {},
        ));

        let state = create_test_state();
        let ctx = DurableContext::new(state).with_logger(inner);

        ctx.log_info_with("Test message", &[("order_id", "123"), ("amount", "99.99")]);

        let captured = captured_info.lock().unwrap();
        let info = captured.as_ref().unwrap();

        // Verify extra fields are present
        assert_eq!(info.extra.len(), 2);
        assert!(info
            .extra
            .contains(&("order_id".to_string(), "123".to_string())));
        assert!(info
            .extra
            .contains(&("amount".to_string(), "99.99".to_string())));
    }

    #[test]
    fn test_log_debug_with_extra_fields() {
        use std::sync::Mutex;

        let captured_info = Arc::new(Mutex::new(None::<LogInfo>));

        let captured_info_clone = captured_info.clone();
        let inner = Arc::new(custom_logger(
            move |_, info: &LogInfo| {
                *captured_info_clone.lock().unwrap() = Some(info.clone());
            },
            |_, _| {},
            |_, _| {},
            |_, _| {},
        ));

        let state = create_test_state();
        let ctx = DurableContext::new(state).with_logger(inner);

        ctx.log_debug_with("Debug message", &[("key", "value")]);

        let captured = captured_info.lock().unwrap();
        let info = captured.as_ref().unwrap();

        assert_eq!(info.extra.len(), 1);
        assert!(info
            .extra
            .contains(&("key".to_string(), "value".to_string())));
    }

    #[test]
    fn test_log_warn_with_extra_fields() {
        use std::sync::Mutex;

        let captured_info = Arc::new(Mutex::new(None::<LogInfo>));

        let captured_info_clone = captured_info.clone();
        let inner = Arc::new(custom_logger(
            |_, _| {},
            |_, _| {},
            move |_, info: &LogInfo| {
                *captured_info_clone.lock().unwrap() = Some(info.clone());
            },
            |_, _| {},
        ));

        let state = create_test_state();
        let ctx = DurableContext::new(state).with_logger(inner);

        ctx.log_warn_with("Warning message", &[("retry", "3")]);

        let captured = captured_info.lock().unwrap();
        let info = captured.as_ref().unwrap();

        assert_eq!(info.extra.len(), 1);
        assert!(info.extra.contains(&("retry".to_string(), "3".to_string())));
    }

    #[test]
    fn test_log_error_with_extra_fields() {
        use std::sync::Mutex;

        let captured_info = Arc::new(Mutex::new(None::<LogInfo>));

        let captured_info_clone = captured_info.clone();
        let inner = Arc::new(custom_logger(
            |_, _| {},
            |_, _| {},
            |_, _| {},
            move |_, info: &LogInfo| {
                *captured_info_clone.lock().unwrap() = Some(info.clone());
            },
        ));

        let state = create_test_state();
        let ctx = DurableContext::new(state).with_logger(inner);

        ctx.log_error_with(
            "Error message",
            &[("error_code", "E001"), ("details", "Something went wrong")],
        );

        let captured = captured_info.lock().unwrap();
        let info = captured.as_ref().unwrap();

        assert_eq!(info.extra.len(), 2);
        assert!(info
            .extra
            .contains(&("error_code".to_string(), "E001".to_string())));
        assert!(info
            .extra
            .contains(&("details".to_string(), "Something went wrong".to_string())));
    }

    #[test]
    fn test_log_methods_include_parent_id_in_child_context() {
        use std::sync::Mutex;

        let captured_info = Arc::new(Mutex::new(None::<LogInfo>));

        let captured_info_clone = captured_info.clone();
        let inner: Arc<dyn Logger> = Arc::new(custom_logger(
            |_, _| {},
            move |_, info: &LogInfo| {
                *captured_info_clone.lock().unwrap() = Some(info.clone());
            },
            |_, _| {},
            |_, _| {},
        ));

        let state = create_test_state();
        let ctx = DurableContext::new(state).with_logger(inner.clone());

        let parent_op_id = ctx.next_operation_id();
        let child_ctx = ctx.create_child_context(&parent_op_id).with_logger(inner);

        child_ctx.log_info("Child context message");

        let captured = captured_info.lock().unwrap();
        let info = captured.as_ref().unwrap();

        // Verify parent_id is included
        assert_eq!(info.parent_id, Some(parent_op_id));
    }

    // ========================================================================
    // Runtime Logger Reconfiguration Tests (Req 13.1, 13.2)
    // ========================================================================

    #[test]
    fn test_configure_logger_swaps_logger() {
        // Req 13.1: configure_logger swaps the logger, subsequent calls use new logger
        use std::sync::atomic::{AtomicUsize, Ordering};

        let original_count = Arc::new(AtomicUsize::new(0));
        let new_count = Arc::new(AtomicUsize::new(0));

        let original_logger: Arc<dyn Logger> = Arc::new(custom_logger(
            |_, _| {},
            {
                let count = original_count.clone();
                move |_, _| {
                    count.fetch_add(1, Ordering::SeqCst);
                }
            },
            |_, _| {},
            |_, _| {},
        ));

        let new_logger: Arc<dyn Logger> = Arc::new(custom_logger(
            |_, _| {},
            {
                let count = new_count.clone();
                move |_, _| {
                    count.fetch_add(1, Ordering::SeqCst);
                }
            },
            |_, _| {},
            |_, _| {},
        ));

        let state = create_test_state();
        let ctx = DurableContext::new(state).with_logger(original_logger);

        // Log with original logger
        ctx.log_info("before swap");
        assert_eq!(original_count.load(Ordering::SeqCst), 1);
        assert_eq!(new_count.load(Ordering::SeqCst), 0);

        // Swap logger at runtime (note: &self, not &mut self)
        ctx.configure_logger(new_logger);

        // Subsequent calls use the new logger
        ctx.log_info("after swap");
        assert_eq!(original_count.load(Ordering::SeqCst), 1); // unchanged
        assert_eq!(new_count.load(Ordering::SeqCst), 1);
    }

    #[test]
    fn test_original_logger_used_when_configure_logger_not_called() {
        // Req 13.2: Original logger used when configure_logger not called
        use std::sync::atomic::{AtomicUsize, Ordering};

        let original_count = Arc::new(AtomicUsize::new(0));

        let original_logger: Arc<dyn Logger> = Arc::new(custom_logger(
            |_, _| {},
            {
                let count = original_count.clone();
                move |_, _| {
                    count.fetch_add(1, Ordering::SeqCst);
                }
            },
            |_, _| {},
            |_, _| {},
        ));

        let state = create_test_state();
        let ctx = DurableContext::new(state).with_logger(original_logger);

        // Log multiple times without calling configure_logger
        ctx.log_info("message 1");
        ctx.log_info("message 2");
        ctx.log_info("message 3");

        assert_eq!(original_count.load(Ordering::SeqCst), 3);
    }

    #[test]
    fn test_configure_logger_affects_child_contexts() {
        // Verify that child contexts sharing the same RwLock see the swapped logger
        use std::sync::atomic::{AtomicUsize, Ordering};

        let original_count = Arc::new(AtomicUsize::new(0));
        let new_count = Arc::new(AtomicUsize::new(0));

        let original_logger: Arc<dyn Logger> = Arc::new(custom_logger(
            |_, _| {},
            {
                let count = original_count.clone();
                move |_, _| {
                    count.fetch_add(1, Ordering::SeqCst);
                }
            },
            |_, _| {},
            |_, _| {},
        ));

        let new_logger: Arc<dyn Logger> = Arc::new(custom_logger(
            |_, _| {},
            {
                let count = new_count.clone();
                move |_, _| {
                    count.fetch_add(1, Ordering::SeqCst);
                }
            },
            |_, _| {},
            |_, _| {},
        ));

        let state = create_test_state();
        let ctx = DurableContext::new(state).with_logger(original_logger);
        let parent_op_id = ctx.next_operation_id();
        let child_ctx = ctx.create_child_context(&parent_op_id);

        // Child uses original logger
        child_ctx.log_info("child before swap");
        assert_eq!(original_count.load(Ordering::SeqCst), 1);

        // Swap on parent
        ctx.configure_logger(new_logger);

        // Child now uses the new logger (shared RwLock)
        child_ctx.log_info("child after swap");
        assert_eq!(new_count.load(Ordering::SeqCst), 1);
        assert_eq!(original_count.load(Ordering::SeqCst), 1); // unchanged
    }
}

#[cfg(test)]
mod property_tests {
    use super::*;
    use proptest::prelude::*;

    // Property 3: Operation ID Determinism
    // *For any* DurableContext with a given parent_id and step counter state,
    // calling create_operation_id() SHALL produce the same ID when called
    // in the same sequence position across multiple executions.
    // **Validates: Requirements 1.10**
    proptest! {
        #![proptest_config(ProptestConfig::with_cases(100))]

        /// Feature: durable-execution-rust-sdk, Property 3: Operation ID Determinism
        /// Validates: Requirements 1.10
        #[test]
        fn prop_operation_id_determinism(
            base_id in "[a-zA-Z0-9:/-]{1,100}",
            counter in 0u64..10000u64,
        ) {
            // Generate the same ID twice with the same inputs
            let id1 = generate_operation_id(&base_id, counter);
            let id2 = generate_operation_id(&base_id, counter);

            // IDs must be identical for the same inputs
            prop_assert_eq!(&id1, &id2, "Same base_id and counter must produce identical IDs");

            // ID must be a valid hex string of expected length
            prop_assert_eq!(id1.len(), 32, "ID must be 32 hex characters");
            prop_assert!(id1.chars().all(|c| c.is_ascii_hexdigit()), "ID must be valid hex");
        }

        /// Feature: durable-execution-rust-sdk, Property 3: Operation ID Determinism (Generator)
        /// Validates: Requirements 1.10
        #[test]
        fn prop_operation_id_generator_determinism(
            base_id in "[a-zA-Z0-9:/-]{1,100}",
            num_ids in 1usize..50usize,
        ) {
            // Create two generators with the same base ID
            let gen1 = OperationIdGenerator::new(&base_id);
            let gen2 = OperationIdGenerator::new(&base_id);

            // Generate the same sequence of IDs from both
            let ids1: Vec<String> = (0..num_ids).map(|_| gen1.next_id()).collect();
            let ids2: Vec<String> = (0..num_ids).map(|_| gen2.next_id()).collect();

            // Sequences must be identical
            prop_assert_eq!(&ids1, &ids2, "Same generator sequence must produce identical IDs");

            // All IDs in a sequence must be unique
            let unique_count = {
                let mut set = std::collections::HashSet::new();
                for id in &ids1 {
                    set.insert(id.clone());
                }
                set.len()
            };
            prop_assert_eq!(unique_count, num_ids, "All IDs in sequence must be unique");
        }

        /// Feature: durable-execution-rust-sdk, Property 3: Operation ID Determinism (Replay Simulation)
        /// Validates: Requirements 1.10
        #[test]
        fn prop_operation_id_replay_determinism(
            base_id in "[a-zA-Z0-9:/-]{1,100}",
            operations in prop::collection::vec(0u32..3u32, 1..20),
        ) {
            // Simulate two executions with the same sequence of operations
            // Each operation type increments the counter

            let gen1 = OperationIdGenerator::new(&base_id);
            let gen2 = OperationIdGenerator::new(&base_id);

            let mut ids1 = Vec::new();
            let mut ids2 = Vec::new();

            // First "execution"
            for op_type in &operations {
                // Each operation generates an ID
                ids1.push(gen1.next_id());

                // Some operations might generate additional child IDs
                if *op_type == 2 {
                    let parent_id = ids1.last().unwrap().clone();
                    let child_gen = gen1.create_child(parent_id);
                    ids1.push(child_gen.next_id());
                }
            }

            // Second "execution" (replay) - must produce same IDs
            for op_type in &operations {
                ids2.push(gen2.next_id());

                if *op_type == 2 {
                    let parent_id = ids2.last().unwrap().clone();
                    let child_gen = gen2.create_child(parent_id);
                    ids2.push(child_gen.next_id());
                }
            }

            // Both executions must produce identical ID sequences
            prop_assert_eq!(&ids1, &ids2, "Replay must produce identical operation IDs");
        }
    }

    // Property 5: Concurrent ID Generation Uniqueness
    // *For any* number of concurrent tasks generating operation IDs from the same
    // DurableContext, all generated IDs SHALL be unique.
    // **Validates: Requirements 17.3**
    mod concurrent_id_tests {
        use super::*;
        use std::sync::Arc;
        use std::thread;

        proptest! {
            #![proptest_config(ProptestConfig::with_cases(100))]

            /// Feature: durable-execution-rust-sdk, Property 5: Concurrent ID Generation Uniqueness
            /// Validates: Requirements 17.3
            #[test]
            fn prop_concurrent_id_uniqueness(
                base_id in "[a-zA-Z0-9:/-]{1,100}",
                num_threads in 2usize..10usize,
                ids_per_thread in 10usize..100usize,
            ) {
                let gen = Arc::new(OperationIdGenerator::new(&base_id));
                let mut handles = vec![];

                // Spawn multiple threads that concurrently generate IDs
                for _ in 0..num_threads {
                    let gen_clone = gen.clone();
                    let count = ids_per_thread;
                    handles.push(thread::spawn(move || {
                        let mut ids = Vec::with_capacity(count);
                        for _ in 0..count {
                            ids.push(gen_clone.next_id());
                        }
                        ids
                    }));
                }

                // Collect all IDs from all threads
                let mut all_ids = Vec::new();
                for handle in handles {
                    all_ids.extend(handle.join().unwrap());
                }

                let total_expected = num_threads * ids_per_thread;

                // Verify we got the expected number of IDs
                prop_assert_eq!(all_ids.len(), total_expected, "Should have generated {} IDs", total_expected);

                // Verify all IDs are unique
                let unique_count = {
                    let mut set = std::collections::HashSet::new();
                    for id in &all_ids {
                        set.insert(id.clone());
                    }
                    set.len()
                };

                prop_assert_eq!(
                    unique_count,
                    total_expected,
                    "All {} IDs must be unique, but only {} were unique",
                    total_expected,
                    unique_count
                );

                // Verify the counter was incremented correctly
                prop_assert_eq!(
                    gen.current_counter() as usize,
                    total_expected,
                    "Counter should equal total IDs generated"
                );
            }

            /// Feature: durable-execution-rust-sdk, Property 5: Concurrent ID Generation Uniqueness (Stress)
            /// Validates: Requirements 17.3
            #[test]
            fn prop_concurrent_id_uniqueness_stress(
                base_id in "[a-zA-Z0-9:/-]{1,50}",
            ) {
                // Fixed high-concurrency stress test
                let num_threads = 20;
                let ids_per_thread = 500;

                let gen = Arc::new(OperationIdGenerator::new(&base_id));
                let mut handles = vec![];

                for _ in 0..num_threads {
                    let gen_clone = gen.clone();
                    handles.push(thread::spawn(move || {
                        let mut ids = Vec::with_capacity(ids_per_thread);
                        for _ in 0..ids_per_thread {
                            ids.push(gen_clone.next_id());
                        }
                        ids
                    }));
                }

                let mut all_ids = Vec::new();
                for handle in handles {
                    all_ids.extend(handle.join().unwrap());
                }

                let total_expected = num_threads * ids_per_thread;

                // Verify all IDs are unique
                let unique_count = {
                    let mut set = std::collections::HashSet::new();
                    for id in &all_ids {
                        set.insert(id.clone());
                    }
                    set.len()
                };

                prop_assert_eq!(
                    unique_count,
                    total_expected,
                    "All {} IDs must be unique under high concurrency",
                    total_expected
                );
            }
        }
    }

    // Property 9: Logging Methods Automatic Context
    // *For any* call to log_info, log_debug, log_warn, or log_error on DurableContext,
    // the resulting log message SHALL include durable_execution_arn and parent_id
    // (when available) without the caller needing to specify them.
    // **Validates: Requirements 4.5**
    mod logging_automatic_context_tests {
        use super::*;
        use std::sync::{Arc, Mutex};

        proptest! {
            #![proptest_config(ProptestConfig::with_cases(100))]

            /// Feature: sdk-ergonomics-improvements, Property 9: Logging Methods Automatic Context
            /// Validates: Requirements 4.5
            #[test]
            fn prop_logging_automatic_context(
                message in "[a-zA-Z0-9 ]{1,100}",
                log_level in 0u8..4u8,
            ) {
                use crate::client::MockDurableServiceClient;
                use crate::lambda::InitialExecutionState;

                let captured_info = Arc::new(Mutex::new(None::<LogInfo>));
                let captured_info_clone = captured_info.clone();

                // Create a logger that captures the LogInfo
                let inner = Arc::new(custom_logger(
                    {
                        let captured = captured_info_clone.clone();
                        move |_, info: &LogInfo| {
                            *captured.lock().unwrap() = Some(info.clone());
                        }
                    },
                    {
                        let captured = captured_info_clone.clone();
                        move |_, info: &LogInfo| {
                            *captured.lock().unwrap() = Some(info.clone());
                        }
                    },
                    {
                        let captured = captured_info_clone.clone();
                        move |_, info: &LogInfo| {
                            *captured.lock().unwrap() = Some(info.clone());
                        }
                    },
                    {
                        let captured = captured_info_clone.clone();
                        move |_, info: &LogInfo| {
                            *captured.lock().unwrap() = Some(info.clone());
                        }
                    },
                ));

                let client: crate::client::SharedDurableServiceClient = Arc::new(MockDurableServiceClient::new());
                let state = Arc::new(ExecutionState::new(
                    "arn:aws:lambda:us-east-1:123456789012:function:test:durable:abc123",
                    "token-123",
                    InitialExecutionState::new(),
                    client,
                ));
                let ctx = DurableContext::new(state).with_logger(inner);

                // Call the appropriate log method based on log_level
                match log_level {
                    0 => ctx.log_debug(&message),
                    1 => ctx.log_info(&message),
                    2 => ctx.log_warn(&message),
                    _ => ctx.log_error(&message),
                }

                // Verify the captured LogInfo has the automatic context
                let captured = captured_info.lock().unwrap();
                let info = captured.as_ref().expect("LogInfo should be captured");

                // durable_execution_arn must always be present
                prop_assert!(
                    info.durable_execution_arn.is_some(),
                    "durable_execution_arn must be automatically included"
                );
                prop_assert_eq!(
                    info.durable_execution_arn.as_ref().unwrap(),
                    "arn:aws:lambda:us-east-1:123456789012:function:test:durable:abc123",
                    "durable_execution_arn must match the context's ARN"
                );
            }

            /// Feature: sdk-ergonomics-improvements, Property 9: Logging Methods Automatic Context (Child Context)
            /// Validates: Requirements 4.5
            #[test]
            fn prop_logging_automatic_context_child(
                message in "[a-zA-Z0-9 ]{1,100}",
                log_level in 0u8..4u8,
            ) {
                use crate::client::MockDurableServiceClient;
                use crate::lambda::InitialExecutionState;

                let captured_info = Arc::new(Mutex::new(None::<LogInfo>));
                let captured_info_clone = captured_info.clone();

                // Create a logger that captures the LogInfo
                let inner: Arc<dyn Logger> = Arc::new(custom_logger(
                    {
                        let captured = captured_info_clone.clone();
                        move |_, info: &LogInfo| {
                            *captured.lock().unwrap() = Some(info.clone());
                        }
                    },
                    {
                        let captured = captured_info_clone.clone();
                        move |_, info: &LogInfo| {
                            *captured.lock().unwrap() = Some(info.clone());
                        }
                    },
                    {
                        let captured = captured_info_clone.clone();
                        move |_, info: &LogInfo| {
                            *captured.lock().unwrap() = Some(info.clone());
                        }
                    },
                    {
                        let captured = captured_info_clone.clone();
                        move |_, info: &LogInfo| {
                            *captured.lock().unwrap() = Some(info.clone());
                        }
                    },
                ));

                let client: crate::client::SharedDurableServiceClient = Arc::new(MockDurableServiceClient::new());
                let state = Arc::new(ExecutionState::new(
                    "arn:aws:lambda:us-east-1:123456789012:function:test:durable:abc123",
                    "token-123",
                    InitialExecutionState::new(),
                    client,
                ));
                let ctx = DurableContext::new(state).with_logger(inner.clone());

                // Create a child context
                let parent_op_id = ctx.next_operation_id();
                let child_ctx = ctx.create_child_context(&parent_op_id).with_logger(inner);

                // Call the appropriate log method based on log_level
                match log_level {
                    0 => child_ctx.log_debug(&message),
                    1 => child_ctx.log_info(&message),
                    2 => child_ctx.log_warn(&message),
                    _ => child_ctx.log_error(&message),
                }

                // Verify the captured LogInfo has the automatic context including parent_id
                let captured = captured_info.lock().unwrap();
                let info = captured.as_ref().expect("LogInfo should be captured");

                // durable_execution_arn must always be present
                prop_assert!(
                    info.durable_execution_arn.is_some(),
                    "durable_execution_arn must be automatically included in child context"
                );

                // parent_id must be present in child context
                prop_assert!(
                    info.parent_id.is_some(),
                    "parent_id must be automatically included in child context"
                );
                prop_assert_eq!(
                    info.parent_id.as_ref().unwrap(),
                    &parent_op_id,
                    "parent_id must match the parent operation ID"
                );
            }
        }
    }

    // Property 10: Logging Methods Extra Fields
    // *For any* call to log_info_with, log_debug_with, log_warn_with, or log_error_with
    // with extra fields, those fields SHALL appear in the resulting log output.
    // **Validates: Requirements 4.6**
    mod logging_extra_fields_tests {
        use super::*;
        use std::sync::{Arc, Mutex};

        proptest! {
            #![proptest_config(ProptestConfig::with_cases(100))]

            /// Feature: sdk-ergonomics-improvements, Property 10: Logging Methods Extra Fields
            /// Validates: Requirements 4.6
            #[test]
            fn prop_logging_extra_fields(
                message in "[a-zA-Z0-9 ]{1,100}",
                log_level in 0u8..4u8,
                key1 in "[a-zA-Z_][a-zA-Z0-9_]{0,20}",
                value1 in "[a-zA-Z0-9]{1,50}",
                key2 in "[a-zA-Z_][a-zA-Z0-9_]{0,20}",
                value2 in "[a-zA-Z0-9]{1,50}",
            ) {
                use crate::client::MockDurableServiceClient;
                use crate::lambda::InitialExecutionState;

                let captured_info = Arc::new(Mutex::new(None::<LogInfo>));
                let captured_info_clone = captured_info.clone();

                // Create a logger that captures the LogInfo
                let inner = Arc::new(custom_logger(
                    {
                        let captured = captured_info_clone.clone();
                        move |_, info: &LogInfo| {
                            *captured.lock().unwrap() = Some(info.clone());
                        }
                    },
                    {
                        let captured = captured_info_clone.clone();
                        move |_, info: &LogInfo| {
                            *captured.lock().unwrap() = Some(info.clone());
                        }
                    },
                    {
                        let captured = captured_info_clone.clone();
                        move |_, info: &LogInfo| {
                            *captured.lock().unwrap() = Some(info.clone());
                        }
                    },
                    {
                        let captured = captured_info_clone.clone();
                        move |_, info: &LogInfo| {
                            *captured.lock().unwrap() = Some(info.clone());
                        }
                    },
                ));

                let client: crate::client::SharedDurableServiceClient = Arc::new(MockDurableServiceClient::new());
                let state = Arc::new(ExecutionState::new(
                    "arn:aws:lambda:us-east-1:123456789012:function:test:durable:abc123",
                    "token-123",
                    InitialExecutionState::new(),
                    client,
                ));
                let ctx = DurableContext::new(state).with_logger(inner);

                // Create extra fields
                let fields: Vec<(&str, &str)> = vec![(&key1, &value1), (&key2, &value2)];

                // Call the appropriate log method based on log_level
                match log_level {
                    0 => ctx.log_debug_with(&message, &fields),
                    1 => ctx.log_info_with(&message, &fields),
                    2 => ctx.log_warn_with(&message, &fields),
                    _ => ctx.log_error_with(&message, &fields),
                }

                // Verify the captured LogInfo has the extra fields
                let captured = captured_info.lock().unwrap();
                let info = captured.as_ref().expect("LogInfo should be captured");

                // Extra fields must be present
                prop_assert_eq!(
                    info.extra.len(),
                    2,
                    "Extra fields must be included in the log output"
                );

                // Verify each field is present
                prop_assert!(
                    info.extra.contains(&(key1.clone(), value1.clone())),
                    "First extra field must be present: {}={}", key1, value1
                );
                prop_assert!(
                    info.extra.contains(&(key2.clone(), value2.clone())),
                    "Second extra field must be present: {}={}", key2, value2
                );
            }

            /// Feature: sdk-ergonomics-improvements, Property 10: Logging Methods Extra Fields (Empty)
            /// Validates: Requirements 4.6
            #[test]
            fn prop_logging_extra_fields_empty(
                message in "[a-zA-Z0-9 ]{1,100}",
                log_level in 0u8..4u8,
            ) {
                use crate::client::MockDurableServiceClient;
                use crate::lambda::InitialExecutionState;

                let captured_info = Arc::new(Mutex::new(None::<LogInfo>));
                let captured_info_clone = captured_info.clone();

                // Create a logger that captures the LogInfo
                let inner = Arc::new(custom_logger(
                    {
                        let captured = captured_info_clone.clone();
                        move |_, info: &LogInfo| {
                            *captured.lock().unwrap() = Some(info.clone());
                        }
                    },
                    {
                        let captured = captured_info_clone.clone();
                        move |_, info: &LogInfo| {
                            *captured.lock().unwrap() = Some(info.clone());
                        }
                    },
                    {
                        let captured = captured_info_clone.clone();
                        move |_, info: &LogInfo| {
                            *captured.lock().unwrap() = Some(info.clone());
                        }
                    },
                    {
                        let captured = captured_info_clone.clone();
                        move |_, info: &LogInfo| {
                            *captured.lock().unwrap() = Some(info.clone());
                        }
                    },
                ));

                let client: crate::client::SharedDurableServiceClient = Arc::new(MockDurableServiceClient::new());
                let state = Arc::new(ExecutionState::new(
                    "arn:aws:lambda:us-east-1:123456789012:function:test:durable:abc123",
                    "token-123",
                    InitialExecutionState::new(),
                    client,
                ));
                let ctx = DurableContext::new(state).with_logger(inner);

                // Call the appropriate log method with empty fields
                let empty_fields: &[(&str, &str)] = &[];
                match log_level {
                    0 => ctx.log_debug_with(&message, empty_fields),
                    1 => ctx.log_info_with(&message, empty_fields),
                    2 => ctx.log_warn_with(&message, empty_fields),
                    _ => ctx.log_error_with(&message, empty_fields),
                }

                // Verify the captured LogInfo has no extra fields
                let captured = captured_info.lock().unwrap();
                let info = captured.as_ref().expect("LogInfo should be captured");

                prop_assert!(
                    info.extra.is_empty(),
                    "Extra fields should be empty when none are provided"
                );

                // But automatic context should still be present
                prop_assert!(
                    info.durable_execution_arn.is_some(),
                    "durable_execution_arn must still be present even with empty extra fields"
                );
            }
        }
    }
}

#[cfg(test)]
mod sealed_trait_tests {
    use super::*;
    use std::sync::atomic::{AtomicUsize, Ordering};

    /// Tests for sealed Logger trait implementations.
    ///
    /// The Logger trait is sealed, meaning it cannot be implemented outside this crate.
    /// This is enforced at compile time by requiring the private `Sealed` supertrait.
    /// External crates attempting to implement Logger will get a compile error:
    /// "the trait bound `MyType: Sealed` is not satisfied"
    ///
    /// These tests verify that the internal implementations work correctly.
    mod logger_tests {
        use super::*;

        #[test]
        fn test_tracing_logger_implements_logger() {
            // TracingLogger should implement Logger (compile-time check)
            let logger: &dyn Logger = &TracingLogger;
            let info = LogInfo::default();

            // These calls should not panic
            logger.debug("test debug", &info);
            logger.info("test info", &info);
            logger.warn("test warn", &info);
            logger.error("test error", &info);
        }

        #[test]
        fn test_replay_aware_logger_implements_logger() {
            // ReplayAwareLogger should implement Logger (compile-time check)
            let inner = Arc::new(TracingLogger);
            let logger = ReplayAwareLogger::new(inner, ReplayLoggingConfig::AllowAll);
            let logger_ref: &dyn Logger = &logger;
            let info = LogInfo::default();

            // These calls should not panic
            logger_ref.debug("test debug", &info);
            logger_ref.info("test info", &info);
            logger_ref.warn("test warn", &info);
            logger_ref.error("test error", &info);
        }

        #[test]
        fn test_custom_logger_implements_logger() {
            // CustomLogger should implement Logger (compile-time check)
            let call_count = Arc::new(AtomicUsize::new(0));
            let count_clone = call_count.clone();

            let logger = custom_logger(
                {
                    let count = count_clone.clone();
                    move |_msg, _info| {
                        count.fetch_add(1, Ordering::SeqCst);
                    }
                },
                {
                    let count = count_clone.clone();
                    move |_msg, _info| {
                        count.fetch_add(1, Ordering::SeqCst);
                    }
                },
                {
                    let count = count_clone.clone();
                    move |_msg, _info| {
                        count.fetch_add(1, Ordering::SeqCst);
                    }
                },
                {
                    let count = count_clone.clone();
                    move |_msg, _info| {
                        count.fetch_add(1, Ordering::SeqCst);
                    }
                },
            );

            let logger_ref: &dyn Logger = &logger;
            let info = LogInfo::default();

            logger_ref.debug("test", &info);
            logger_ref.info("test", &info);
            logger_ref.warn("test", &info);
            logger_ref.error("test", &info);

            assert_eq!(call_count.load(Ordering::SeqCst), 4);
        }

        #[test]
        fn test_simple_custom_logger() {
            let call_count = Arc::new(AtomicUsize::new(0));
            let count_clone = call_count.clone();

            let logger = simple_custom_logger(move |_level, _msg, _info| {
                count_clone.fetch_add(1, Ordering::SeqCst);
            });

            let info = LogInfo::default();

            logger.debug("test", &info);
            logger.info("test", &info);
            logger.warn("test", &info);
            logger.error("test", &info);

            assert_eq!(call_count.load(Ordering::SeqCst), 4);
        }

        #[test]
        fn test_custom_logger_receives_correct_messages() {
            let messages = Arc::new(std::sync::Mutex::new(Vec::new()));
            let messages_clone = messages.clone();

            let logger = simple_custom_logger(move |level, msg, _info| {
                messages_clone
                    .lock()
                    .unwrap()
                    .push(format!("[{}] {}", level, msg));
            });

            let info = LogInfo::default();

            logger.debug("debug message", &info);
            logger.info("info message", &info);
            logger.warn("warn message", &info);
            logger.error("error message", &info);

            let logged = messages.lock().unwrap();
            assert_eq!(logged.len(), 4);
            assert_eq!(logged[0], "[DEBUG] debug message");
            assert_eq!(logged[1], "[INFO] info message");
            assert_eq!(logged[2], "[WARN] warn message");
            assert_eq!(logged[3], "[ERROR] error message");
        }

        #[test]
        fn test_custom_logger_receives_log_info() {
            let received_info = Arc::new(std::sync::Mutex::new(None));
            let info_clone = received_info.clone();

            let logger = simple_custom_logger(move |_level, _msg, info| {
                *info_clone.lock().unwrap() = Some(info.clone());
            });

            let info = LogInfo::new("arn:aws:test")
                .with_operation_id("op-123")
                .with_parent_id("parent-456")
                .with_replay(true);

            logger.info("test", &info);

            let received = received_info.lock().unwrap().clone().unwrap();
            assert_eq!(
                received.durable_execution_arn,
                Some("arn:aws:test".to_string())
            );
            assert_eq!(received.operation_id, Some("op-123".to_string()));
            assert_eq!(received.parent_id, Some("parent-456".to_string()));
            assert!(received.is_replay);
        }

        #[test]
        fn test_replay_aware_logger_suppresses_during_replay() {
            let call_count = Arc::new(AtomicUsize::new(0));
            let count_clone = call_count.clone();

            let inner_logger = Arc::new(custom_logger(
                {
                    let count = count_clone.clone();
                    move |_msg, _info| {
                        count.fetch_add(1, Ordering::SeqCst);
                    }
                },
                {
                    let count = count_clone.clone();
                    move |_msg, _info| {
                        count.fetch_add(1, Ordering::SeqCst);
                    }
                },
                {
                    let count = count_clone.clone();
                    move |_msg, _info| {
                        count.fetch_add(1, Ordering::SeqCst);
                    }
                },
                {
                    let count = count_clone.clone();
                    move |_msg, _info| {
                        count.fetch_add(1, Ordering::SeqCst);
                    }
                },
            ));

            let logger = ReplayAwareLogger::new(inner_logger, ReplayLoggingConfig::SuppressAll);

            // Non-replay logs should pass through
            let non_replay_info = LogInfo::default().with_replay(false);
            logger.debug("test", &non_replay_info);
            logger.info("test", &non_replay_info);
            logger.warn("test", &non_replay_info);
            logger.error("test", &non_replay_info);
            assert_eq!(call_count.load(Ordering::SeqCst), 4);

            // Replay logs should be suppressed
            let replay_info = LogInfo::default().with_replay(true);
            logger.debug("test", &replay_info);
            logger.info("test", &replay_info);
            logger.warn("test", &replay_info);
            logger.error("test", &replay_info);
            assert_eq!(call_count.load(Ordering::SeqCst), 4); // Still 4, no new calls
        }

        #[test]
        fn test_replay_aware_logger_errors_only_during_replay() {
            let call_count = Arc::new(AtomicUsize::new(0));
            let count_clone = call_count.clone();

            let inner_logger = Arc::new(custom_logger(
                {
                    let count = count_clone.clone();
                    move |_msg, _info| {
                        count.fetch_add(1, Ordering::SeqCst);
                    }
                },
                {
                    let count = count_clone.clone();
                    move |_msg, _info| {
                        count.fetch_add(1, Ordering::SeqCst);
                    }
                },
                {
                    let count = count_clone.clone();
                    move |_msg, _info| {
                        count.fetch_add(1, Ordering::SeqCst);
                    }
                },
                {
                    let count = count_clone.clone();
                    move |_msg, _info| {
                        count.fetch_add(1, Ordering::SeqCst);
                    }
                },
            ));

            let logger = ReplayAwareLogger::new(inner_logger, ReplayLoggingConfig::ErrorsOnly);

            let replay_info = LogInfo::default().with_replay(true);
            logger.debug("test", &replay_info);
            logger.info("test", &replay_info);
            logger.warn("test", &replay_info);
            logger.error("test", &replay_info);

            // Only error should pass through during replay
            assert_eq!(call_count.load(Ordering::SeqCst), 1);
        }
    }
}