rialo_types/

rex_info.rs

// Copyright (c) Subzero Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

//! # REX Module
//!
//! This module contains all core REX-related types and structures.
//!
//! ## REX Identification
//! - [`RexId`] - Uniquely identifies a REX instance using a nonce and creator
//!
//! ## REX Configuration & State
//! - [`RexInfo`] - Complete REX definition and configuration
//! - [`RexEntry`] - REX registry entry with metadata
//!
//! ## REX Values
//! - [`RexValue`] - String value (plain or encrypted)
//! - [`RexValueBody`] - Binary value (plain or encrypted)
//!
//! ## REX Targets & Updates
//! - [`TargetRexProgram`] - Defines what the REX system should query (HTTP, Time, etc.)
//! - [`RexUpdateResult`] - Result of a REX update with signature
//!
//! ## REX Requests & Scheduling
//! - [`RexRequest`] - Request parameters for REX execution
//! - [`UpdateFrequency`] - How often the REX system should run
//! - [`StartingTimestamp`] - When the REX system should start

use std::{
    collections::BTreeMap,
    convert::Infallible,
    fmt,
    ops::Deref,
    str::FromStr,
    sync::Arc,
    time::{Duration, SystemTime, UNIX_EPOCH},
};

use borsh::{BorshDeserialize, BorshSerialize};
#[cfg(feature = "non-pdk")]
use clap::Subcommand;
#[cfg(feature = "non-pdk")]
use fastcrypto::encoding::{Base64, Encoding};
use rialo_cli_representable::Representable;
use rialo_limits::{max_rex_output_serialized_bytes, MIN_VIABLE_LIMIT_OF_REX_OUTPUT_SIZE};
use rialo_s_pubkey::Pubkey;
use serde::{Deserialize, Serialize};
use serde_big_array::BigArray;
#[cfg(feature = "non-pdk")]
use url::Url;

use crate::{
    time::{DurationMs, TimestampMs},
    websocket_op::WebSocketOperation,
    AttestationReport, AuthorityKeyBytes, Headers, HttpFilter, Nonce, RexDutyConfig,
};

/// Lowest allowed update period for periodic REX requests, in milliseconds.
///
/// This is a pragmatic lower bound to avoid excessive scheduling / load.
const MIN_UPDATE_PERIOD_MS: DurationMs = DurationMs::new(50);

// ============================================================================
// REX Identification
// ============================================================================

/// REX identifier that uniquely identifies a REX instance using a nonce and creator.
///
/// # String Parsing
///
/// `RexId` implements `FromStr` which expects a JSON format:
/// ```json
/// {"nonce":"<nonce_value>","creator":"<base58_pubkey>"}
/// ```
///
/// This JSON format is used for CLI parsing and other string-based inputs.
#[derive(
    Debug,
    Default,
    Clone,
    Copy,
    PartialEq,
    Eq,
    Hash,
    PartialOrd,
    Ord,
    Serialize,
    Deserialize,
    BorshSerialize,
    BorshDeserialize,
)]
pub struct RexId {
    pub nonce: Nonce,
    pub creator: Pubkey,
}

impl RexId {
    /// Create a new RexId from a nonce and creator
    pub fn new(creator: Pubkey, nonce: impl Into<Nonce>) -> Self {
        Self {
            nonce: nonce.into(),
            creator,
        }
    }
}

impl fmt::Display for RexId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}:{}", &self.nonce, &self.creator)
    }
}

impl FromStr for RexId {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        // Parse JSON format: {"nonce":"...","creator":"..."}
        serde_json::from_str(s).map_err(|e| format!("Failed to parse RexId: {}", e))
    }
}

// ============================================================================
// REX Configuration & State
// ============================================================================

/// Represents a REX definition and configuration.
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone, Representable)]
#[representable(human_readable = "rex_info_human_readable")]
pub struct RexInfo {
    /// Unique identifier
    pub id: RexId,
    /// Description for the REX instance.
    pub description: String,
    /// When the REX instance should bring back data.
    pub update_frequency: UpdateFrequency,
    /// URLs for the REX instance.
    pub target_rex_programs: Vec<TargetRexProgram>,
    /// Timestamp in which the REX instance will start.
    pub starting_timestamp: StartingTimestamp,
    /// Whether the REX instance is active.
    pub is_active: bool,
    /// When the REX instance was created.
    pub created_at_ms: i64,
    /// Number of validators to assign per REX request.
    #[serde(default = "default_validators_per_duty")]
    pub validators_per_duty: u32,
    /// Delay sufficient for the REX task to complete, in milliseconds.
    #[serde(default = "default_rex_request_delay_ms")]
    pub request_delay_ms: DurationMs,
}

fn rex_info_human_readable(info: &RexInfo) -> String {
    let mut out = String::new();

    out.push_str(&format!("REX ID: {}\n", info.id));
    out.push_str(&format!("Description: {}\n", info.description));
    out.push_str(&format!("Active: {}\n", info.is_active));
    out.push_str(&format!(
        "Starting Timestamp: {:?}\n",
        info.starting_timestamp
    ));
    out.push_str(&format!("Update Frequency: {:?}\n", info.update_frequency));
    out.push_str(&format!("Created At: {}\n", info.created_at_ms));
    out.push_str(&format!(
        "Validators Per Duty: {}\n",
        info.validators_per_duty
    ));
    out.push_str(&format!("REX Request Delay: {}\n", info.request_delay_ms));

    if !info.target_rex_programs.is_empty() {
        out.push_str(&format!(
            "\nTarget REX operations ({}):\n",
            info.target_rex_programs.len()
        ));
        for (i, target) in info.target_rex_programs.iter().enumerate() {
            out.push_str(&format!("  {}. {:?}\n", i + 1, target));
        }
    }

    out
}

impl Default for RexInfo {
    fn default() -> Self {
        Self {
            id: RexId::default(),
            description: String::new(),
            update_frequency: UpdateFrequency::default(),
            target_rex_programs: Vec::new(),
            starting_timestamp: StartingTimestamp::default(),
            is_active: false,
            created_at_ms: 0,
            validators_per_duty: default_validators_per_duty(),
            request_delay_ms: default_rex_request_delay_ms(),
        }
    }
}

impl RexInfo {
    /// Returns true if the REX should start as soon as possible (ASAP),
    /// rather than at a specific timestamp.
    pub fn is_asap(&self) -> bool {
        matches!(self.starting_timestamp, StartingTimestamp::Asap)
    }

    pub fn target_timestamp(&self) -> Option<TimestampMs> {
        match self.starting_timestamp {
            StartingTimestamp::Timestamp(timestamp) => Some(timestamp),
            StartingTimestamp::Asap => None,
        }
    }

