wavs-types 2.0.0-alpha.6

use alloy_primitives::{hex, LogData};
use iri_string::types::UriString;
use semver::Version;
use serde::{Deserialize, Serialize};
use std::collections::{BTreeMap, BTreeSet};
use std::num::{NonZeroU32, NonZeroU64};
use std::str::FromStr;
use utoipa::ToSchema;
use wasm_pkg_common::package::PackageRef;

use crate::{ByteArray, ComponentDigest, ServiceDigest, Timestamp};

use super::{ChainKey, ServiceId, WorkflowId};

/// Service validation is a runtime check, and depends on:
///
/// 1. All service handlers on a given chain use the same service manager
/// 2. All service managers on non-source chains properly mirror the operator set of the source
/// 3. All components are legitimate (e.g. can be downloaded, match the provided digest, execute as expected, etc.)
#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, ToSchema)]
#[serde(rename_all = "snake_case")]
pub struct Service {
    /// This is any utf-8 string, for human-readable display.
    pub name: String,

    /// We support multiple workflows in one service with unique service-scoped IDs.
    pub workflows: BTreeMap<WorkflowId, Workflow>,

    pub status: ServiceStatus,

    pub manager: ServiceManager,
}

impl Service {
    // this is only used for local/tests, but we want to keep it consistent
    pub fn hash(&self) -> anyhow::Result<ServiceDigest> {
        let service_bytes = serde_json::to_vec(self)?;
        Ok(ServiceDigest::hash(&service_bytes))
    }

    pub fn id(&self) -> ServiceId {
        ServiceId::from(&self.manager)
    }
}

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, ToSchema, PartialOrd, Ord)]
#[serde(rename_all = "snake_case")]
pub enum ServiceManager {
    Evm {
        chain: ChainKey,
        #[schema(value_type = String)]
        address: alloy_primitives::Address,
    },
    Cosmos {
        chain: ChainKey,
        #[schema(value_type = String)]
        address: layer_climb_address::CosmosAddr,
    },
}

impl From<&ServiceManager> for ServiceId {
    fn from(manager: &ServiceManager) -> Self {
        match manager {
            ServiceManager::Evm { chain, address } => {
                let mut bytes = Vec::new();
                bytes.extend_from_slice(b"evm");
                bytes.extend_from_slice(chain.to_string().as_bytes());
                bytes.extend_from_slice(address.as_slice());
                ServiceId::hash(bytes)
            }
            ServiceManager::Cosmos { chain, address } => {
                let mut bytes = Vec::new();
                bytes.extend_from_slice(b"cosmos");
                bytes.extend_from_slice(chain.to_string().as_bytes());
                bytes.extend_from_slice(&address.to_vec());
                ServiceId::hash(bytes)
            }
        }
    }
}

impl ServiceManager {
    pub fn chain(&self) -> &ChainKey {
        match self {
            ServiceManager::Evm { chain, .. } => chain,
            ServiceManager::Cosmos { chain, .. } => chain,
        }
    }

    pub fn evm_address_unchecked(&self) -> alloy_primitives::Address {
        match self {
            ServiceManager::Evm { address, .. } => *address,
            ServiceManager::Cosmos { .. } => {
                panic!("ServiceManager is not EVM type, cannot get EVM address")
            }
        }
    }

    pub fn cosmos_address_unchecked(&self) -> layer_climb_address::CosmosAddr {
        match self {
            ServiceManager::Cosmos { address, .. } => address.clone(),
            ServiceManager::Evm { .. } => {
                panic!("ServiceManager is not Cosmos type, cannot get Cosmos address")
            }
        }
    }
}

impl Service {
    pub fn new_simple(
        name: Option<String>,
        trigger: Trigger,
        source: ComponentSource,
        submit: Submit,
        manager: ServiceManager,
    ) -> Self {
        let workflow_id = WorkflowId::default();

        let workflow = Workflow {
            trigger,
            component: Component::new(source),
            submit,
        };

        let workflows = BTreeMap::from([(workflow_id, workflow)]);

        Self {
            name: name.unwrap_or_else(|| "Unknown".to_string()),
            workflows,
            status: ServiceStatus::Active,
            manager,
        }
    }
}

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, ToSchema)]
#[serde(rename_all = "snake_case")]
pub struct Component {
    pub source: ComponentSource,

