// 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::{
    websocket_op::WebSocketOperation, AttestationReport, AuthorityKeyBytes, Headers, HttpFilter,
    Nonce, RexDutyConfig,
};

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

/// 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: TimestampMs = 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: TimestampMs,
}

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,
        }
    }

    /// 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: 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 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
        }
    }
}

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

fn default_rex_request_delay_ms() -> TimestampMs {
    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
// ============================================================================

/// 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.
///
/// # 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 RexValue {
    Plain(Vec<u8>),
    Encrypted(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![])
    }
}

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.
    ///
    /// The variant is always present for serialization stability; the `wasm-rex` feature
    /// gate is on *execution* code only.
    Wasm {
        /// Account pubkey containing the deployed WASM component
        #[cfg_attr(feature = "non-pdk", clap(long))]
        bytecode_account: Pubkey,
        /// Optional Base64-encoded input data to be passed to the WASM module
        #[cfg_attr(feature = "non-pdk", clap(long))]
        input_base64: Option<String>,
    } = 8,
}

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: 0,
            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(TimestampMs),
    /// 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(TimestampMs, 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(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_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(0),
            // 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(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(5, 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(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_rex_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_rex_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 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(1000);
        info.update_frequency = UpdateFrequency::LimitedPeriodic(300, 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 - 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 + 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"));
    }
}