    /// True if any target program carries a DKG-encrypted payload, so this REX
    /// needs the threshold-decryption committee + combine pipeline before any
    /// result can be produced. Such duties get an extended collection window
    /// (`DKG_ENCRYPTED_PIPELINE_OVERHEAD_MS + request_delay_ms`) rather than the
    /// plain ASAP timeout. See [`TargetRexProgram::is_dkg_encrypted`].
    pub fn requires_dkg_decryption(&self) -> bool {
        self.target_rex_programs
            .iter()
            .any(TargetRexProgram::is_dkg_encrypted)
    }

    /// Validates all REX configuration fields and their consistency.
    pub fn validate(&self) -> Result<(), String> {
        match self.starting_timestamp {
            StartingTimestamp::Asap => {
                if !matches!(self.update_frequency, UpdateFrequency::OneShot) {
                    return Err("ASAP REX requests cannot be periodic".to_string());
                }
            }
            StartingTimestamp::Timestamp(starting_timestamp) => {
                match self.update_frequency {
                    UpdateFrequency::OneShot => {}
                    UpdateFrequency::Periodic(period)
                    | UpdateFrequency::LimitedPeriodic(period, _) => {
                        validate_periodic_frequency(period)?;

                        // Additional validation for LimitedPeriodic
                        if let UpdateFrequency::LimitedPeriodic(_, end_timestamp) =
                            self.update_frequency
                        {
                            if starting_timestamp >= end_timestamp {
                                return Err("end_timestamp of a LimitedPeriodic REX should be above starting_timestamp".to_string());
                            }
                        }
                    }
                }
            }
        }

        // Validate `target_rex_programs`.
        if self.target_rex_programs.is_empty() {
            return Err("RexTargets cannot be empty".to_string());
        }

        // Validate `rex_request_delay`.
        if self.request_delay_ms < RexDutyConfig::MIN_REX_REQUEST_DELAY {
            return Err(format!(
                "rex_request_delay cannot be below {}",
                RexDutyConfig::MIN_REX_REQUEST_DELAY
            ));
        }
        if self.request_delay_ms > RexDutyConfig::MAX_REX_REQUEST_DELAY_MS {
            return Err(format!(
                "rex_request_delay cannot be above {}",
                RexDutyConfig::MAX_REX_REQUEST_DELAY_MS
            ));
        }

        // Validate `validators_per_duty`.
        if self.validators_per_duty == 0 {
            return Err("validators_per_duty cannot be 0".to_string());
        }
        let max_rex_output_size = max_rex_output_serialized_bytes(self.validators_per_duty);
        if max_rex_output_size < MIN_VIABLE_LIMIT_OF_REX_OUTPUT_SIZE {
            return Err(format!("validators_per_duty is too high, results in max size of REX updates that is too low: {max_rex_output_size} vs {MIN_VIABLE_LIMIT_OF_REX_OUTPUT_SIZE}"));
        }

        Ok(())
    }

    /// Extracts the WebSocket operation from the REX targets, if any.
    pub fn websocket_op(&self) -> Option<WebSocketOperation> {
        self.target_rex_programs
            .first()
            .and_then(|target| target.websocket_op())
    }
}

fn validate_periodic_frequency(period_ms: DurationMs) -> Result<(), String> {
    if period_ms == DurationMs::ZERO {
        return Err("update frequency cannot be zero".to_string());
    }

    if period_ms < MIN_UPDATE_PERIOD_MS {
        return Err(format!(
            "update frequency {period_ms} cannot be below {MIN_UPDATE_PERIOD_MS}"
        ));
    }

    Ok(())
}

impl TargetRexProgram {
    /// Extracts the WebSocket operation if this is a WebSocket target.
    pub fn websocket_op(&self) -> Option<WebSocketOperation> {
        if let TargetRexProgram::WebSocket(ws_op) = self {
            Some(ws_op.clone())
        } else {
            None
        }
    }

    /// Ciphertext bytes of every field on this program that can carry an
    /// encrypted payload: the encrypted URL, body, WebSocket `Send` message(s),
    /// or WASM input(s). Single source of truth for *which fields are ciphertext
    /// carriers*; the version-gated routing (`Router::dkg_payload_bytes`) and
    /// the on-chain ciphertext-match check (`encrypted_payload_iter`) build on
    /// this and additionally require the `DKG_PAYLOAD_VERSION` byte (a constant
    /// that lives above this crate, so it is applied by those callers).
    pub fn encrypted_payloads(&self) -> Box<dyn Iterator<Item = &[u8]> + '_> {
        fn enc(v: &RexValue) -> Option<&[u8]> {
            v.is_encrypted().then(|| v.as_bytes())
        }
        // Matched exhaustively (no `_`) on purpose: this is the single source
        // for which program fields can carry ciphertext, so a newly-added
        // `TargetRexProgram` variant must fail to compile here until it is
        // classified — otherwise it would silently bypass DKG detection in both
        // the off-chain router and the on-chain ciphertext-match check.
        // (`#[non_exhaustive]` only forces a wildcard in *other* crates; this is
        // the defining crate, so the guardrail holds.)
        match self {
            TargetRexProgram::HttpGet { url, .. } if url.is_encrypted() => {
                Box::new(std::iter::once(url.as_bytes()))
            }
            TargetRexProgram::HttpGet { .. } => Box::new(std::iter::empty()),
            TargetRexProgram::HttpPost { body, .. } => Box::new(enc(body).into_iter()),
            TargetRexProgram::WebSocket(WebSocketOperation::Send { messages, .. }) => {
                Box::new(messages.iter().filter_map(enc))
            }
            TargetRexProgram::WebSocket(_) => Box::new(std::iter::empty()),
            TargetRexProgram::Wasm { input, .. } => Box::new(input.iter().filter_map(enc)),
            TargetRexProgram::Time
            | TargetRexProgram::Number
            | TargetRexProgram::SecretKeyGeneration { .. }
            | TargetRexProgram::SecretKeyEncryption { .. }
            | TargetRexProgram::SecretKeyDecryption { .. } => Box::new(std::iter::empty()),
        }
    }

    /// True if this program carries any encrypted payload, meaning the duty
    /// cannot be served by a direct fetch and must go through the
    /// threshold-decryption committee + combine pipeline before any result can
    /// exist. Used to size the encrypted duty's collection window
    /// (`DutyRequest::inner`); does not inspect `DKG_PAYLOAD_VERSION` —
    /// encryption alone is sufficient (and conservative) for that purpose.
    pub fn is_dkg_encrypted(&self) -> bool {
        self.encrypted_payloads().next().is_some()
    }
}

