rialo-types 0.2.0

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

//! # Oracle Module
//!
//! This module contains all core oracle-related types and structures.
//!
//! ## Oracle Identification
//! - [`OracleId`] - Uniquely identifies an oracle using a nonce and creator
//!
//! ## Oracle Configuration & State
//! - [`OracleInfo`] - Complete oracle definition and configuration
//! - [`OracleEntry`] - Oracle registry entry with metadata
//!
//! ## Oracle Values
//! - [`OracleValue`] - String value (plain or encrypted)
//! - [`OracleValueBody`] - Binary value (plain or encrypted)
//!
//! ## Oracle Targets & Updates
//! - [`TargetOracle`] - Defines what the oracle should query (HTTP, Time, etc.)
//! - [`OracleUpdateResult`] - Result of an oracle update with signature
//!
//! ## Oracle Requests & Scheduling
//! - [`OracleRequest`] - Request parameters for oracle execution
//! - [`UpdateFrequency`] - How often the oracle should run
//! - [`StartingTimestamp`] - When the oracle 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;
use rialo_cli_representable::Representable;
use rialo_limits::{max_oracle_output_serialized_bytes, MIN_VIABLE_LIMIT_OF_ORACLE_OUTPUT_SIZE};
use rialo_s_compute_budget::compute_budget_limits::{MAX_COMPUTE_UNIT_LIMIT, MAX_HEAP_FRAME_BYTES};
use rialo_s_pubkey::Pubkey;
use serde::{Deserialize, Serialize};
use serde_big_array::BigArray;
#[cfg(feature = "non-pdk")]
use url::Url;

use crate::{AttestationReport, AuthorityKeyBytes, Headers, HttpFilter, Nonce, OracleDutyConfig};

/// Type alias for timestamp in milliseconds
// TODO: Unify with BlockTimestampMs in fourier.
pub type TimestampMs = u64;

/// Lowest allowed update period for periodic oracles, in milliseconds.
///
/// This is a pragmatic lower bound to avoid excessive scheduling / load.
const MIN_UPDATE_PERIOD_MS: TimestampMs = 50;

// ============================================================================
// Oracle Identification
// ============================================================================

/// Oracle identifier that uniquely identifies an oracle using a nonce and creator.
///
/// # String Parsing
///
/// `OracleId` 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 OracleId {
    pub nonce: Nonce,
    pub creator: Pubkey,
}

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

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

impl FromStr for OracleId {
    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 OracleId: {}", e))
    }
}

/// Parse an OracleId from a string (for clap CLI parsing)
#[cfg(feature = "non-pdk")]
fn parse_oracle_id(s: &str) -> Result<OracleId, String> {
    OracleId::from_str(s)
}

// ============================================================================
// Oracle Configuration & State
// ============================================================================

/// Represents an oracle's definition and configuration
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone, Representable)]
#[representable(human_readable = "oracle_info_human_readable")]
pub struct OracleInfo {
    /// Unique identifier
    pub id: OracleId,
    /// Description for the oracle
    pub description: String,
    /// When the oracle should bring back data
    pub update_frequency: UpdateFrequency,
    /// URLs for the oracle
    pub target_oracles: Vec<TargetOracle>,
    /// Timestamp in which the oracle will start
    pub starting_timestamp: StartingTimestamp,
    /// Whether oracle is active
    pub is_active: bool,
    /// When oracle was created
    pub created_at_ms: i64,
    /// Number of validators to assign per oracle request
    #[serde(default = "default_validators_per_duty")]
    pub validators_per_duty: u32,
    /// Delay sufficient for the oracle to complete, in milliseconds.
    #[serde(default = "default_oracle_request_delay_ms")]
    pub oracle_request_delay_ms: TimestampMs,
    /// Optionally, change the compute usage limit.
    pub compute_units_limit: Option<u32>,
    /// Optionally, change the heap size limit.
    pub heap_size_limit: Option<u32>,
}

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

    out.push_str(&format!("Oracle 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!(
        "Oracle Request Delay: {}\n",
        info.oracle_request_delay_ms
    ));

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

    if let Some(compute_units) = info.compute_units_limit {
        out.push_str(&format!("\nCompute Units Limit: {}\n", compute_units));
    }

    if let Some(heap_size) = info.heap_size_limit {
        out.push_str(&format!("Heap Size Limit: {}\n", heap_size));
    }

    out
}

impl Default for OracleInfo {
    fn default() -> Self {
        Self {
            id: OracleId::default(),
            description: String::new(),
            update_frequency: UpdateFrequency::default(),
            target_oracles: Vec::new(),
            starting_timestamp: StartingTimestamp::default(),
            is_active: false,
            created_at_ms: 0,
            validators_per_duty: default_validators_per_duty(),
            oracle_request_delay_ms: default_oracle_request_delay_ms(),
            compute_units_limit: None,
            heap_size_limit: None,
        }
    }
}