    // What permissions this component has.
    // These are currently not enforced, you can pass in Default::default() for now
    pub permissions: Permissions,

    /// The maximum amount of compute metering to allow for a single component execution
    /// If not supplied, will be `Workflow::DEFAULT_FUEL_LIMIT`
    pub fuel_limit: Option<u64>,

    /// The maximum amount of time to allow for a single component execution, in seconds
    /// If not supplied, default will be `Workflow::DEFAULT_TIME_LIMIT_SECONDS`
    pub time_limit_seconds: Option<u64>,

    /// Key-value pairs that are accessible in the components via host bindings.
    pub config: BTreeMap<String, String>,

    /// External env variable keys to be read from the system host on execute (i.e. API keys).
    /// Must be prefixed with `WAVS_ENV_`.
    pub env_keys: BTreeSet<String>,
}

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, ToSchema)]
#[serde(rename_all = "snake_case")]
pub enum ComponentSource {
    /// The wasm bytecode provided at fixed url, digest provided to ensure no tampering
    Download {
        #[schema(value_type = String)]
        uri: UriString,
        digest: ComponentDigest,
    },
    /// The wasm bytecode downloaded from a standard registry, digest provided to ensure no tampering
    Registry {
        #[serde(flatten)]
        registry: Registry,
    },
    /// An already deployed component
    Digest(ComponentDigest),
}

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, ToSchema)]
pub struct Registry {
    pub digest: ComponentDigest,
    /// Optional domain to use for a registry (such as ghcr.io)
    /// if default of wa.dev (or whatever wavs uses in the future)
    /// is not desired by user
    pub domain: Option<String>,
    /// Optional semver value, if absent then latest is used
    #[schema(value_type = Option<String>)]
    pub version: Option<Version>,
    /// Package identifier of form <namespace>:<packagename>
    #[schema(value_type = String)]
    pub package: PackageRef,
}

impl ComponentSource {
    pub fn digest(&self) -> &ComponentDigest {
        match self {
            ComponentSource::Download { digest, .. } => digest,
            ComponentSource::Registry { registry } => &registry.digest,
            ComponentSource::Digest(digest) => digest,
        }
    }
}

// FIXME: happy for a better name.
/// This captures the triggers we listen to, the components we run, and how we submit the result
#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, ToSchema)]
#[serde(rename_all = "snake_case")]
pub struct Workflow {
    /// The trigger that fires this workflow
    pub trigger: Trigger,

    /// The component to run when the trigger fires
    pub component: Component,

    /// How to submit the result of the component.
    pub submit: Submit,
}

impl Workflow {
    pub const DEFAULT_FUEL_LIMIT: u64 = u64::MAX;
    pub const DEFAULT_TIME_LIMIT_SECONDS: u64 = u64::MAX;
}

// The TriggerManager reacts to these triggers
#[derive(Hash, Serialize, Deserialize, Clone, Debug, PartialEq, Eq, ToSchema)]
#[serde(rename_all = "snake_case")]
pub enum Trigger {
    // A contract that emits an event
    CosmosContractEvent {
        #[schema(value_type = String)]
        address: layer_climb_address::CosmosAddr,
        chain: ChainKey,
        event_type: String,
    },
    EvmContractEvent {
        #[schema(value_type = String)]
        address: alloy_primitives::Address,
        chain: ChainKey,
        event_hash: ByteArray<32>,
    },
    BlockInterval {
        /// The chain to use for the block interval
        chain: ChainKey,
        /// Number of blocks to wait between each execution
        #[schema(value_type = u32)]
        n_blocks: NonZeroU32,
        /// Optional start block height indicating when the interval begins.
        #[schema(value_type = Option<u64>)]
        start_block: Option<NonZeroU64>,
        /// Optional end block height indicating when the interval begins.
        #[schema(value_type = Option<u64>)]
        end_block: Option<NonZeroU64>,
    },
    Cron {
        /// A cron expression defining the schedule for execution.
        schedule: String,
        /// Optional start time (timestamp in nanoseconds) indicating when the schedule begins.
        start_time: Option<Timestamp>,
        /// Optional end time (timestamp in nanoseconds) indicating when the schedule ends.
        end_time: Option<Timestamp>,
    },
    // not a real trigger, just for testing
    Manual,
}