fn default_validators_per_duty() -> u32 {
    RexDutyConfig::DEFAULT_VALIDATORS_PER_DUTY
}

fn default_rex_request_delay_ms() -> DurationMs {
    RexDutyConfig::DEFAULT_REQUEST_DELAY_MS
}

/// Represents an REX registry entry containing REX information and metadata.
///
/// This struct stores REX data along with a hash of the data for change detection
/// and tracking information about when the entry was last modified.
#[derive(Debug, Eq, PartialEq, Clone, Serialize, Deserialize)]
pub struct RexEntry {
    rex_info: Arc<RexInfo>,
    data_hash: [u8; RexEntry::HASH_LENGTH],
    last_modified_timestamp: u64,
}

impl RexEntry {
    const HASH_LENGTH: usize = 32;

    /// Creates a new RexEntry instance.
    ///
    /// # Arguments
    ///
    /// * `data`: The account data as a byte vector.
    /// * `data_hash`: The hash of the account data.
    /// * `last_modified_timestamp`: The round in which the account was last modified.
    pub fn new(
        rex_info: RexInfo,
        data_hash: [u8; Self::HASH_LENGTH],
        last_modified_round: u64,
    ) -> Self {
        Self {
            rex_info: Arc::new(rex_info),
            data_hash,
            last_modified_timestamp: last_modified_round,
        }
    }

    /// Retrieves the account data.
    ///
    /// # Returns
    ///
    /// The account data as an RexInfo.
    pub fn rex_info(&self) -> Arc<RexInfo> {
        self.rex_info.clone()
    }

    /// Retrieves the last modified round.
    ///
    /// # Returns
    ///
    /// The last modified round as a `u64`.
    pub fn last_modified_timestamp(&self) -> u64 {
        self.last_modified_timestamp
    }

    /// Retrieves the hash of the account data.
    ///
    /// # Returns
    ///
    /// The hash of the account data as a byte array.
    pub fn data_hash(&self) -> &[u8; Self::HASH_LENGTH] {
        &self.data_hash
    }
}

// ============================================================================
// REX Values
// ============================================================================

/// Deserializes a `Vec<u8>` from either a JSON string or a JSON byte array.
///
/// This allows CLI users and scripts to pass `{"Plain": "hello"}` instead of
/// the canonical `{"Plain": [104, 101, 108, 108, 111]}` byte-array format.
/// Both formats are accepted; strings are converted via UTF-8 `.as_bytes()`.
///
/// Uses `deserialize_bytes` rather than `deserialize_any` for compatibility
/// with non-self-describing binary formats (e.g., bincode used for on-chain data).
/// For JSON, `serde_json`'s `deserialize_bytes` handles both strings (via
/// `visit_str`) and arrays (via `visit_seq`) directly. For bincode, it reads
/// raw bytes via `visit_bytes`.
fn deserialize_bytes_or_string<'de, D>(deserializer: D) -> Result<Vec<u8>, D::Error>
where
    D: serde::Deserializer<'de>,
{
    struct BytesOrString;

    impl<'de> serde::de::Visitor<'de> for BytesOrString {
        type Value = Vec<u8>;

        fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            f.write_str("a byte array, byte slice, or a string")
        }

        fn visit_bytes<E: serde::de::Error>(self, bytes: &[u8]) -> Result<Vec<u8>, E> {
            Ok(bytes.to_vec())
        }

        fn visit_byte_buf<E: serde::de::Error>(self, bytes: Vec<u8>) -> Result<Vec<u8>, E> {
            Ok(bytes)
        }

        fn visit_str<E: serde::de::Error>(self, s: &str) -> Result<Vec<u8>, E> {
            Ok(s.as_bytes().to_vec())
        }

        fn visit_seq<A: serde::de::SeqAccess<'de>>(self, mut seq: A) -> Result<Vec<u8>, A::Error> {
            let mut bytes = Vec::with_capacity(seq.size_hint().unwrap_or(0));
            while let Some(b) = seq.next_element()? {
                bytes.push(b);
            }
            Ok(bytes)
        }
    }

    deserializer.deserialize_bytes(BytesOrString)
}