impl OracleInfo {
    /// Returns true if the oracle 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,
        }
    }

    /// Validates all oracle 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 oracles 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 oracle should be above starting_timestamp".to_string());
                            }
                        }
                    }
                }
            }
        }

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

        // Validate `oracle_request_delay`.
        if self.oracle_request_delay_ms < OracleDutyConfig::MIN_ORACLE_REQUEST_DELAY {
            return Err(format!(
                "oracle_request_delay cannot be below {}",
                OracleDutyConfig::MIN_ORACLE_REQUEST_DELAY
            ));
        }
        if self.oracle_request_delay_ms > OracleDutyConfig::MAX_ORACLE_REQUEST_DELAY_MS {
            return Err(format!(
                "oracle_request_delay cannot be above {}",
                OracleDutyConfig::MAX_ORACLE_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_oracle_output_size = max_oracle_output_serialized_bytes(self.validators_per_duty);
        if max_oracle_output_size < MIN_VIABLE_LIMIT_OF_ORACLE_OUTPUT_SIZE {
            return Err(format!("validators_per_duty is too high, results in max size of oracle updates that is too low: {max_oracle_output_size} vs {MIN_VIABLE_LIMIT_OF_ORACLE_OUTPUT_SIZE}"));
        }

        // Validate `compute_units_limit`.
        if let Some(compute_units_limit) = self.compute_units_limit {
            if compute_units_limit == 0 {
                return Err("compute_usage_limit cannot be Some(0)".to_string());
            }
            if compute_units_limit > MAX_COMPUTE_UNIT_LIMIT {
                return Err(format!("compute_usage_limit cannot be above MAX_COMPUTE_UNIT_LIMIT={MAX_COMPUTE_UNIT_LIMIT}"));
            }
        }

        // Validate `heap_size_limit`.
        if let Some(heap_size_limit) = self.heap_size_limit {
            if heap_size_limit == 0 {
                return Err("heap_size_limit cannot be Some(0)".to_string());
            }
            if heap_size_limit > MAX_HEAP_FRAME_BYTES {
                return Err(format!(
                    "heap_size_limit cannot be above MAX_HEAP_FRAME_BYTES={MAX_HEAP_FRAME_BYTES}"
                ));
            }
        }

        Ok(())
    }

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

fn validate_periodic_frequency(period_ms: TimestampMs) -> Result<(), String> {
    if period_ms == 0 {
        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 TargetOracle {
    /// Extracts the WebSocket operation if this is a WebSocket target.
    pub fn websocket_op(&self) -> Option<WebSocketOperation> {
        if let TargetOracle::WebSocket(ws_op) = self {
            Some(ws_op.clone())
        } else {
            None
        }
    }
}

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

fn default_oracle_request_delay_ms() -> TimestampMs {
    OracleDutyConfig::DEFAULT_ORACLE_REQUEST_DELAY_MS
}

/// Represents an oracle registry entry containing oracle information and metadata.
///
/// This struct stores oracle 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 OracleEntry {
    oracle_info: Arc<OracleInfo>,
    data_hash: [u8; OracleEntry::HASH_LENGTH],
    last_modified_timestamp: u64,
}

impl OracleEntry {
    const HASH_LENGTH: usize = 32;

    /// Creates a new OracleEntry 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(
        oracle_info: OracleInfo,
        data_hash: [u8; Self::HASH_LENGTH],
        last_modified_round: u64,
    ) -> Self {
        Self {
            oracle_info: Arc::new(oracle_info),
            data_hash,
            last_modified_timestamp: last_modified_round,
        }
    }

    /// Retrieves the account data.
    ///
    /// # Returns
    ///
    /// The account data as an OracleInfo.
    pub fn oracle_info(&self) -> Arc<OracleInfo> {
        self.oracle_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
    }
}

// ============================================================================
// Oracle Values
// ============================================================================

/// Represents a value that can be either plain text or encrypted.
///
/// This enum is used to handle oracle 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.
///
/// # Variants
/// * `Plain(String)` - A plain text value that is not encrypted
/// * `Encrypted(String)` - A base64-encoded encrypted value that requires decryption within a TEE
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
pub enum OracleValue {
    Plain(String),
    Encrypted(String),
}

impl Default for OracleValue {
    fn default() -> Self {
        OracleValue::Plain(String::new())
    }
}

/// A wrapper type around `OracleValue` specifically for handling URLs in oracle 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::OracleUrl;
///
/// // Create from plain text URL
/// let plain_url = OracleUrl::from("https://example.com");
///
/// // Create from encrypted URL
/// let encrypted_url = OracleUrl::from("enc://encrypted_data");
/// ```
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
pub struct OracleUrl(OracleValue);

impl fmt::Display for OracleUrl {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self.0 {
            OracleValue::Plain(ref s) => write!(f, "{}", s),
            OracleValue::Encrypted(ref s) => write!(f, "enc://{}", s),
        }
    }
}

impl Deref for OracleUrl {
    type Target = OracleValue;

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

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

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

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

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

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

/// Represents a value that can be either plain data or encrypted.
///
/// This enum is used to handle oracle 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.
///
/// # Variants
/// * `Plain(Vec<u8>)` - Data that is not encrypted
/// * `Encrypted(Vec<u8>)` - A encrypted data that requires decryption within a TEE
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
pub enum OracleValueBody {
    Plain(Vec<u8>),
    Encrypted(Vec<u8>),
}

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

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

// ============================================================================
// Oracle Targets
// ============================================================================

/// Enum that represents the target of the oracle request.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, strum_macros::AsRefStr)]
#[cfg_attr(feature = "non-pdk", derive(Subcommand))]
pub enum TargetOracle {
    /// HTTP GET request to the oracle. The URL is specified in the OracleDutyRequest.
    /// 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!(OracleUrl)
        ))]
        url: OracleUrl,
        #[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()))]
        headers: Headers,
    },
    /// HTTP POST request to the oracle. The URL is specified in the OracleDutyRequest.
    /// The URL must be an HTTPS URL. This oracle 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!(OracleUrl)
        ))]
        url: OracleUrl,
        #[cfg_attr(feature = "non-pdk", clap(long))]
        filter: Option<Vec<HttpFilter>>,
        #[cfg_attr(feature = "non-pdk", clap(
            long,
            value_parser = clap::value_parser!(OracleValueBody)
        ))]
        body: OracleValueBody,
        #[cfg_attr(feature = "non-pdk", clap(long))]
        content_type: String,
        #[cfg_attr(feature = "non-pdk", clap(long, default_value_t = Headers::default()))]
        headers: Headers,
    },
    /// Get the current time from several NIST servers.
    Time,
    /// Get the current price of several select cryptocurrencies.
    PriceReactor,
    /// Only used for testing purposes, to simulate an oracle that always returns a fixed value.
    /// TODO: remove this variant with a cfg testing flag.
    Number,
    /// 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>,
    },
    /// 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,
    },
    /// 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,
    },
    /// Reports prices for a batch of stocks based on the websocket price updates from massive.com
    Stonks,
    /// WebSocket oracle operations for persistent connections
    #[cfg_attr(feature = "non-pdk", clap(subcommand))]
    WebSocket(WebSocketOperation),
}

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

// ============================================================================
// WebSocket Operations
// ============================================================================

// ============================================================================
// WebSocket Types
// ============================================================================

/// Estimated size in bytes for system messages (OversizedWarning).
/// This accounts for the struct fields and some padding.
pub const SYSTEM_MESSAGE_SIZE: usize = 128;

/// Read mode for WebSocket buffer.
///
/// Determines how messages are retrieved from the buffer.
#[derive(
    Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize, BorshSerialize, BorshDeserialize,
)]
pub enum WebSocketReadMode {
    /// Return only the most recent message (default, backward compatible).
    #[default]
    Latest,
    /// Return all accumulated messages.
    All,
    /// Return messages newer than the given index.
    ///
    /// **Note:** Index wraparound at `u64::MAX` is not handled. See DataBuffer
    /// documentation for limitations. At 1 million messages per second,
    /// wraparound would take ~584,000 years.
    FromIndex(u64),
}