/// The data that came from the trigger and is passed to the component after being converted into the WIT-friendly type
#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, ToSchema)]
pub enum TriggerData {
    CosmosContractEvent {
        /// The address of the contract that emitted the event
        #[schema(value_type = String)]
        contract_address: layer_climb_address::CosmosAddr,
        /// The chain where the event was emitted
        chain: ChainKey,
        /// The data that was emitted by the contract
        #[schema(value_type = Object)]
        event: cosmwasm_std::Event,
        /// The block height where the event was emitted
        block_height: u64,
        /// The index of the event in this block, required for unique identification
        event_index: u64,
    },
    EvmContractEvent {
        /// The chain where the event was emitted
        chain: ChainKey,
        /// The address of the contract that emitted the event
        #[schema(value_type = String)]
        contract_address: alloy_primitives::Address,
        /// The log data
        #[schema(value_type = Object)]
        log_data: LogData,
        /// The transaction hash where the event was emitted
        #[schema(value_type = String)]
        tx_hash: alloy_primitives::TxHash,
        /// The block height where the event was emitted
        block_number: u64,
        /// The index of the log in the block
        log_index: u64,
        // these are all optional because they may not be present in the log and we don't need them
        /// Hash of the block the transaction that emitted this log was mined in
        #[schema(value_type = String)]
        block_hash: alloy_primitives::B256,
        /// The timestamp of the block containing this event, as proposed in https://github.com/ethereum/execution-apis/issues/295
        /// This field is optional since nodes are not required to include it in event logs.
        /// If not provided, applications may need to fetch the block header directly to obtain the timestamp.
        block_timestamp: Option<u64>,
        /// Index of the Transaction in the block
        tx_index: u64,
    },
    BlockInterval {
        /// The chain where the blocks are checked
        chain: ChainKey,
        /// The block height where the event was emitted
        block_height: u64,
    },
    Cron {
        /// The trigger time
        trigger_time: Timestamp,
    },
    Raw(Vec<u8>),
}

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

impl TriggerData {
    pub fn new_raw(data: impl AsRef<[u8]>) -> Self {
        TriggerData::Raw(data.as_ref().to_vec())
    }

    pub fn trigger_type(&self) -> &str {
        match self {
            TriggerData::CosmosContractEvent { .. } => "cosmos_contract_event",
            TriggerData::EvmContractEvent { .. } => "evm_contract_event",
            TriggerData::BlockInterval { .. } => "block_interval",
            TriggerData::Cron { .. } => "cron",
            TriggerData::Raw(_) => "manual",
        }
    }

    pub fn chain(&self) -> Option<&ChainKey> {
        match self {
            TriggerData::CosmosContractEvent { chain, .. }
            | TriggerData::EvmContractEvent { chain, .. }
            | TriggerData::BlockInterval { chain, .. } => Some(chain),
            TriggerData::Cron { .. } | TriggerData::Raw(_) => None,
        }
    }
}

/// A bundle of the trigger and the associated data needed to take action on it
#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, bincode::Decode, bincode::Encode)]
pub struct TriggerAction {
    #[bincode(with_serde)]
    /// Identify which trigger this came from
    pub config: TriggerConfig,