/// Represents a value that can be either plain text or encrypted.
///
/// This enum is used to handle REX data that may contain sensitive information
/// that needs to be encrypted when transmitted or stored, while also supporting
/// plain text values for non-sensitive data.
///
/// Both variants accept JSON strings or byte arrays during deserialization.
/// Serialization always outputs the canonical byte-array format.
///
/// # Variants
/// * `Plain(Vec<u8>)` - A plain text value that is not encrypted.
///   Accepts a JSON string (e.g., `{"Plain": "hello"}`) or byte array.
/// * `Encrypted(Vec<u8>)` - Encrypted ciphertext. When passed as a JSON string,
///   the value should be base64-encoded ciphertext for TEE decryption.
///   No base64 validation is performed at deserialization time.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
pub enum RexValue {
    Plain(#[serde(deserialize_with = "deserialize_bytes_or_string")] Vec<u8>),
    Encrypted(#[serde(deserialize_with = "deserialize_bytes_or_string")] Vec<u8>),
}

impl RexValue {
    /// Create a plain value from raw bytes.
    pub fn plain(data: Vec<u8>) -> Self {
        RexValue::Plain(data)
    }

    /// Create a plain value from a string (convenience method).
    ///
    /// This is the preferred way to create an `RexValue` from string data
    /// after the migration from string-based `RexValue`.
    pub fn plain_string(s: impl Into<String>) -> Self {
        RexValue::Plain(s.into().into_bytes())
    }

    /// Create an encrypted value from raw ciphertext bytes.
    pub fn encrypted(ciphertext: Vec<u8>) -> Self {
        RexValue::Encrypted(ciphertext)
    }

    /// Get the inner data as a string slice, if it's valid UTF-8.
    ///
    /// Returns `Some(&str)` if the data is valid UTF-8, `None` otherwise.
    /// Only works on `Plain` variants; `Encrypted` always returns `None`.
    pub fn as_string(&self) -> Option<&str> {
        match self {
            RexValue::Plain(bytes) => std::str::from_utf8(bytes).ok(),
            RexValue::Encrypted(_) => None,
        }
    }

    /// Get the inner data as a byte slice.
    ///
    /// Works on both `Plain` and `Encrypted` variants.
    pub fn as_bytes(&self) -> &[u8] {
        match self {
            RexValue::Plain(bytes) | RexValue::Encrypted(bytes) => bytes,
        }
    }

    /// Returns `true` if this is a `Plain` variant.
    pub fn is_plain(&self) -> bool {
        matches!(self, RexValue::Plain(_))
    }

    /// Returns `true` if this is an `Encrypted` variant.
    pub fn is_encrypted(&self) -> bool {
        matches!(self, RexValue::Encrypted(_))
    }
}

impl Default for RexValue {
    fn default() -> Self {
        RexValue::Plain(vec![])
    }
}

/// Convert an argument into a [`RexValue`] bound to the WASM rex function's
/// expected parameter type `T`.
///
/// Parameterizing the trait by the *target* type lets two non-overlapping
/// impls do compile-time dispatch:
///
/// - `impl<T: BorshSerialize> IntoRexValueFor<T> for T` — plain values are
///   borsh-serialized and emitted as `RexValue::Plain`.
/// - `impl<T> IntoRexValueFor<T> for EncryptedInput<T>` — encrypted values
///   forward their ciphertext as `RexValue::Encrypted`.
///
/// The two impls cover distinct `(Self, T)` pairs (`(u64, u64)` vs
/// `(EncryptedInput<u64>, u64)`), so they don't overlap and no specialization
/// is required. Venus codegen emits `IntoRexValueFor::<#param_ty>` for each
/// positional argument, so type mismatches (e.g. `EncryptedInput<String>`
/// passed to a `u64` rex parameter) surface as ordinary `rustc` errors from
/// the trait resolver rather than requiring the macro to run its own
/// precheck.
pub trait IntoRexValueFor<T> {
    fn into_rex_value_for(self) -> RexValue;
}

impl<T: BorshSerialize> IntoRexValueFor<T> for T {
    fn into_rex_value_for(self) -> RexValue {
        RexValue::Plain(borsh::to_vec(&self).expect("borsh serialize failed"))
    }
}

impl<T> IntoRexValueFor<T> for EncryptedInput<T> {
    fn into_rex_value_for(self) -> RexValue {
        RexValue::Encrypted(self.into_bytes())
    }
}

impl FromStr for RexValue {
    type Err = Infallible;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        Ok(RexValue::Plain(s.as_bytes().to_vec()))
    }
}

/// Type alias for backward compatibility.
///
/// `RexValueBody` has been consolidated into `RexValue`.
/// This alias is provided for migration but will be removed in a future version.
#[deprecated(since = "0.2.0", note = "Use RexValue instead")]
pub type RexValueBody = RexValue;

/// A wrapper type around `RexValue` specifically for handling URLs in REX configurations.
///
/// This type provides a convenient way to handle both plain text and encrypted URLs:
/// - Plain text URLs are stored directly as strings
/// - Encrypted URLs are stored as base64-encoded encrypted strings prefixed with "enc://"
///
/// The type implements common traits like Display and FromStr for easy conversion and
/// formatting, and integrates with the url crate when the "non-pdk" feature is enabled.
///
/// # Examples
///
/// ```
/// use std::str::FromStr;
/// use rialo_types::RexUrl;
///
/// // Create from plain text URL
/// let plain_url = RexUrl::from("https://example.com");
///
/// // Create from encrypted URL
/// let encrypted_url = RexUrl::from("enc://encrypted_data");
/// ```
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
pub struct RexUrl(RexValue);

impl fmt::Display for RexUrl {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match &self.0 {
            RexValue::Plain(bytes) => {
                // Try to display as UTF-8 string, fall back to hex format
                match std::str::from_utf8(bytes) {
                    Ok(s) => write!(f, "{}", s),
                    #[cfg(feature = "non-pdk")]
                    Err(_) => write!(f, "<binary:{}>", Base64::encode(bytes)),
                    #[cfg(not(feature = "non-pdk"))]
                    Err(_) => write!(f, "<binary:{}>", hex::encode(bytes)),
                }
            }
            RexValue::Encrypted(bytes) => {
                // Try to display as UTF-8 string, fall back to hex format
                match std::str::from_utf8(bytes) {
                    Ok(s) => write!(f, "enc://{}", s),
                    #[cfg(feature = "non-pdk")]
                    Err(_) => write!(f, "enc://<binary:{}>", Base64::encode(bytes)),
                    #[cfg(not(feature = "non-pdk"))]
                    Err(_) => write!(f, "enc://<binary:{}>", hex::encode(bytes)),
                }
            }
        }
    }
}

impl Deref for RexUrl {
    type Target = RexValue;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

#[cfg(feature = "non-pdk")]
impl From<Url> for RexUrl {
    fn from(url: Url) -> Self {
        url.to_string().into()
    }
}

#[cfg(feature = "non-pdk")]
impl From<&Url> for RexUrl {
    fn from(url: &Url) -> Self {
        Self(RexValue::Plain(url.to_string().into_bytes()))
    }
}

impl From<String> for RexUrl {
    fn from(url: String) -> Self {
        url.as_str().into()
    }
}

impl From<&str> for RexUrl {
    fn from(s: &str) -> Self {
        if let Some(encrypted) = s.strip_prefix("enc://") {
            Self(RexValue::Encrypted(encrypted.into()))
        } else {
            // If it isn't prefixed with `enc://`, treat it as plain text
            Self(RexValue::Plain(s.into()))
        }
    }
}

impl FromStr for RexUrl {
    type Err = Infallible;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        Ok(s.into())
    }
}

impl From<RexValue> for RexUrl {
    fn from(value: RexValue) -> Self {
        Self(value)
    }
}

// ============================================================================
// REX Targets
// ============================================================================