/// Message content type - either user data or system notification.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
pub enum MessageContent {
    /// Normal message data from WebSocket.
    Data(Vec<u8>),
    /// System-generated message for oversized content that was rejected.
    OversizedWarning {
        /// Size of the rejected message in bytes.
        original_size: usize,
        /// Buffer size limit in bytes.
        limit: usize,
        /// When the oversized message was received (RFC3339 with microseconds).
        original_timestamp: String,
    },
}

impl MessageContent {
    /// Returns the byte size of this content for buffer accounting.
    ///
    /// For data messages, this is the actual data length.
    /// For system messages, this is a fixed size constant.
    pub fn byte_size(&self) -> usize {
        match self {
            MessageContent::Data(data) => data.len(),
            MessageContent::OversizedWarning { .. } => SYSTEM_MESSAGE_SIZE,
        }
    }

    /// Returns true if this is a system message.
    pub fn is_system(&self) -> bool {
        matches!(self, MessageContent::OversizedWarning { .. })
    }

    /// Returns the data bytes if this is a Data message, None otherwise.
    pub fn as_data(&self) -> Option<&[u8]> {
        match self {
            MessageContent::Data(data) => Some(data),
            MessageContent::OversizedWarning { .. } => None,
        }
    }
}

/// A single buffered message with metadata.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
pub struct BufferedMessage {
    /// Unique, monotonically increasing index (never reused, persists across reconnects).
    pub index: u64,
    /// The message content.
    pub content: MessageContent,
    /// When this message was received/created (RFC3339 with microseconds).
    pub received_at: String,
}

impl BufferedMessage {
    /// Returns the size of this message in bytes (for buffer accounting).
    pub fn byte_size(&self) -> usize {
        self.content.byte_size()
    }

    /// Returns true if this is a system message.
    pub fn is_system_message(&self) -> bool {
        self.content.is_system()
    }

    /// Returns the data bytes if this contains a Data message, None otherwise.
    pub fn as_data(&self) -> Option<&[u8]> {
        self.content.as_data()
    }
}

/// Response for WebSocket read operations.
///
/// This struct captures the read result with messages and buffer state information.
///
/// # Gap Detection
///
/// Consumers can detect gaps by comparing message indices. Each `BufferedMessage`
/// contains an `index` field that increases monotonically. If there's a gap between
/// indices, some messages were evicted from the buffer.
///
/// # Convenience Methods
///
/// For common access patterns, use the convenience methods instead of accessing
/// fields directly:
///
/// ```ignore
/// // Instead of:
/// let data = response.messages.first().unwrap().content.as_data().unwrap();
///
/// // Use:
/// let data = response.first_data().unwrap();
/// ```
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
pub struct WebSocketReadResponse {
    /// All requested messages (based on read mode).
    pub messages: Vec<BufferedMessage>,
    /// Current highest index in buffer (for tracking).
    pub latest_index: Option<u64>,
    /// Current lowest index in buffer (helpful for gap detection).
    pub oldest_index: Option<u64>,
}

impl WebSocketReadResponse {
    /// Returns true if there are any messages in the response.
    pub fn has_messages(&self) -> bool {
        !self.messages.is_empty()
    }

    /// Returns the number of messages in the response.
    pub fn message_count(&self) -> usize {
        self.messages.len()
    }

    /// Get the first (oldest) message, if any.
    pub fn first_message(&self) -> Option<&BufferedMessage> {
        self.messages.first()
    }

    /// Get the last (most recent) message, if any.
    pub fn latest_message(&self) -> Option<&BufferedMessage> {
        self.messages.last()
    }

    /// Get the data from the first message (convenience for common pattern).
    ///
    /// Returns `None` if there are no messages or if the first message is a system message.
    pub fn first_data(&self) -> Option<&[u8]> {
        self.messages.first().and_then(|m| m.as_data())
    }

    /// Get the data from the latest message (convenience for common pattern).
    ///
    /// Returns `None` if there are no messages or if the latest message is a system message.
    pub fn latest_data(&self) -> Option<&[u8]> {
        self.messages.last().and_then(|m| m.as_data())
    }

    /// Iterate over only the data messages (filtering out system messages).
    pub fn data_messages(&self) -> impl Iterator<Item = &BufferedMessage> {
        self.messages.iter().filter(|m| !m.is_system_message())
    }

    /// Iterate over only the data bytes (convenience for processing).
    ///
    /// This filters out system messages and returns only the raw data bytes.
    pub fn iter_data(&self) -> impl Iterator<Item = &[u8]> {
        self.messages.iter().filter_map(|m| m.as_data())
    }

    /// Returns the number of data messages (excluding system messages).
    pub fn data_message_count(&self) -> usize {
        self.messages
            .iter()
            .filter(|m| !m.is_system_message())
            .count()
    }

    /// Returns true if any message is a system message (like OversizedWarning).
    pub fn has_system_messages(&self) -> bool {
        self.messages.iter().any(|m| m.is_system_message())
    }
}