    #[bincode(with_serde)]
    /// The data that came from the trigger
    pub data: TriggerData,
}

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)]
// Trigger with metadata so it can be identified in relation to services and workflows
pub struct TriggerConfig {
    pub service_id: ServiceId,
    pub workflow_id: WorkflowId,
    pub trigger: Trigger,
}

// TODO - rename this? Trigger is a noun, Submit is a verb.. feels a bit weird
#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, ToSchema)]
#[serde(rename_all = "snake_case")]
pub enum Submit {
    // useful for when the component just does something with its own state
    None,
    Aggregator {
        /// The aggregator endpoint
        url: String,
        /// component dynamically determines the destination
        component: Box<Component>,
        signature_kind: SignatureKind,
    },
}

/// Defines the signature configuration for cryptographic operations in WAVS.
///
/// This struct separates the cryptographic algorithm from the message formatting
/// to provide flexibility in signature schemes while maintaining compatibility
/// across different blockchain ecosystems.
///
/// ## Why Separate Algorithm and Prefix?
///
/// The separation of `algorithm` and `prefix` serves several important purposes:
///
/// 1. **Algorithm Independence**: The same cryptographic algorithm (e.g., secp256k1)
///    can be used with different message formatting schemes. This allows the same
///    private key to work across different contexts.
///
/// 2. **Ethereum Compatibility**: Some signatures need EIP-191 prefixing for
///    Ethereum compatibility, while others work with raw message hashes. The
///    optional prefix allows both modes.
///
/// 3. **Future Extensibility**: As new signature algorithms (BLS12-381, Ed25519, etc.)
///    and prefix schemes are added, this structure can accommodate them without
///    breaking changes.
#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, ToSchema)]
pub struct SignatureKind {
    /// The cryptographic algorithm used for signature generation and verification.
    ///
    /// This determines the elliptic curve and mathematical operations used,
    /// but not how the message is formatted before signing.
    pub algorithm: SignatureAlgorithm,

    /// Optional message prefix scheme applied before signing.
    ///
    /// When `Some(prefix)`, the message is formatted according to the specified
    /// scheme (e.g., EIP-191 for Ethereum compatibility). When `None`, the raw
    /// message hash is signed directly.
    pub prefix: Option<SignaturePrefix>,
}

impl SignatureKind {
    pub fn evm_default() -> Self {
        Self {
            algorithm: SignatureAlgorithm::Secp256k1,
            prefix: Some(SignaturePrefix::Eip191),
        }
    }
}

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, ToSchema)]
#[serde(rename_all = "snake_case")]
pub enum SignatureAlgorithm {
    Secp256k1,
    // Future: Bls12381, Ed25519, Secp256r1, etc.
}

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, ToSchema)]
#[serde(rename_all = "snake_case")]
pub enum SignaturePrefix {
    Eip191,
}

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, Copy, ToSchema)]
#[serde(rename_all = "snake_case")]
pub enum ServiceStatus {
    Active,
    // Service is paused, no workflows will be executed
    // however the service can still be queried for AVS Key etc.
    Paused,
}

impl FromStr for ServiceStatus {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_lowercase().as_str() {
            "active" => Ok(ServiceStatus::Active),
            "paused" => Ok(ServiceStatus::Paused),
            _ => Err(anyhow::anyhow!("Invalid service status: {}", s)),
        }
    }
}

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, ToSchema)]
#[serde(default, rename_all = "snake_case")]
#[derive(Default)]
pub struct Permissions {
    /// If it can talk to http hosts on the network
    pub allowed_http_hosts: AllowedHostPermission,
    /// If it can write to it's own local directory in the filesystem
    pub file_system: bool,
}

#[test]
fn permission_defaults() {
    let permissions_json: Permissions = serde_json::from_str("{}").unwrap();
    let permissions_default: Permissions = Permissions::default();

    assert_eq!(permissions_json, permissions_default);
    assert_eq!(
        permissions_default.allowed_http_hosts,
        AllowedHostPermission::None
    );
    assert!(!permissions_default.file_system);
}

// TODO: remove / change defaults?