/// Enum that represents the target REX Program of the REX request.
///
/// Discriminant values are fixed to ensure stable serialization across builds
/// with different feature flags.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, strum_macros::AsRefStr)]
#[cfg_attr(feature = "non-pdk", derive(Subcommand))]
#[repr(u16)]
#[non_exhaustive]
pub enum TargetRexProgram {
    /// HTTP GET request to the REX. The URL is specified in the RexDutyRequest.
    /// The URL must be an HTTPS URL.
    /// The filter is optional, and if provided, it is used to filter the response.
    HttpGet {
        #[cfg_attr(feature = "non-pdk", clap(
            long = "target-url",
            value_parser = clap::value_parser!(RexUrl)
        ))]
        url: RexUrl,
        #[cfg_attr(feature = "non-pdk", clap(long, default_value = None))]
        filter: Option<Vec<HttpFilter>>,
        #[cfg_attr(feature = "non-pdk", clap(
            long,
            default_value_t = Headers::default(),
            action = clap::ArgAction::Append
        ))]
        headers: Headers,
    } = 0,
    /// HTTP POST request to the REX. The URL is specified in the RexDutyRequest.
    /// The URL must be an HTTPS URL. This REX type is restricted to single-validator assignment
    /// to prevent duplicate operations on non-idempotent endpoints.
    HttpPost {
        #[cfg_attr(feature = "non-pdk", clap(
            long = "target-url",
            value_parser = clap::value_parser!(RexUrl)
        ))]
        url: RexUrl,
        #[cfg_attr(feature = "non-pdk", clap(long))]
        filter: Option<Vec<HttpFilter>>,
        #[cfg_attr(feature = "non-pdk", clap(
            long,
            value_parser = clap::value_parser!(RexValue)
        ))]
        body: RexValue,
        #[cfg_attr(feature = "non-pdk", clap(long))]
        content_type: String,
        #[cfg_attr(feature = "non-pdk", clap(
            long,
            default_value_t = Headers::default(),
            action = clap::ArgAction::Append
        ))]
        headers: Headers,
    } = 1,
    /// Get the current time from several NIST servers.
    Time = 2,
    /// Only used for testing purposes, to simulate an REX that always returns a fixed value.
    /// TODO: remove this variant with a cfg testing flag.
    Number = 3,
    /// Generate a shared secret key within a committee of TEEs.
    /// The manager TEE generates the key and distributes it to all committee members.
    SecretKeyGeneration {
        #[cfg_attr(feature = "non-pdk", clap(long))]
        committee_id: String,
        #[cfg_attr(feature = "non-pdk", clap(long))]
        committee_members: Vec<String>,
    } = 4,
    /// Encrypt a secret key for a target TEE using their public key.
    /// This is used to share keys with TEEs outside the original committee.
    SecretKeyEncryption {
        #[cfg_attr(feature = "non-pdk", clap(long))]
        target_tee_id: String,
        #[cfg_attr(feature = "non-pdk", clap(long))]
        secret_data: Vec<u8>,
        #[cfg_attr(feature = "non-pdk", clap(long))]
        committee_id: String,
    } = 5,
    /// Decrypt a secret key that was encrypted for this TEE.
    /// This is used by TEEs to access keys shared with them.
    SecretKeyDecryption {
        #[cfg_attr(feature = "non-pdk", clap(long))]
        encrypted_data: Vec<u8>,
        #[cfg_attr(feature = "non-pdk", clap(long))]
        source_committee_id: String,
    } = 6,
    /// WebSocket REX operations for persistent connections
    #[cfg_attr(feature = "non-pdk", clap(subcommand))]
    WebSocket(WebSocketOperation) = 7,
    /// Execute WASM bytecode inside the TEE.
    /// The bytecode must be a WASI-compatible component deployed to an on-chain account.
    Wasm {
        /// Account pubkey containing the deployed WASM component(s)
        #[cfg_attr(feature = "non-pdk", clap(long))]
        bytecode_account: Pubkey,
        /// Per-argument input data for the WASM module.
        /// Each entry is one argument, independently Plain (borsh-serialized)
        /// or Encrypted (HPKE ciphertext the TEE decrypts). The TEE decrypts
        /// each Encrypted entry and concatenates all bytes into a borsh tuple.
        /// From CLI: pass base64-encoded bytes (single plain argument).
        #[cfg_attr(feature = "non-pdk", clap(long, value_parser = parse_base64_rex_value))]
        input: Vec<RexValue>,
        /// Index into the program table within the account (defaults to 0)
        #[cfg_attr(feature = "non-pdk", clap(long))]
        program_index: Option<u32>,
    } = 8,
}

/// Pre-encrypted input bytes for TEE-decrypted REX arguments.
///
/// On-chain this is opaque ciphertext (`Vec<u8>`). The TEE decrypts it
/// transparently before passing the plaintext to the REX handler (e.g.,
/// a WASM component). Venus codegen recognizes this type and emits
/// `RexValue::Encrypted(...)` instead of borsh-serializing individual args.
///
/// The type parameter `T` is a phantom type that Venus codegen uses to
/// verify the encrypted value matches the expected function parameter type
/// at compile time. It has no effect on serialization — `EncryptedInput<u64>`
/// and `EncryptedInput<String>` serialize identically as opaque bytes.
///
/// Create with `rialo_cdk::encrypt_input` on the client side.
///
/// # Examples
///
/// ```ignore
/// // Typed: Venus verifies EncryptedInput<u64> matches `amount: u64`
/// state { encrypted_amount: EncryptedInput<u64> }
///
/// // Untyped: no compile-time check (T defaults to ())
/// state { encrypted_blob: EncryptedInput }
/// ```
pub struct EncryptedInput<T = ()> {
    ciphertext: Vec<u8>,
    _phantom: std::marker::PhantomData<T>,
}

// Manual derives because #[derive] adds bounds on T that we don't want.
// EncryptedInput<T> should be Clone/Default/etc regardless of T.

impl<T> Clone for EncryptedInput<T> {
    fn clone(&self) -> Self {
        Self {
            ciphertext: self.ciphertext.clone(),
            _phantom: std::marker::PhantomData,
        }
    }
}

impl<T> Default for EncryptedInput<T> {
    fn default() -> Self {
        Self {
            ciphertext: Vec::new(),
            _phantom: std::marker::PhantomData,
        }
    }
}

impl<T> std::fmt::Debug for EncryptedInput<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_tuple("EncryptedInput")
            .field(&format!("[{} bytes]", self.ciphertext.len()))
            .finish()
    }
}

impl<T> PartialEq for EncryptedInput<T> {
    fn eq(&self, other: &Self) -> bool {
        self.ciphertext == other.ciphertext
    }
}

impl<T> Eq for EncryptedInput<T> {}

impl<T> BorshSerialize for EncryptedInput<T> {
    fn serialize<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
        borsh::BorshSerialize::serialize(&self.ciphertext, writer)
    }
}

impl<T> BorshDeserialize for EncryptedInput<T> {
    fn deserialize_reader<R: std::io::Read>(reader: &mut R) -> std::io::Result<Self> {
        Ok(Self {
            ciphertext: borsh::BorshDeserialize::deserialize_reader(reader)?,
            _phantom: std::marker::PhantomData,
        })
    }
}

impl<T> Serialize for EncryptedInput<T> {
    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        serde::Serialize::serialize(&self.ciphertext, serializer)
    }
}

impl<'de, T> Deserialize<'de> for EncryptedInput<T> {
    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        Ok(Self {
            ciphertext: <Vec<u8> as serde::Deserialize>::deserialize(deserializer)?,
            _phantom: std::marker::PhantomData,
        })
    }
}

impl<T> EncryptedInput<T> {
    pub fn new(ciphertext: Vec<u8>) -> Self {
        Self {
            ciphertext,
            _phantom: std::marker::PhantomData,
        }
    }

    pub fn as_bytes(&self) -> &[u8] {
        &self.ciphertext
    }