/// WebSocket operation types for the WebSocket oracle.
///
/// This enum defines the different operations that can be performed on WebSocket connections.
/// All WebSocket operations are handled through a single `TargetOracle::WebSocket` variant.
///
/// # Variants
/// * `Connect` - Establish a new WebSocket connection to an external server
/// * `Read` - Read the latest data from an existing WebSocket connection
///
/// # Future Extensions
/// Additional operations like `Send` and `Close` may be added in the future.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
#[cfg_attr(feature = "non-pdk", derive(Subcommand))]
pub enum WebSocketOperation {
    /// Establish a new WebSocket connection.
    ///
    /// When this operation is executed:
    /// - A validator is randomly selected to handle the connection
    /// - The TEE establishes the WebSocket connection to the specified URL
    /// - On success, the connection is registered in the WebSocket Registry
    /// - The assigned validator is recorded for future routing
    Connect {
        /// WebSocket URL (typically wss://...)
        #[cfg_attr(feature = "non-pdk", clap(
            long = "target-url",
            value_parser = clap::value_parser!(OracleUrl)
        ))]
        url: OracleUrl,
        /// The OracleId that represents this connection request
        #[cfg_attr(feature = "non-pdk", clap(long, value_parser = parse_oracle_id))]
        oracle_id: OracleId,
    },
    /// Read data from an existing WebSocket connection.
    ///
    /// When this operation is executed:
    /// - The operation is automatically routed to the TEE holding the connection
    /// - Returns messages from the connection's buffer based on the read mode
    /// - The referenced connection must exist and be active
    ///
    /// # Read Modes
    ///
    /// - `Latest` (default): Returns only the most recent message
    /// - `All`: Returns all accumulated messages in the buffer
    /// - `FromIndex(n)`: Returns all messages with index > n, with gap detection
    ///
    /// # Validation
    ///
    /// The handler implementation must verify that:
    /// - `connection_oracle_id` references an existing Connect operation in the WebSocket Registry
    /// - The connection is still active and hasn't been closed
    /// - The requesting validator has permission to read from this connection
    ///
    /// # Error Handling
    ///
    /// If validation fails, the oracle should return an appropriate error indicating:
    /// - `ConnectionNotFound` - if the referenced oracle ID doesn't exist
    /// - `ConnectionClosed` - if the connection has been terminated
    /// - `InvalidConnectionType` - if the referenced oracle is not a WebSocket Connect operation
    Read {
        /// References the OracleId of the Connect operation that established the connection
        #[cfg_attr(feature = "non-pdk", clap(long, value_parser = parse_oracle_id))]
        connection_oracle_id: OracleId,
        /// The read mode determining which messages to return (defaults to Latest for backward compatibility)
        #[serde(default)]
        #[cfg_attr(feature = "non-pdk", clap(skip))]
        mode: WebSocketReadMode,
    },
    /// Send messages to an existing WebSocket connection.
    ///
    /// When this operation is executed:
    /// - The operation is automatically routed to the TEE holding the connection
    /// - Each `OracleValue` in `messages` is sent as a separate WebSocket message
    /// - The inner value is extracted: `Plain(String)` → sent as-is, `Encrypted(String)` → decrypted then sent
    /// - Returns the number of messages successfully sent
    /// - The referenced connection must exist and be active
    ///
    /// # Validation
    ///
    /// The handler implementation must verify that:
    /// - `connection_oracle_id` references an existing Connect operation in the WebSocket Registry
    /// - The connection is still active and hasn't been closed
    /// - The requesting validator has permission to send to this connection
    ///
    /// # Error Handling
    ///
    /// If validation fails, the oracle should return an appropriate error indicating:
    /// - `ConnectionNotFound` - if the referenced oracle ID doesn't exist
    /// - `ConnectionClosed` - if the connection has been terminated
    /// - `InvalidConnectionType` - if the referenced oracle is not a WebSocket Connect operation
    /// - `SecretDecryptionFailed` - if an encrypted message cannot be decrypted
    Send {
        /// References the OracleId of the Connect operation that established the connection
        #[cfg_attr(feature = "non-pdk", clap(long, value_parser = parse_oracle_id))]
        connection_oracle_id: OracleId,
        /// Messages to send to the WebSocket connection
        /// Each OracleValue will be sent as a separate WebSocket message
        #[cfg_attr(feature = "non-pdk", clap(skip))]
        messages: Vec<OracleValue>,
    },
    /// Close an existing WebSocket connection.
    ///
    /// When this operation is executed:
    /// - The operation is automatically routed to the TEE holding the connection
    /// - The WebSocket connection is gracefully closed
    /// - The connection status in the WebSocket Registry is updated to `Closed`
    /// - The referenced connection must exist and be active
    ///
    /// # Validation
    ///
    /// The handler implementation must verify that:
    /// - `connection_oracle_id` references an existing Connect operation in the WebSocket Registry
    /// - The requesting validator has permission to close this connection
    ///
    /// # Error Handling
    ///
    /// If validation fails, the oracle should return an appropriate error indicating:
    /// - `ConnectionNotFound` - if the referenced oracle ID doesn't exist
    /// - `ConnectionAlreadyClosed` - if the connection has already been terminated
    /// - `InvalidConnectionType` - if the referenced oracle is not a WebSocket Connect operation
    Close {
        /// References the OracleId of the Connect operation to close
        #[cfg_attr(feature = "non-pdk", clap(long, value_parser = parse_oracle_id))]
        connection_oracle_id: OracleId,
    },
}

impl WebSocketOperation {
    /// Create a new Connect operation.
    pub fn connect(url: impl Into<OracleUrl>, oracle_id: OracleId) -> Self {
        WebSocketOperation::Connect {
            url: url.into(),
            oracle_id,
        }
    }

    /// Create a new Read operation with default Latest mode.
    pub fn read(connection_oracle_id: OracleId) -> Self {
        WebSocketOperation::Read {
            connection_oracle_id,
            mode: WebSocketReadMode::default(),
        }
    }

    /// Create a new Read operation with a specific mode.
    pub fn read_with_mode(connection_oracle_id: OracleId, mode: WebSocketReadMode) -> Self {
        WebSocketOperation::Read {
            connection_oracle_id,
            mode,
        }
    }

    /// Create a new Send operation.
    pub fn send(connection_oracle_id: OracleId, messages: Vec<OracleValue>) -> Self {
        WebSocketOperation::Send {
            connection_oracle_id,
            messages,
        }
    }

    /// Create a new Close operation.
    pub fn close(connection_oracle_id: OracleId) -> Self {
        WebSocketOperation::Close {
            connection_oracle_id,
        }
    }

    /// Returns true if this is a Connect operation.
    pub fn is_connect(&self) -> bool {
        matches!(self, WebSocketOperation::Connect { .. })
    }

    /// Returns true if this is a Read operation.
    pub fn is_read(&self) -> bool {
        matches!(self, WebSocketOperation::Read { .. })
    }

    /// Returns true if this is a Send operation.
    pub fn is_send(&self) -> bool {
        matches!(self, WebSocketOperation::Send { .. })
    }

    /// Returns true if this is a Close operation.
    pub fn is_close(&self) -> bool {
        matches!(self, WebSocketOperation::Close { .. })
    }

    /// Returns the URL if this is a Connect operation, None otherwise.
    pub fn url(&self) -> Option<&OracleUrl> {
        match self {
            WebSocketOperation::Connect { url, .. } => Some(url),
            _ => None,
        }
    }

    /// Returns the oracle_id if this is a Connect operation, None otherwise.
    pub fn oracle_id(&self) -> Option<&OracleId> {
        match self {
            WebSocketOperation::Connect { oracle_id, .. } => Some(oracle_id),
            _ => None,
        }
    }