#[derive(Serialize, Deserialize, Clone, Default, Debug, PartialEq, Eq, ToSchema)]
#[serde(rename_all = "snake_case")]
pub enum AllowedHostPermission {
    All,
    Only(Vec<String>),
    #[default]
    None,
}

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)]
#[serde(default, rename_all = "snake_case")]
#[derive(Default)]
pub struct WasmResponse {
    #[serde(with = "hex")]
    pub payload: Vec<u8>,
    pub ordering: Option<u64>,
}

// TODO - these shouldn't be needed in main code... gate behind `debug_assertions`
// will need to go through use-cases of `test-utils`, maybe move into layer-tests or something
mod test_ext {
    use std::{
        collections::{BTreeMap, BTreeSet},
        num::NonZeroU32,
    };

    use crate::{
        ByteArray, ChainKey, ChainKeyError, ComponentSource, ServiceId, WorkflowId, WorkflowIdError,
    };

    use super::{Component, Trigger, TriggerConfig};

    impl Component {
        pub fn new(source: ComponentSource) -> Component {
            Self {
                source,
                permissions: Default::default(),
                fuel_limit: None,
                time_limit_seconds: None,
                config: BTreeMap::new(),
                env_keys: BTreeSet::new(),
            }
        }
    }

    impl Trigger {
        pub fn cosmos_contract_event(
            address: layer_climb_address::CosmosAddr,
            chain: impl TryInto<ChainKey, Error = ChainKeyError>,
            event_type: impl ToString,
        ) -> Self {
            Trigger::CosmosContractEvent {
                address,
                chain: chain.try_into().unwrap(),
                event_type: event_type.to_string(),
            }
        }
        pub fn evm_contract_event(
            address: alloy_primitives::Address,
            chain: impl TryInto<ChainKey, Error = ChainKeyError>,
            event_hash: ByteArray<32>,
        ) -> Self {
            Trigger::EvmContractEvent {
                address,
                chain: chain.try_into().unwrap(),
                event_hash,
            }
        }
    }

    impl TriggerConfig {
        pub fn cosmos_contract_event(
            service_id: ServiceId,
            workflow_id: impl TryInto<WorkflowId, Error = WorkflowIdError>,
            contract_address: layer_climb_address::CosmosAddr,
            chain: impl TryInto<ChainKey, Error = ChainKeyError>,
            event_type: impl ToString,
        ) -> Self {
            Self {
                service_id,
                workflow_id: workflow_id.try_into().unwrap(),
                trigger: Trigger::cosmos_contract_event(contract_address, chain, event_type),
            }
        }

        pub fn evm_contract_event(
            service_id: ServiceId,
            workflow_id: impl TryInto<WorkflowId, Error = WorkflowIdError>,
            contract_address: alloy_primitives::Address,
            chain: impl TryInto<ChainKey, Error = ChainKeyError>,
            event_hash: ByteArray<32>,
        ) -> Self {
            Self {
                service_id,
                workflow_id: workflow_id.try_into().unwrap(),
                trigger: Trigger::evm_contract_event(contract_address, chain, event_hash),
            }
        }

        pub fn block_interval_event(
            service_id: ServiceId,
            workflow_id: impl TryInto<WorkflowId, Error = WorkflowIdError>,
            chain: impl TryInto<ChainKey, Error = ChainKeyError>,
            n_blocks: NonZeroU32,
        ) -> Self {
            Self {
                service_id,
                workflow_id: workflow_id.try_into().unwrap(),
                trigger: Trigger::BlockInterval {
                    chain: chain.try_into().unwrap(),
                    n_blocks,
                    start_block: None,
                    end_block: None,
                },
            }
        }

        #[cfg(test)]
        pub fn manual(
            service_id: ServiceId,
            workflow_id: impl TryInto<WorkflowId, Error = WorkflowIdError>,
        ) -> Self {
            Self {
                service_id,
                workflow_id: workflow_id.try_into().unwrap(),
                trigger: Trigger::Manual,
            }
        }
    }
}