    pub fn into_bytes(self) -> Vec<u8> {
        self.ciphertext
    }
}

impl<T> From<Vec<u8>> for EncryptedInput<T> {
    fn from(v: Vec<u8>) -> Self {
        Self::new(v)
    }
}

/// Parse CLI input as base64-encoded bytes into a plain RexValue for the WASM `--input` flag.
#[cfg(feature = "non-pdk")]
fn parse_base64_rex_value(s: &str) -> Result<RexValue, String> {
    use fastcrypto::encoding::{Base64, Encoding};
    Base64::decode(s)
        .map(RexValue::Plain)
        .map_err(|e| format!("Invalid base64 input: {e}"))
}

impl TargetRexProgram {
    /// Returns true if this is a WebSocket operation.
    pub fn is_websocket(&self) -> bool {
        matches!(self, TargetRexProgram::WebSocket(_))
    }
}

impl FromStr for TargetRexProgram {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s == "Time" {
            Ok(TargetRexProgram::Time)
        } else if s == "SecretKeyGeneration" {
            Err("SecretKeyGeneration REX requires committee_id and committee_members parameters. Use the appropriate API to create this REX type.".to_string())
        } else if s == "SecretKeyEncryption" {
            Err("SecretKeyEncryption REX requires target_tee_id, secret_data, and committee_id parameters. Use the appropriate API to create this REX type.".to_string())
        } else if s == "SecretKeyDecryption" {
            Err("SecretKeyDecryption REX requires encrypted_data and source_committee_id parameters. Use the appropriate API to create this REX type.".to_string())
        } else if s == "number" {
            Err("The 'number' REX is only for testing purposes and should not be used in production.".to_string())
        } else {
            if let Some(rest) = s.strip_prefix("HttpGet:") {
                let parts: Vec<&str> = rest.splitn(2, '|').collect();
                if parts.is_empty() {
                    return Err(
                        "Invalid HttpGet format. Use 'HttpGet:<url>[|<filter>]'.".to_string()
                    );
                }

                let url = parts[0].to_string();
                let filter = if parts.len() > 1 && !parts[1].is_empty() {
                    Some(vec![HttpFilter::from_str(parts[1])?])
                } else {
                    None
                };

                // Validate URL (only when url crate is available)
                #[cfg(feature = "non-pdk")]
                if Url::parse(&url).is_err() {
                    return Err(format!("Invalid URL: {url}"));
                }

                return Ok(TargetRexProgram::HttpGet {
                    url: url.into(),
                    filter,
                    headers: Headers::default(),
                });
            }

            Err(format!("Unknown TargetRexProgram type: {s}"))
        }
    }
}

// ============================================================================
// REX Updates
// ============================================================================

/// 32‑byte Blake3 hash of the raw request payload observed by the REX service.
///
/// This binds a response to the exact input it was computed from and is included
/// in the signature transcript as `input_commitment || blake3(response_value)`.
pub type InputCommitmentBytes = [u8; 32];

/// Raw 64‑byte Ed25519 signature over the REX response transcript.
///
/// The transcript signed by the REX is `input_commitment || blake3(response_value)`.
pub type SignatureBytes = [u8; 64];

/// A structure representing the result of an REX update
/// This structure is designed to fit within Solana's transaction size limit
/// TODO: <https://linear.app/subzero-labs/issue/SUB-449/audit-transaction-sizes-in-the-REX-subsystem>
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct RexUpdateResult {
    /// The identifier of the REX (should be kept under 100 bytes)
    pub rex_id: RexId,

    /// The round in which this result should be proposed
    pub target_timestamp: TimestampMs,

    /// Blake3 hash of `rex_result` (32 bytes)
    #[serde(with = "BigArray")]
    pub response_hash: [u8; 32],

    /// Blake3 hash commitment of the original request input bytes (32 bytes)
    /// This is included in the signature transcript to bind the response to
    /// the exact request payload observed by the REX service.
    #[serde(with = "BigArray")]
    pub input_commitment: InputCommitmentBytes,

    /// Signature of the hash of the REX response, base64 encoded.
    #[serde(with = "BigArray")]
    pub signature: SignatureBytes,

    /// Response data (up to MAX_TRANSACTION_SIZE bytes in size)
    pub rex_result: Vec<u8>,

    /// Optional attestation report, if available
    /// This can be used to provide additional context or verification of the REX result.
    pub attestation_report: Option<AttestationReport>,

    /// Validator's protocol public key bytes of the validator that executed the edge rquest.
    #[serde(with = "BigArray")]
    pub authority_key: AuthorityKeyBytes,
}

impl RexUpdateResult {
    /// Create a RexUpdateResult
    pub fn new(
        rex_id: RexId,
        target_timestamp: TimestampMs,
        rex_result: Vec<u8>,
        input_commitment: InputCommitmentBytes,
        signature: SignatureBytes,
        attestation_report: Option<AttestationReport>,
        authority_key: AuthorityKeyBytes,
    ) -> Result<Self, &'static str> {
        let hash = blake3::hash(&rex_result);

        // This is a temporary solution to avoid having to deal with the size of the REX result
        #[cfg(feature = "non-pdk")]
        let rex_result = if rex_result.len() > rialo_limits::MAX_TRANSACTION_SIZE as usize {
            tracing::error!(
                "REX result size {} exceeds maximum size {}, dropping the result.",
                rex_result.len(),
                rialo_limits::MAX_TRANSACTION_SIZE
            );
            return Err("REX result exceeds maximum size");
        } else {
            // Use the REX result as-is
            rex_result
        };

        Ok(Self {
            rex_id,
            target_timestamp,
            response_hash: *hash.as_bytes(),
            input_commitment,
            signature,
            rex_result,
            attestation_report,
            authority_key,
        })
    }
}

impl Default for RexUpdateResult {
    fn default() -> Self {
        Self {
            rex_id: RexId::default(),
            target_timestamp: TimestampMs::ZERO,
            response_hash: [0; 32],
            input_commitment: [0xee; 32],
            signature: [0; 64],
            rex_result: vec![],
            attestation_report: None,
            authority_key: [0xff; 96],
        }
    }
}

// ============================================================================
// REX Requests & Scheduling
// ============================================================================

/// Represents a request to an REX with parameters that can be used to specify the query.
/// The parameters are stored as a BTreeMap to allow for flexible key-value pairs.
/// Contains fields that uniquely identify and configure the request.
#[derive(BorshSerialize, BorshDeserialize, Debug, PartialEq, Eq)]
pub struct RexRequest {
    /// Structured fields of the request.
    pub rex_id: Option<RexId>,
    pub target_timestamp: Option<TimestampMs>,
    pub authority_key: AuthorityKeyBytes,
    pub include_attestation: bool,
    pub max_output_size: u32,