    /// Returns the connection_oracle_id if this is a Read, Send, or Close operation, None otherwise.
    pub fn connection_oracle_id(&self) -> Option<&OracleId> {
        match self {
            WebSocketOperation::Read {
                connection_oracle_id,
                ..
            }
            | WebSocketOperation::Send {
                connection_oracle_id,
                ..
            }
            | WebSocketOperation::Close {
                connection_oracle_id,
            } => Some(connection_oracle_id),
            _ => None,
        }
    }

    /// Returns the read mode if this is a Read operation, None otherwise.
    pub fn read_mode(&self) -> Option<&WebSocketReadMode> {
        match self {
            WebSocketOperation::Read { mode, .. } => Some(mode),
            _ => None,
        }
    }

    /// Returns the messages if this is a Send operation, None otherwise.
    pub fn messages(&self) -> Option<&[OracleValue]> {
        match self {
            WebSocketOperation::Send { messages, .. } => Some(messages),
            _ => None,
        }
    }
}

impl fmt::Display for WebSocketOperation {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            WebSocketOperation::Connect { url, oracle_id } => {
                write!(f, "Connect(url={url}, oracle_id={oracle_id})")
            }
            WebSocketOperation::Read {
                connection_oracle_id,
                mode,
            } => {
                write!(
                    f,
                    "Read(connection_oracle_id={}, mode={:?})",
                    connection_oracle_id, mode
                )
            }
            WebSocketOperation::Send {
                connection_oracle_id,
                messages,
            } => {
                write!(
                    f,
                    "Send(connection_oracle_id={}, messages_count={})",
                    connection_oracle_id,
                    messages.len()
                )
            }
            WebSocketOperation::Close {
                connection_oracle_id,
            } => {
                write!(f, "Close(connection_oracle_id={})", connection_oracle_id)
            }
        }
    }
}

impl FromStr for TargetOracle {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s == "Time" {
            Ok(TargetOracle::Time)
        } else if s == "PriceReactor" {
            Ok(TargetOracle::PriceReactor)
        } else if s == "SecretKeyGeneration" {
            Err("SecretKeyGeneration oracle requires committee_id and committee_members parameters. Use the appropriate API to create this oracle type.".to_string())
        } else if s == "SecretKeyEncryption" {
            Err("SecretKeyEncryption oracle requires target_tee_id, secret_data, and committee_id parameters. Use the appropriate API to create this oracle type.".to_string())
        } else if s == "SecretKeyDecryption" {
            Err("SecretKeyDecryption oracle requires encrypted_data and source_committee_id parameters. Use the appropriate API to create this oracle type.".to_string())
        } else if s == "number" {
            Err("The 'number' oracle 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(TargetOracle::HttpGet {
                    url: url.into(),
                    filter,
                    headers: Headers::default(),
                });
            }

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

// ============================================================================
// Oracle Updates
// ============================================================================

/// 32‑byte Blake3 hash of the raw request payload observed by the oracle 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 oracle response transcript.
///
/// The transcript signed by the oracle is `input_commitment || blake3(response_value)`.
pub type SignatureBytes = [u8; 64];

/// A structure representing the result of an oracle 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-oracle-subsystem>
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct OracleUpdateResult {
    /// The identifier of the oracle (should be kept under 100 bytes)
    pub oracle_id: OracleId,

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

    /// Blake3 hash of `oracle_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 oracle service.
    #[serde(with = "BigArray")]
    pub input_commitment: InputCommitmentBytes,

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

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

    /// Optional attestation report, if available
    /// This can be used to provide additional context or verification of the oracle 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 OracleUpdateResult {
    /// Create an OracleUpdateResult
    pub fn new(
        oracle_id: OracleId,
        target_timestamp: TimestampMs,
        oracle_result: Vec<u8>,
        input_commitment: InputCommitmentBytes,
        signature: SignatureBytes,
        attestation_report: Option<AttestationReport>,
        authority_key: AuthorityKeyBytes,
    ) -> Result<Self, &'static str> {
        let hash = blake3::hash(&oracle_result);

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

        #[cfg(not(feature = "non-pdk"))]
        let oracle_result = oracle_result;

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

impl Default for OracleUpdateResult {
    fn default() -> Self {
        Self {
            oracle_id: OracleId::default(),
            target_timestamp: 0,
            response_hash: [0; 32],
            input_commitment: [0xee; 32],
            signature: [0; 64],
            oracle_result: vec![],
            attestation_report: None,
            authority_key: [0xff; 96],
        }
    }
}

// ============================================================================
// Oracle Requests & Scheduling
// ============================================================================

/// Represents a request to an oracle 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 OracleRequest {
    /// Structured fields of the request.
    pub oracle_id: Option<OracleId>,
    pub target_timestamp: Option<TimestampMs>,
    pub authority_key: AuthorityKeyBytes,
    pub include_attestation: bool,
    pub max_oracle_output_size: u32,

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

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

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

/// The frequency of the oracle update.
#[derive(Debug, Default, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum UpdateFrequency {
    /// The oracle will update once.
    #[default]
    OneShot,
    /// The oracle will update every N milliseconds.
    Periodic(TimestampMs),
    /// The oracle 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(TimestampMs, TimestampMs),
}

