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

use rialo_limits::max_rex_output_serialized_bytes;
use serde::{Deserialize, Serialize};

use crate::{
    rex_info::{RexInfo, TimestampMs},
    websocket_op::WebSocketOperation,
};

/// Configuration for REX duties.
///
/// This struct defines the parameters that control how REX duties are scheduled,
/// executed, and finalized within the system.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct RexDutyConfig {
    /// The number of validators assigned to each REX duty.
    pub validators_per_duty: u32,
    /// The time in milliseconds needed for the fulfillment of a REX request.
    pub request_delay_ms: TimestampMs,
    /// Max size of REX output in bytes.
    pub max_output_size: u32,
    /// Whether this duty is ASAP (execute as soon as possible) rather than scheduled.
    pub is_asap: bool,
    /// WebSocket operation information needed to update the WebSocketRegistry with the
    /// on-chain connection state when this duty completes.
    pub websocket_op: Option<WebSocketOperation>,
}

impl Default for RexDutyConfig {
    fn default() -> Self {
        Self {
            // 3-validator committees gives us 99.996% success rate leveraging TEE guarantees
            validators_per_duty: Self::DEFAULT_VALIDATORS_PER_DUTY,
            // Time to execute a REX task, determined experimentally
            request_delay_ms: Self::DEFAULT_REQUEST_DELAY_MS,
            // Safe size for the default number of validators.
            max_output_size: Self::DEFAULT_MAX_OUTPUT_SIZE,
            is_asap: false,
            websocket_op: None,
        }
    }
}

impl RexDutyConfig {
    /// Constructs an instance based on RexInfo.
    pub fn from_rex_info(rex_info: &RexInfo) -> Self {
        Self {
            validators_per_duty: rex_info.validators_per_duty,
            request_delay_ms: rex_info.request_delay_ms,
            max_output_size: max_rex_output_serialized_bytes(rex_info.validators_per_duty) as u32,
            is_asap: rex_info.is_asap(),
            websocket_op: rex_info.websocket_op(),
        }
    }

    /// Safe max size limit for the default number of validators.
    pub const DEFAULT_MAX_OUTPUT_SIZE: u32 = 10 * 1024;

    /// A reasonably safe value to make one-off REX requests succeed with >22 bits of security.
    /// This value is based on the example from the Rialo whitepaper, that is N = 130, D = 43.
    pub const ONE_OFF_VALIDATORS_PER_DUTY: u32 = 13;

    /// A small value for REX requests that is reasonably cheap.
    /// Note that by itself it doesn't provide a response reliably enough
    /// and needs retries in case of failure.
    pub const DEFAULT_VALIDATORS_PER_DUTY: u32 = 5;

    /// Most REX requests should complete within a few blocks.
    pub const DEFAULT_REQUEST_DELAY_MS: TimestampMs = 300; // 0.3s

    /// The minimum allowed delay for REX requests, in milliseconds.
    pub const MIN_REX_REQUEST_DELAY: TimestampMs = 100; // 0.1s

    /// The maximum allowed delay for REX requests, in milliseconds.
    pub const MAX_REX_REQUEST_DELAY_MS: TimestampMs = 240_000; // 240s

    /// Default maximum number of ASAP duties to process in a single consensus round
    pub const DEFAULT_MAX_ASAP_DUTIES_PER_COMMIT: usize = 5;

    /// Default timeout for ASAP duties in milliseconds
    pub const DEFAULT_ASAP_TIMEOUT_MS: TimestampMs = 5_000; // 5s

    /// Sets the REX request delay (in milliseconds).
    ///
    /// # Arguments
    ///
    /// * `delay` - The time in milliseconds to wait before issuing a REX request.
    pub fn set_rex_request_delay_ms(mut self, delay: TimestampMs) -> Self {
        self.request_delay_ms = delay;
        self
    }

    /// Sets the desired number of validators per duty.
    pub fn set_validators_per_duty(mut self, validators_per_duty: u32) -> RexDutyConfig {
        self.validators_per_duty = validators_per_duty;
        self
    }
}