    /// Request-specific extra data for the request.
    pub params: BTreeMap<String, String>,
}

impl Default for RexRequest {
    fn default() -> Self {
        Self {
            rex_id: None,
            target_timestamp: None,
            authority_key: [0; 96],
            include_attestation: true,
            max_output_size: 0,
            params: BTreeMap::default(),
        }
    }
}

impl RexRequest {
    pub fn input_commitment(&self) -> Result<blake3::Hash, &'static str> {
        let request_bytes = borsh::to_vec(self).map_err(|_| "Failed to serialize RexRequest")?;
        Ok(blake3::hash(&request_bytes))
    }
}

/// The frequency of the REX update.
#[derive(Debug, Default, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum UpdateFrequency {
    /// The REX will update once.
    #[default]
    OneShot,
    /// The REX will update every N milliseconds.
    Periodic(DurationMs),
    /// The REX will update every N milliseconds, but only up to end_timestamp_ms.
    /// First parameter is the frequency in ms, second parameter is the end timestamp in ms.
    LimitedPeriodic(DurationMs, TimestampMs),
}

/// Starting timestamp configuration for REX requests
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone, Copy)]
pub enum StartingTimestamp {
    /// Start at a specific timestamp in milliseconds
    Timestamp(TimestampMs),
    /// Start as soon as possible
    Asap,
}

impl Default for StartingTimestamp {
    fn default() -> Self {
        StartingTimestamp::Timestamp(TimestampMs::ZERO)
    }
}

impl UpdateFrequency {
    /// Creates a periodic update frequency from a [`Duration`].
    pub fn periodic(duration: Duration) -> Self {
        Self::Periodic(DurationMs::new(duration.as_millis() as u64))
    }
}