/// Starting timestamp configuration for oracles
#[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(0)
    }
}

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

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

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

    fn base_valid_oracle_info() -> OracleInfo {
        OracleInfo {
            description: "test".to_string(),
            target_oracles: vec![TargetOracle::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(0),
            // Keep other fields as default
            ..OracleInfo::default()
        }
    }

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

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

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

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

    #[test]
    fn test_periodic_with_zero_period_is_invalid() {
        let mut info = base_valid_oracle_info();
        info.starting_timestamp = StartingTimestamp::Timestamp(1);
        info.update_frequency = UpdateFrequency::Periodic(0);
        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_oracle_info();
        info.starting_timestamp = StartingTimestamp::Timestamp(1);
        info.update_frequency = UpdateFrequency::LimitedPeriodic(0, 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_oracle_info();
        info.starting_timestamp = StartingTimestamp::Timestamp(500);
        info.update_frequency = UpdateFrequency::LimitedPeriodic(300, 500);
        let err = info.validate().unwrap_err();
        assert!(err.contains(
            "end_timestamp of a LimitedPeriodic oracle should be above starting_timestamp"
        ));
    }

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

    #[test]
    fn test_target_oracles_cannot_be_empty() {
        let mut info = base_valid_oracle_info();
        info.target_oracles.clear();
        let err = info.validate().unwrap_err();
        assert!(err.contains("OracleTargets cannot be empty"));
    }

    #[test]
    fn test_oracle_request_delay_bounds() {
        // Below minimum
        let mut info = base_valid_oracle_info();
        info.oracle_request_delay_ms = OracleDutyConfig::MIN_ORACLE_REQUEST_DELAY - 1;
        let err = info.validate().unwrap_err();
        assert!(err.contains(&format!(
            "oracle_request_delay cannot be below {}",
            OracleDutyConfig::MIN_ORACLE_REQUEST_DELAY
        )));

        // Above maximum
        let mut info = base_valid_oracle_info();
        info.oracle_request_delay_ms = OracleDutyConfig::MAX_ORACLE_REQUEST_DELAY_MS + 1;
        let err = info.validate().unwrap_err();
        assert!(err.contains(&format!(
            "oracle_request_delay cannot be above {}",
            OracleDutyConfig::MAX_ORACLE_REQUEST_DELAY_MS
        )));
    }

    #[test]
    fn test_validators_per_duty_cannot_be_zero() {
        let mut info = base_valid_oracle_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_oracle_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_compute_units_limit_checks() {
        // Zero is invalid
        let mut info = base_valid_oracle_info();
        info.compute_units_limit = Some(0);
        let err = info.validate().unwrap_err();
        assert!(err.contains("compute_usage_limit cannot be Some(0)"));

        // Heap size above the max is invalid
        let mut info = base_valid_oracle_info();
        info.compute_units_limit = Some(MAX_COMPUTE_UNIT_LIMIT + 1);
        let err = info.validate().unwrap_err();
        assert!(err.contains(&format!(
            "compute_usage_limit cannot be above MAX_COMPUTE_UNIT_LIMIT={}",
            MAX_COMPUTE_UNIT_LIMIT
        )));
    }

    #[test]
    fn test_heap_size_limit_checks() {
        // Zero is invalid
        let mut info = base_valid_oracle_info();
        info.heap_size_limit = Some(0);
        let err = info.validate().unwrap_err();
        assert!(err.contains("heap_size_limit cannot be Some(0)"));

        // Heap size above the max is invalid
        let mut info = base_valid_oracle_info();
        info.heap_size_limit = Some(MAX_HEAP_FRAME_BYTES + 1);
        let err = info.validate().unwrap_err();
        assert!(err.contains(&format!(
            "heap_size_limit cannot be above MAX_HEAP_FRAME_BYTES={}",
            MAX_HEAP_FRAME_BYTES
        )));
    }

    // ========================================================================
    // WebSocket Operation Tests
    // ========================================================================

    #[test]
    fn test_websocket_connect_operation_creation() {
        let oracle_id = OracleId::new(Pubkey::default(), 1u64);
        let op = WebSocketOperation::connect("wss://example.com/stream", oracle_id);
        assert!(op.is_connect());
        assert!(!op.is_read());

        if let WebSocketOperation::Connect {
            url,
            oracle_id: op_oracle_id,
        } = op
        {
            assert_eq!(url.to_string(), "wss://example.com/stream");
            assert_eq!(op_oracle_id, oracle_id);
        } else {
            panic!("Expected Connect variant");
        }
    }

    #[test]
    fn test_websocket_read_operation_creation() {
        let oracle_id = OracleId::new(Pubkey::default(), 42u64);
        let op = WebSocketOperation::read(oracle_id);
        assert!(op.is_read());
        assert!(!op.is_connect());

        if let WebSocketOperation::Read {
            connection_oracle_id,
            mode,
        } = op
        {
            assert_eq!(connection_oracle_id, oracle_id);
            assert_eq!(mode, WebSocketReadMode::Latest);
        } else {
            panic!("Expected Read variant");
        }
    }

    #[test]
    fn test_websocket_operation_serde_roundtrip_connect() {
        let oracle_id = OracleId::new(Pubkey::default(), 1u64);
        let op = WebSocketOperation::Connect {
            url: "wss://example.com/stream".into(),
            oracle_id,
        };

        // Serialize to JSON
        let json = serde_json::to_string(&op).expect("Failed to serialize");
        // Deserialize back
        let deserialized: WebSocketOperation =
            serde_json::from_str(&json).expect("Failed to deserialize");

        assert_eq!(op, deserialized);
    }

    #[test]
    fn test_websocket_operation_serde_roundtrip_read() {
        let oracle_id = OracleId::new(Pubkey::default(), 123u64);
        let op = WebSocketOperation::Read {
            connection_oracle_id: oracle_id,
            mode: WebSocketReadMode::default(),
        };

        // Serialize to JSON
        let json = serde_json::to_string(&op).expect("Failed to serialize");
        // Deserialize back
        let deserialized: WebSocketOperation =
            serde_json::from_str(&json).expect("Failed to deserialize");

        assert_eq!(op, deserialized);
    }

    #[test]
    fn test_target_oracle_websocket_serde_roundtrip() {
        let oracle_id = OracleId::new(Pubkey::default(), 1u64);
        let target = TargetOracle::WebSocket(WebSocketOperation::Connect {
            url: "wss://example.com/stream".into(),
            oracle_id,
        });

        // Serialize to JSON
        let json = serde_json::to_string(&target).expect("Failed to serialize");
        // Deserialize back
        let deserialized: TargetOracle =
            serde_json::from_str(&json).expect("Failed to deserialize");

        assert_eq!(target, deserialized);
    }

    #[test]
    fn test_websocket_operation_display() {
        let oracle_id = OracleId::new(Pubkey::default(), 1u64);
        let connect_op = WebSocketOperation::Connect {
            url: "wss://example.com".into(),
            oracle_id,
        };
        let display = format!("{}", connect_op);
        assert!(display.contains("Connect"));
        assert!(display.contains("wss://example.com"));

        let oracle_id = OracleId::new(Pubkey::default(), 1u64);
        let read_op = WebSocketOperation::Read {
            connection_oracle_id: oracle_id,
            mode: WebSocketReadMode::default(),
        };
        let display = format!("{}", read_op);
        assert!(display.contains("Read"));
        assert!(display.contains("connection_oracle_id"));
    }

    #[test]
    fn test_websocket_connect_with_encrypted_url() {
        let oracle_id = OracleId::new(Pubkey::default(), 1u64);
        let op = WebSocketOperation::Connect {
            url: "enc://encrypted_websocket_url".into(),
            oracle_id,
        };

        if let WebSocketOperation::Connect { url, .. } = &op {
            assert_eq!(url.to_string(), "enc://encrypted_websocket_url");
        }

        // Verify serde roundtrip preserves encrypted URL
        let json = serde_json::to_string(&op).expect("Failed to serialize");
        let deserialized: WebSocketOperation =
            serde_json::from_str(&json).expect("Failed to deserialize");
        assert_eq!(op, deserialized);
    }

    #[test]
    fn test_oracle_info_with_websocket_target() {
        let oracle_id = OracleId::new(Pubkey::default(), 1u64);
        let mut info = base_valid_oracle_info();
        info.target_oracles = vec![TargetOracle::WebSocket(WebSocketOperation::Connect {
            url: "wss://example.com/stream".into(),
            oracle_id,
        })];

        // Should validate successfully
        assert!(info.validate().is_ok());
    }

    #[test]
    fn test_websocket_operation_borsh_roundtrip_connect() {
        let oracle_id = OracleId::new(Pubkey::default(), 1u64);
        let op = WebSocketOperation::Connect {
            url: "wss://example.com/stream".into(),
            oracle_id,
        };

        // Serialize to Borsh
        let bytes = borsh::to_vec(&op).expect("Failed to serialize");
        // Deserialize back
        let deserialized: WebSocketOperation =
            borsh::from_slice(&bytes).expect("Failed to deserialize");

        assert_eq!(op, deserialized);
    }

    #[test]
    fn test_websocket_operation_borsh_roundtrip_read() {
        let oracle_id = OracleId::new(Pubkey::default(), 123u64);
        let op = WebSocketOperation::Read {
            connection_oracle_id: oracle_id,
            mode: WebSocketReadMode::default(),
        };

        // Serialize to Borsh
        let bytes = borsh::to_vec(&op).expect("Failed to serialize");
        // Deserialize back
        let deserialized: WebSocketOperation =
            borsh::from_slice(&bytes).expect("Failed to deserialize");

        assert_eq!(op, deserialized);
    }

    #[test]
    fn test_websocket_operation_borsh_roundtrip_encrypted_url() {
        let oracle_id = OracleId::new(Pubkey::default(), 1u64);
        let op = WebSocketOperation::Connect {
            url: "enc://encrypted_websocket_url".into(),
            oracle_id,
        };

        // Serialize to Borsh
        let bytes = borsh::to_vec(&op).expect("Failed to serialize");
        // Deserialize back
        let deserialized: WebSocketOperation =
            borsh::from_slice(&bytes).expect("Failed to deserialize");

        assert_eq!(op, deserialized);
    }

    #[test]
    fn test_websocket_operation_url_getter() {
        // Connect operation should return Some(url)
        let oracle_id = OracleId::new(Pubkey::default(), 1u64);
        let connect_op = WebSocketOperation::Connect {
            url: "wss://example.com/stream".into(),
            oracle_id,
        };
        assert!(connect_op.url().is_some());
        assert_eq!(
            connect_op.url().unwrap().to_string(),
            "wss://example.com/stream"
        );

        // Read operation should return None
        let oracle_id = OracleId::new(Pubkey::default(), 42u64);
        let read_op = WebSocketOperation::Read {
            connection_oracle_id: oracle_id,
            mode: WebSocketReadMode::default(),
        };
        assert!(read_op.url().is_none());
    }

    #[test]
    fn test_websocket_operation_connection_oracle_id_getter() {
        // Read operation should return Some(connection_oracle_id)
        let oracle_id = OracleId::new(Pubkey::default(), 42u64);
        let read_op = WebSocketOperation::Read {
            connection_oracle_id: oracle_id,
            mode: WebSocketReadMode::default(),
        };
        assert!(read_op.connection_oracle_id().is_some());
        assert_eq!(*read_op.connection_oracle_id().unwrap(), oracle_id);

        // Connect operation should return None
        let oracle_id = OracleId::new(Pubkey::default(), 1u64);
        let connect_op = WebSocketOperation::Connect {
            url: "wss://example.com/stream".into(),
            oracle_id,
        };
        assert!(connect_op.connection_oracle_id().is_none());

        // Send operation should return Some(connection_oracle_id)
        let oracle_id = OracleId::new(Pubkey::default(), 99u64);
        let send_op = WebSocketOperation::Send {
            connection_oracle_id: oracle_id,
            messages: vec![],
        };
        assert!(send_op.connection_oracle_id().is_some());
        assert_eq!(*send_op.connection_oracle_id().unwrap(), oracle_id);
    }

    // ========================================================================
    // WebSocket Send Operation Tests
    // ========================================================================

    #[test]
    fn test_websocket_send_operation_creation() {
        let oracle_id = OracleId::new(Pubkey::default(), 5u64);
        let messages = vec![
            OracleValue::Plain("message1".to_string()),
            OracleValue::Plain("message2".to_string()),
        ];
        let op = WebSocketOperation::send(oracle_id, messages.clone());
        assert!(op.is_send());
        assert!(!op.is_connect());
        assert!(!op.is_read());

        if let WebSocketOperation::Send {
            connection_oracle_id,
            messages: op_messages,
        } = op
        {
            assert_eq!(connection_oracle_id, oracle_id);
            assert_eq!(op_messages, messages);
        } else {
            panic!("Expected Send variant");
        }
    }

    #[test]
    fn test_websocket_send_messages_getter() {
        let oracle_id = OracleId::new(Pubkey::default(), 10u64);
        let messages = vec![OracleValue::Plain("test".to_string())];
        let send_op = WebSocketOperation::Send {
            connection_oracle_id: oracle_id,
            messages: messages.clone(),
        };

        assert!(send_op.messages().is_some());
        assert_eq!(send_op.messages().unwrap(), &messages[..]);

        // Connect and Read should return None
        let connect_op = WebSocketOperation::Connect {
            url: "wss://example.com".into(),
            oracle_id,
        };
        assert!(connect_op.messages().is_none());

        let read_op = WebSocketOperation::Read {
            connection_oracle_id: oracle_id,
            mode: WebSocketReadMode::default(),
        };
        assert!(read_op.messages().is_none());
    }

    #[test]
    fn test_websocket_operation_serde_roundtrip_send() {
        let oracle_id = OracleId::new(Pubkey::default(), 7u64);
        let messages = vec![
            OracleValue::Plain("plain_message".to_string()),
            OracleValue::Encrypted("encrypted_data".to_string()),
        ];
        let op = WebSocketOperation::Send {
            connection_oracle_id: oracle_id,
            messages,
        };

        // Serialize to JSON
        let json = serde_json::to_string(&op).expect("Failed to serialize");
        // Deserialize back
        let deserialized: WebSocketOperation =
            serde_json::from_str(&json).expect("Failed to deserialize");

        assert_eq!(op, deserialized);
    }

    #[test]
    fn test_websocket_operation_borsh_roundtrip_send() {
        let oracle_id = OracleId::new(Pubkey::default(), 8u64);
        let messages = vec![
            OracleValue::Plain("hello".to_string()),
            OracleValue::Plain("world".to_string()),
        ];
        let op = WebSocketOperation::Send {
            connection_oracle_id: oracle_id,
            messages,
        };

        // Serialize to Borsh
        let bytes = borsh::to_vec(&op).expect("Failed to serialize");
        // Deserialize back
        let deserialized: WebSocketOperation =
            borsh::from_slice(&bytes).expect("Failed to deserialize");

        assert_eq!(op, deserialized);
    }

    #[test]
    fn test_websocket_send_display() {
        let oracle_id = OracleId::new(Pubkey::default(), 20u64);
        let messages = vec![
            OracleValue::Plain("msg1".to_string()),
            OracleValue::Plain("msg2".to_string()),
            OracleValue::Plain("msg3".to_string()),
        ];
        let send_op = WebSocketOperation::Send {
            connection_oracle_id: oracle_id,
            messages,
        };

        let display = format!("{}", send_op);
        assert!(display.contains("Send"));
        assert!(display.contains("connection_oracle_id"));
        assert!(display.contains("messages_count=3"));
    }

    #[test]
    fn test_websocket_send_empty_messages() {
        let oracle_id = OracleId::new(Pubkey::default(), 15u64);
        let op = WebSocketOperation::send(oracle_id, vec![]);

        assert!(op.is_send());
        assert_eq!(op.messages().unwrap().len(), 0);

        let display = format!("{}", op);
        assert!(display.contains("messages_count=0"));
    }

    #[test]
    fn test_websocket_send_with_encrypted_messages() {
        let oracle_id = OracleId::new(Pubkey::default(), 25u64);
        let messages = vec![
            OracleValue::Encrypted("encrypted1".to_string()),
            OracleValue::Encrypted("encrypted2".to_string()),
        ];
        let op = WebSocketOperation::Send {
            connection_oracle_id: oracle_id,
            messages: messages.clone(),
        };

        // Verify serde roundtrip preserves encrypted messages
        let json = serde_json::to_string(&op).expect("Failed to serialize");
        let deserialized: WebSocketOperation =
            serde_json::from_str(&json).expect("Failed to deserialize");
        assert_eq!(op, deserialized);

        // Verify borsh roundtrip preserves encrypted messages
        let bytes = borsh::to_vec(&op).expect("Failed to serialize");
        let deserialized: WebSocketOperation =
            borsh::from_slice(&bytes).expect("Failed to deserialize");
        assert_eq!(op, deserialized);
    }

    #[test]
    fn test_target_oracle_websocket_send_serde_roundtrip() {
        let oracle_id = OracleId::new(Pubkey::default(), 30u64);
        let messages = vec![OracleValue::Plain("test_message".to_string())];
        let target = TargetOracle::WebSocket(WebSocketOperation::Send {
            connection_oracle_id: oracle_id,
            messages,
        });

        // Serialize to JSON
        let json = serde_json::to_string(&target).expect("Failed to serialize");
        // Deserialize back
        let deserialized: TargetOracle =
            serde_json::from_str(&json).expect("Failed to deserialize");

        assert_eq!(target, deserialized);
    }

    // ========================================================================
    // WebSocket Close Operation Tests
    // ========================================================================

    #[test]
    fn test_websocket_close_operation_creation() {
        let oracle_id = OracleId::new(Pubkey::default(), 50u64);
        let op = WebSocketOperation::close(oracle_id);
        assert!(op.is_close());
        assert!(!op.is_connect());
        assert!(!op.is_read());
        assert!(!op.is_send());

        if let WebSocketOperation::Close {
            connection_oracle_id,
        } = op
        {
            assert_eq!(connection_oracle_id, oracle_id);
        } else {
            panic!("Expected Close variant");
        }
    }

    #[test]
    fn test_websocket_close_connection_oracle_id_getter() {
        let oracle_id = OracleId::new(Pubkey::default(), 55u64);
        let close_op = WebSocketOperation::Close {
            connection_oracle_id: oracle_id,
        };
        assert!(close_op.connection_oracle_id().is_some());
        assert_eq!(*close_op.connection_oracle_id().unwrap(), oracle_id);
    }

    #[test]
    fn test_websocket_operation_serde_roundtrip_close() {
        let oracle_id = OracleId::new(Pubkey::default(), 60u64);
        let op = WebSocketOperation::Close {
            connection_oracle_id: oracle_id,
        };

        // Serialize to JSON
        let json = serde_json::to_string(&op).expect("Failed to serialize");
        // Deserialize back
        let deserialized: WebSocketOperation =
            serde_json::from_str(&json).expect("Failed to deserialize");

        assert_eq!(op, deserialized);
    }

    #[test]
    fn test_websocket_operation_borsh_roundtrip_close() {
        let oracle_id = OracleId::new(Pubkey::default(), 65u64);
        let op = WebSocketOperation::Close {
            connection_oracle_id: oracle_id,
        };

        // Serialize to Borsh
        let bytes = borsh::to_vec(&op).expect("Failed to serialize");
        // Deserialize back
        let deserialized: WebSocketOperation =
            borsh::from_slice(&bytes).expect("Failed to deserialize");

        assert_eq!(op, deserialized);
    }

    #[test]
    fn test_websocket_close_display() {
        let oracle_id = OracleId::new(Pubkey::default(), 70u64);
        let close_op = WebSocketOperation::Close {
            connection_oracle_id: oracle_id,
        };

        let display = format!("{}", close_op);
        assert!(display.contains("Close"));
        assert!(display.contains("connection_oracle_id"));
        assert!(display.contains(&oracle_id.to_string()));
    }

    #[test]
    fn test_target_oracle_websocket_close_serde_roundtrip() {
        let oracle_id = OracleId::new(Pubkey::default(), 75u64);
        let target = TargetOracle::WebSocket(WebSocketOperation::Close {
            connection_oracle_id: oracle_id,
        });

        // Serialize to JSON
        let json = serde_json::to_string(&target).expect("Failed to serialize");
        // Deserialize back
        let deserialized: TargetOracle =
            serde_json::from_str(&json).expect("Failed to deserialize");

        assert_eq!(target, deserialized);
    }
}