impl StartingTimestamp {
    pub fn start_offset(offset: Duration) -> Self {
        let timestamp = SystemTime::now() + offset;
        Self::Timestamp(TimestampMs::new(
            timestamp.duration_since(UNIX_EPOCH).unwrap().as_millis() as u64,
        ))
    }
}

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

    fn http_post(body: RexValue) -> TargetRexProgram {
        TargetRexProgram::HttpPost {
            url: "https://example.com".into(),
            filter: None,
            body,
            content_type: "application/json".to_string(),
            headers: Headers::default(),
        }
    }

    fn ws_send(messages: Vec<RexValue>) -> TargetRexProgram {
        TargetRexProgram::WebSocket(WebSocketOperation::Send {
            connection_rex_id: RexId::default(),
            messages,
        })
    }

    #[test]
    fn test_is_dkg_encrypted_false_for_plain_carriers() {
        assert!(!TargetRexProgram::Time.is_dkg_encrypted());
        assert!(!TargetRexProgram::HttpGet {
            url: "https://example.com".into(),
            filter: None,
            headers: Headers::default(),
        }
        .is_dkg_encrypted());
        assert!(!http_post(RexValue::plain(vec![1, 2, 3])).is_dkg_encrypted());
        assert!(!ws_send(vec![RexValue::plain(vec![1])]).is_dkg_encrypted());
        assert!(!TargetRexProgram::Wasm {
            bytecode_account: Pubkey::new_unique(),
            input: vec![RexValue::plain(vec![1]), RexValue::plain(vec![2])],
            program_index: None,
        }
        .is_dkg_encrypted());
    }

    /// `encrypted_payloads` is the single source the router and the on-chain
    /// ciphertext-match check both build on, so it must surface the right bytes
    /// from every carrier — not merely report a bool.
    #[test]
    fn test_encrypted_payloads_yields_ciphertext_from_every_carrier() {
        let get = TargetRexProgram::HttpGet {
            url: RexValue::Encrypted(vec![0x02, 1, 2]).into(),
            filter: None,
            headers: Headers::default(),
        };
        assert_eq!(
            get.encrypted_payloads().collect::<Vec<_>>(),
            vec![&[0x02, 1, 2][..]]
        );
        assert!(get.is_dkg_encrypted());

        let post = http_post(RexValue::Encrypted(vec![0x02, 7]));
        assert_eq!(
            post.encrypted_payloads().collect::<Vec<_>>(),
            vec![&[0x02, 7][..]]
        );
        assert!(post.is_dkg_encrypted());

        // WebSocket Send: only the encrypted message(s), in order.
        let ws = ws_send(vec![
            RexValue::plain(vec![9]),
            RexValue::Encrypted(vec![0x02, 8]),
            RexValue::Encrypted(vec![0x02, 5]),
        ]);
        assert_eq!(
            ws.encrypted_payloads().collect::<Vec<_>>(),
            vec![&[0x02, 8][..], &[0x02, 5][..]]
        );
        assert!(ws.is_dkg_encrypted());

        // WASM: a non-first encrypted arg is still surfaced (preserves order).
        let wasm = TargetRexProgram::Wasm {
            bytecode_account: Pubkey::new_unique(),
            input: vec![RexValue::plain(vec![1]), RexValue::Encrypted(vec![0x02, 9])],
            program_index: None,
        };
        assert_eq!(
            wasm.encrypted_payloads().collect::<Vec<_>>(),
            vec![&[0x02, 9][..]]
        );
        assert!(wasm.is_dkg_encrypted());

        // Plain program yields nothing.
        assert!(TargetRexProgram::Time.encrypted_payloads().next().is_none());
    }

    #[test]
    fn test_requires_dkg_decryption_rolls_up_across_programs() {
        let mut info = RexInfo {
            target_rex_programs: vec![TargetRexProgram::Time],
            ..RexInfo::default()
        };
        assert!(!info.requires_dkg_decryption());
        info.target_rex_programs
            .push(http_post(RexValue::Encrypted(vec![0x02, 1])));
        assert!(info.requires_dkg_decryption());
    }

    fn base_valid_rex_info() -> RexInfo {
        RexInfo {
            description: "test".to_string(),
            target_rex_programs: vec![TargetRexProgram::Time],
            // Default is OneShot, which is allowed for both Asap and Timestamp
            update_frequency: UpdateFrequency::OneShot,
            // Start at timestamp 0 by default
            starting_timestamp: StartingTimestamp::Timestamp(TimestampMs::ZERO),
            // Keep other fields as default
            ..RexInfo::default()
        }
    }

    #[test]
    fn test_is_asap_true_and_false() {
        let mut info = base_valid_rex_info();
        assert!(!info.is_asap());
        info.starting_timestamp = StartingTimestamp::Asap;
        assert!(info.is_asap());
    }

    #[test]
    fn test_validate_success_minimal() {
        let info = base_valid_rex_info();
        assert!(info.validate().is_ok());
    }

    #[test]
    fn test_asap_cannot_be_periodic() {
        let mut info = base_valid_rex_info();
        info.starting_timestamp = StartingTimestamp::Asap;
        info.update_frequency = UpdateFrequency::Periodic(DurationMs::new(10));
        let err = info.validate().unwrap_err();
        assert!(err.contains("ASAP REX requests cannot be periodic"));
    }

    #[test]
    fn test_asap_cannot_be_limited_periodic() {
        let mut info = base_valid_rex_info();
        info.starting_timestamp = StartingTimestamp::Asap;
        info.update_frequency =
            UpdateFrequency::LimitedPeriodic(DurationMs::new(5), TimestampMs::new(100));
        let err = info.validate().unwrap_err();
        assert!(err.contains("ASAP REX requests cannot be periodic"));
    }

    #[test]
    fn test_periodic_with_zero_period_is_invalid() {
        let mut info = base_valid_rex_info();
        info.starting_timestamp = StartingTimestamp::Timestamp(TimestampMs::new(1));
        info.update_frequency = UpdateFrequency::Periodic(DurationMs::ZERO);
        let err = info.validate().unwrap_err();
        assert!(err.contains("update frequency cannot be zero"));
    }

    #[test]
    fn test_limited_periodic_with_zero_period_is_invalid() {
        let mut info = base_valid_rex_info();
        info.starting_timestamp = StartingTimestamp::Timestamp(TimestampMs::new(1));
        info.update_frequency =
            UpdateFrequency::LimitedPeriodic(DurationMs::ZERO, TimestampMs::new(100));
        let err = info.validate().unwrap_err();
        assert!(err.contains("update frequency cannot be zero"));
    }

    #[test]
    fn test_limited_periodic_end_timestamp_must_be_above_starting_timestamp() {
        let mut info = base_valid_rex_info();
        info.starting_timestamp = StartingTimestamp::Timestamp(TimestampMs::new(500));
        info.update_frequency =
            UpdateFrequency::LimitedPeriodic(DurationMs::new(300), TimestampMs::new(500));
        let err = info.validate().unwrap_err();
        assert!(err
            .contains("end_timestamp of a LimitedPeriodic REX should be above starting_timestamp"));
    }

    #[test]
    fn test_limited_periodic_end_timestamp_below_starting_timestamp_is_invalid() {
        let mut info = base_valid_rex_info();
        info.starting_timestamp = StartingTimestamp::Timestamp(TimestampMs::new(1000));
        info.update_frequency =
            UpdateFrequency::LimitedPeriodic(DurationMs::new(300), TimestampMs::new(900));
        let err = info.validate().unwrap_err();
        assert!(err
            .contains("end_timestamp of a LimitedPeriodic REX should be above starting_timestamp"));
    }

    #[test]
    fn test_target_rex_programs_cannot_be_empty() {
        let mut info = base_valid_rex_info();
        info.target_rex_programs.clear();
        let err = info.validate().unwrap_err();
        assert!(err.contains("RexTargets cannot be empty"));
    }

    #[test]
    fn test_rex_request_delay_bounds() {
        // Below minimum
        let mut info = base_valid_rex_info();
        info.request_delay_ms = RexDutyConfig::MIN_REX_REQUEST_DELAY - DurationMs::new(1);
        let err = info.validate().unwrap_err();
        assert!(err.contains(&format!(
            "rex_request_delay cannot be below {}",
            RexDutyConfig::MIN_REX_REQUEST_DELAY
        )));

        // Above maximum
        let mut info = base_valid_rex_info();
        info.request_delay_ms = RexDutyConfig::MAX_REX_REQUEST_DELAY_MS + DurationMs::new(1);
        let err = info.validate().unwrap_err();
        assert!(err.contains(&format!(
            "rex_request_delay cannot be above {}",
            RexDutyConfig::MAX_REX_REQUEST_DELAY_MS
        )));
    }

    #[test]
    fn test_validators_per_duty_cannot_be_zero() {
        let mut info = base_valid_rex_info();
        info.validators_per_duty = 0;
        let err = info.validate().unwrap_err();
        assert!(err.contains("validators_per_duty cannot be 0"));
    }

    #[test]
    fn test_validators_per_duty_too_high_results_in_too_low_output_size() {
        let mut info = base_valid_rex_info();
        info.validators_per_duty = 1_000_000; // ridiculously high
        let err = info.validate().unwrap_err();
        assert!(err.contains("validators_per_duty is too high"));
    }

    #[test]
    fn test_rex_value_plain_deserialize_from_byte_array() {
        let json = r#"{"Plain":[104,101,108,108,111]}"#;
        let value: RexValue = serde_json::from_str(json).unwrap();
        assert_eq!(value, RexValue::Plain(b"hello".to_vec()));
    }

    #[test]
    fn test_rex_value_plain_deserialize_from_string() {
        let json = r#"{"Plain":"hello"}"#;
        let value: RexValue = serde_json::from_str(json).unwrap();
        assert_eq!(value, RexValue::Plain(b"hello".to_vec()));
    }

    #[test]
    fn test_rex_value_encrypted_deserialize_from_string() {
        let json = r#"{"Encrypted":"base64data"}"#;
        let value: RexValue = serde_json::from_str(json).unwrap();
        assert_eq!(value, RexValue::Encrypted(b"base64data".to_vec()));
    }

    #[test]
    fn test_rex_value_encrypted_deserialize_from_byte_array() {
        let json = r#"{"Encrypted":[65,66,67]}"#;
        let value: RexValue = serde_json::from_str(json).unwrap();
        assert_eq!(value, RexValue::Encrypted(b"ABC".to_vec()));
    }

    /// Plain borsh-serializable arg → `RexValue::Plain(borsh_bytes)`.
    #[test]
    fn test_into_rex_value_for_plain() {
        let out = <u64 as IntoRexValueFor<u64>>::into_rex_value_for(1_000_000u64);
        assert_eq!(out, RexValue::Plain(borsh::to_vec(&1_000_000u64).unwrap()));
    }

    /// `EncryptedInput<T>` → `RexValue::Encrypted(ciphertext)` with `T`
    /// as the target rex-fn param type.
    #[test]
    fn test_into_rex_value_for_encrypted() {
        let ct = vec![0xAAu8; 56];
        let wrapped = EncryptedInput::<u64>::new(ct.clone());
        let out = <EncryptedInput<u64> as IntoRexValueFor<u64>>::into_rex_value_for(wrapped);
        assert_eq!(out, RexValue::Encrypted(ct));
    }
}