#![allow(non_snake_case)]

use super::common;
use mpc_contract::{
    MpcContract,
    errors::{Error, InvalidState, TeeError},
    primitives::{
        key_state::EpochId,
        participants::{ParticipantId, ParticipantInfo},
        test_utils::{create_node_id, gen_participants, node_id_for},
        thresholds::{
            GovernanceThreshold, GovernanceThresholdParameters,
            ProposedGovernanceThresholdParameters,
        },
    },
    tee::tee_state::{AttestationSubmissionError, NodeId},
};
use near_mpc_contract_interface::types::{
    Attestation, InitConfig, MockAttestation, ProtocolContractState,
};
use std::collections::BTreeMap;

use assert_matches::assert_matches;
use near_account_id::AccountId;
use near_sdk::{test_utils::VMContextBuilder, testing_env};
use rstest::rstest;
use std::time::Duration;
use test_utils::attestation::mock_dto_dstack_attestation;

const SECOND: Duration = Duration::from_secs(1);
const NANOS_IN_SECOND: u64 = SECOND.as_nanos() as u64;

const DEFAULT_PARTICIPANT_COUNT: usize = 3;
const DEFAULT_THRESHOLD_SIZE: u64 = 2;
const DEFAULT_CONTRACT_PROTOCOL_STATE: ContractProtocolState = ContractProtocolState::Running;

enum ContractProtocolState {
    Running,
    Initializing,
    Resharing,
}

struct TestSetupBuilder {
    participant_count: Option<usize>,
    threshold: Option<u64>,
    init_config: Option<InitConfig>,
    contract_protocol_state: Option<ContractProtocolState>,
}

impl TestSetupBuilder {
    fn new() -> Self {
        Self {
            participant_count: None,
            threshold: None,
            init_config: None,
            contract_protocol_state: None,
        }
    }

    fn with_participant_count(mut self, participant_count: usize) -> Self {
        self.participant_count = Some(participant_count);
        self
    }

    fn with_threshold(mut self, threshold: u64) -> Self {
        self.threshold = Some(threshold);
        self
    }

    fn with_init_config(mut self, init_config: InitConfig) -> Self {
        self.init_config = Some(init_config);
        self
    }

    fn with_tee_upgrade_grace_period_seconds(self, seconds: u64) -> Self {
        self.with_init_config(InitConfig {
            tee_upgrade_deadline_duration_seconds: Some(seconds),
            ..Default::default()
        })
    }

    fn with_contract_protocol_state(
        mut self,
        contract_protocol_state: ContractProtocolState,
    ) -> Self {
        self.contract_protocol_state = Some(contract_protocol_state);
        self
    }

    fn build(self) -> TestSetup {
        let participant_count = self.participant_count.unwrap_or(DEFAULT_PARTICIPANT_COUNT);
        let threshold = self.threshold.unwrap_or(DEFAULT_THRESHOLD_SIZE);
        let contract_protocol_state = self
            .contract_protocol_state
            .unwrap_or(DEFAULT_CONTRACT_PROTOCOL_STATE);

        let participants = gen_participants(participant_count);
        let participants_list = participants.participants().clone();
        let parameters =
            GovernanceThresholdParameters::new(participants, GovernanceThreshold::new(threshold))
                .expect("failed to create threshold parameters");
        let contract = common::init_contract(&parameters, self.init_config);

        let mut setup = TestSetup {
            contract,
            participants_list,
        };

        match contract_protocol_state {
            ContractProtocolState::Running => {}
            ContractProtocolState::Initializing => {
                let participants = setup.participants_list.clone();
                common::transition_to_initializing(&mut setup.contract, &participants);
            }
            ContractProtocolState::Resharing => {
                let all_nodes = setup.get_participant_node_ids();
                let threshold_nodes = all_nodes.iter().take(threshold as usize);

                for node_id in threshold_nodes.clone() {
                    setup.submit_attestation_for_node(
                        node_id,
                        Attestation::Mock(MockAttestation::Valid),
                    );
                }

                for node_id in threshold_nodes {
                    testing_env!(common::participant_context(&node_id.account_id));
                    let proposal = ProposedGovernanceThresholdParameters::new(
                        parameters.clone(),
                        BTreeMap::new(),
                    );
                    setup
                        .contract
                        .vote_new_parameters(EpochId::new(6), proposal.into())
                        .unwrap();
                }

                assert_matches!(setup.contract.state(), ProtocolContractState::Running(_));
            }
        };

        setup
    }
}

struct TestSetup {
    contract: MpcContract,
    participants_list: Vec<(AccountId, ParticipantId, ParticipantInfo)>,
}

impl TestSetup {
    fn submit_attestation_for_node(&mut self, node_id: &NodeId, attestation: Attestation) {
        self.try_submit_attestation_for_node(node_id, attestation)
            .unwrap();
    }

    fn try_submit_attestation_for_node(
        &mut self,
        node_id: &NodeId,
        attestation: Attestation,
    ) -> Result<(), mpc_contract::errors::Error> {
        testing_env!(common::participant_context(&node_id.account_id));
        self.contract
            .submit_participant_info(attestation, node_id.tls_public_key.clone())
            .map(|_| ())
    }

    /// Switches testing context to a given participant at a specific timestamp
    fn with_env(&mut self, account_id: &AccountId, timestamp: u64) {
        testing_env!(
            VMContextBuilder::new()
                .block_timestamp(timestamp)
                .signer_account_id(account_id.clone())
                .predecessor_account_id(account_id.clone())
                .build()
        );
    }

    /// Makes all participants vote for a given code hash at a specific timestamp
    fn vote_with_all_participants(&mut self, hash: [u8; 32], timestamp: u64) {
        for (account_id, _, _) in &self.participants_list.clone() {
            self.with_env(account_id, timestamp);
            self.contract.vote_code_hash(hash.into()).unwrap();
        }
    }
    /// Returns the list of NodeIds for all participants. The
    /// `account_public_key` is a placeholder — only the `account_id` (to set
    /// the signer context) and `tls_public_key` (passed to
    /// `submit_participant_info`) are consumed by these tests.
    fn get_participant_node_ids(&self) -> Vec<NodeId> {
        self.participants_list
            .iter()
            .map(|(account_id, _, participant_info)| {
                create_node_id(account_id, &participant_info.tls_public_key)
            })
            .collect()
    }

    fn create_attestation_with_hash_constraint(hash: [u8; 32]) -> Attestation {
        Attestation::Mock(MockAttestation::WithConstraints {
            mpc_docker_image_hash: Some(hash.into()),
            launcher_docker_compose_hash: None,
            expiry_timestamp_seconds: None,
            expected_measurements: None,
        })
    }
}

fn set_system_time(nano_seconds_since_unix_epoch: u64) {
    testing_env!(
        VMContextBuilder::new()
            .block_timestamp(nano_seconds_since_unix_epoch)
            .build()
    );
}

/// **Test that `submit_participant_info` rejects an attempt by one account to overwrite
/// another account's stored attestation entry**, keyed by TLS public key. Without this
/// gate, any caller could replace a legit participant's attestation and lock them out of
/// every attestation-gated method (DoS via cross-participant overwrite).
#[test]
fn submit_participant_info__should_reject_overwrite_from_other_account() {
    // Given: a running contract with two participants who have submitted their
    // attestations. We will treat `participants_list[0]` as the victim.
    const PARTICIPANT_COUNT: usize = 2;
    const THRESHOLD: u64 = 2;

    let mut setup = TestSetupBuilder::new()
        .with_participant_count(PARTICIPANT_COUNT)
        .with_threshold(THRESHOLD)
        .build();

    let participant_nodes = setup.get_participant_node_ids();
    let victim_node = participant_nodes[0].clone();
    let victim_attestation = Attestation::Mock(MockAttestation::Valid);
    setup.submit_attestation_for_node(&victim_node, victim_attestation);

    let stored_before = setup
        .contract
        .get_attestation(victim_node.tls_public_key.clone())
        .unwrap()
        .expect("victim attestation should be stored");

    // When: an unrelated account submits an attestation that targets the victim's TLS key.
    let attacker_node = create_node_id(
        &"attacker.near".parse().unwrap(),
        &victim_node.tls_public_key,
    );
    let attack_result = setup
        .try_submit_attestation_for_node(&attacker_node, Attestation::Mock(MockAttestation::Valid));

    // Then: the contract rejects the call with the TLS-ownership error and the victim's
    // entry is unchanged.
    assert_matches!(
        &attack_result,
        Err(Error::AttestationSubmission(
            AttestationSubmissionError::TlsKeyOwnedByOtherAccount
        ))
    );
    let stored_after = setup
        .contract
        .get_attestation(victim_node.tls_public_key)
        .unwrap()
        .expect("victim attestation should still be stored");
    assert_eq!(stored_before, stored_after);
}

/// A newcomer stores its first attestation with no deposit (contract-funded storage), so the
/// node's function-call access key can self-onboard.
#[test]
fn submit_participant_info__should_store_new_entry_with_zero_deposit() {
    // Given
    let mut setup = TestSetupBuilder::new().build();
    let newcomer = node_id_for(&"newcomer.near".parse().unwrap());
    testing_env!(common::participant_context(&newcomer.account_id));

    // When
    let result = setup
        .contract
        .submit_participant_info(
            Attestation::Mock(MockAttestation::Valid),
            newcomer.tls_public_key.clone(),
        )
        .map(|_| ());

    // Then
    assert_matches!(&result, Ok(()));
    assert!(
        setup
            .contract
            .get_attestation(newcomer.tls_public_key)
            .unwrap()
            .is_some()
    );
}

/// A current participant re-attesting an existing entry succeeds with no attached deposit, so the
/// node's function-call access key can re-attest.
#[test]
fn submit_participant_info__should_reattest_with_zero_deposit() {
    // Given
    let mut setup = TestSetupBuilder::new().build();
    let node = setup.get_participant_node_ids()[0].clone();
    let attestation = Attestation::Mock(MockAttestation::Valid);
    setup.submit_attestation_for_node(&node, attestation.clone());
    let stored_before = setup
        .contract
        .get_attestation(node.tls_public_key.clone())
        .unwrap()
        .expect("participant attestation should be stored");

    // When: the same participant re-attests with no attached deposit.
    testing_env!(common::participant_context(&node.account_id));
    let result = setup
        .contract
        .submit_participant_info(attestation, node.tls_public_key.clone())
        .map(|_| ());

    // Then: the submission succeeds and the stored entry is unchanged.
    assert_matches!(&result, Ok(()));
    let stored_after = setup
        .contract
        .get_attestation(node.tls_public_key)
        .unwrap()
        .expect("participant attestation should still be stored");
    assert_eq!(stored_before, stored_after);
}

/// Test that a `Dstack` submission is rejected when no verifier is configured.
#[test]
fn submit_participant_info__should_reject_dstack_when_verifier_not_configured() {
    // Given
    let mut setup = TestSetupBuilder::new().build();
    let node = setup.get_participant_node_ids()[0].clone();

    // When
    let result = setup.try_submit_attestation_for_node(&node, mock_dto_dstack_attestation());

    // Then
    assert_matches!(
        &result,
        Err(Error::TeeError(TeeError::VerifierNotConfigured))
    );
}

/// **Test that `clean_tee_status()` is vote-only** — attestations for non-participants
/// remain in `stored_attestations` after the call. Attestation pruning is handled by the
/// separate `clean_invalid_attestations` endpoint.
#[test]
fn clean_tee_status__should_not_touch_attestations() {
    // Given
    const PARTICIPANT_COUNT: usize = 2; // After resharing removed one participant
    const THRESHOLD: u64 = 2;

    // Create contract in Running state with 2 current participants
    let mut setup = TestSetupBuilder::new()
        .with_participant_count(PARTICIPANT_COUNT)
        .with_threshold(THRESHOLD)
        .build();

    // Submit TEE info for current 2 participants (all have valid attestations)
    let valid_attestation = Attestation::Mock(MockAttestation::Valid);
    let participant_nodes = setup.get_participant_node_ids();
    for node_id in &participant_nodes {
        setup.submit_attestation_for_node(node_id, valid_attestation.clone());
    }

    // Add TEE account for someone who is NOT a current participant
    let removed_participant_node = node_id_for(&"removed.participant.near".parse().unwrap());
    setup.submit_attestation_for_node(&removed_participant_node, valid_attestation);

    // Verify initial state: 2 participants but 3 TEE accounts
    const INITIAL_TEE_ACCOUNTS: usize = PARTICIPANT_COUNT + 1; // 2 current + 1 stale
    assert_eq!(
        setup.contract.get_tee_accounts().len(),
        INITIAL_TEE_ACCOUNTS
    );

    assert_matches!(
        setup.contract.state(),
        ProtocolContractState::Running(r)
            if r.parameters.participants.participants.len() == PARTICIPANT_COUNT
    );

    // When: clean_tee_status runs.
    setup.contract.clean_tee_status().unwrap();

    // Then: stored attestations are unchanged — only vote maps are touched by this endpoint.
    assert_eq!(
        setup.contract.get_tee_accounts().len(),
        INITIAL_TEE_ACCOUNTS
    );

    // State should remain Running with same participant count
    assert_matches!(
        setup.contract.state(),
        ProtocolContractState::Running(r)
            if r.parameters.participants.participants.len() == PARTICIPANT_COUNT
    );
}

/// **Test that `clean_invalid_attestations()` prunes expired attestations end-to-end via
/// the public endpoint**, including attestations that belong to current participants.
/// Restores the cleanup-path coverage that lived in the old `clean_tee_status` test.
#[test]
fn clean_invalid_attestations__should_remove_expired_entries() {
    // Given
    const PARTICIPANT_COUNT: usize = 2;
    const THRESHOLD: u64 = 2;
    const EXPIRY_SECONDS: u64 = 1_000;
    const NOW_NS: u64 = 5_000 * NANOS_IN_SECOND;

    let mut setup = TestSetupBuilder::new()
        .with_participant_count(PARTICIPANT_COUNT)
        .with_threshold(THRESHOLD)
        .build();

    let expiring_attestation = Attestation::Mock(MockAttestation::WithConstraints {
        mpc_docker_image_hash: None,
        launcher_docker_compose_hash: None,
        expiry_timestamp_seconds: Some(EXPIRY_SECONDS),
        expected_measurements: None,
    });

    // init_running seeds one mock `Valid` attestation per participant. Overwrite the
    // first participant's entry with an expiring one, and add a brand-new entry for an
    // outsider account.
    let participant_node = setup.get_participant_node_ids()[0].clone();
    setup.submit_attestation_for_node(&participant_node, expiring_attestation.clone());

    let stale_node = node_id_for(&"stale.near".parse().unwrap());
    setup.submit_attestation_for_node(&stale_node, expiring_attestation);

    const EXPECTED_STORED: usize = PARTICIPANT_COUNT + 1; // original mocks + outsider entry
    assert_eq!(setup.contract.get_tee_accounts().len(), EXPECTED_STORED);

    // When: time advances past the expiry and cleanup runs with a generous max_scan.
    set_system_time(NOW_NS);
    let removed = setup.contract.clean_invalid_attestations(100).unwrap();

    // Then: both entries with `expiry_timestamp_seconds` in the past are evicted; the
    // second participant's un-overwritten `Valid` mock remains.
    const EXPECTED_REMOVED: u32 = 2;
    assert_eq!(removed, EXPECTED_REMOVED);
    assert_eq!(
        setup.contract.get_tee_accounts().len(),
        EXPECTED_STORED - EXPECTED_REMOVED as usize
    );
}

/// **Test that `clean_invalid_attestations()` rejects calls outside `Running` state** so
/// that keygen / resharing flows (which may reference not-yet-activated attestations)
/// aren't disrupted.
#[test]
fn clean_invalid_attestations__should_reject_when_not_running() {
    // Given: contract sitting in Initializing state.

    let mut setup = TestSetupBuilder::new()
        .with_contract_protocol_state(ContractProtocolState::Initializing)
        .build();

    // When: the cleanup endpoint is invoked.
    let result = setup.contract.clean_invalid_attestations(100);

    // Then: the call errors without mutating state.
    assert_matches!(
        result,
        Err(Error::InvalidState(InvalidState::ProtocolStateNotRunning))
    );
}

macro_rules! assert_allowed_docker_image_hashes {
    ($test_setup:expr_2021, $blocktime_ns:expr_2021, $expected_value:expr_2021 $(,)?) => {{
        set_system_time($blocktime_ns);

        let mut res: Vec<([u8; 32], Option<u64>)> = $test_setup
            .contract
            .allowed_docker_image_hashes()
            .into_iter()
            .map(|entry| (*entry.image_hash, entry.expiry_timestamp_seconds))
            .collect();

        res.reverse();

        assert_eq!(res, $expected_value);
    }};
}

/// **Test for grace-period expiry of older code hashes**
///
/// Verifies that when participants vote for a new image hash, the older
/// hash remains allowed only until the successor’s grace period deadline.
/// At the exact deadline both old and new hashes are valid, but immediately
/// after, only the latest remains.
#[test]
fn only_latest_hash_after_grace_period() {
    const FIRST_ENTRY_TIME_NS: u64 = NANOS_IN_SECOND; // 1s
    const SECOND_ENTRY_TIME_NS: u64 = 4 * NANOS_IN_SECOND; // 1s
    const GRACE_PERIOD_NS: u64 = 10 * NANOS_IN_SECOND; // 10s

    let mut setup = TestSetupBuilder::new()
        .with_tee_upgrade_grace_period_seconds(GRACE_PERIOD_NS / NANOS_IN_SECOND)
        .build();

    let old_hash = [1; 32];
    let successor_hash = [2; 32];

    let old_hash_expiry = Some((SECOND_ENTRY_TIME_NS + GRACE_PERIOD_NS) / NANOS_IN_SECOND);

    setup.vote_with_all_participants(old_hash, FIRST_ENTRY_TIME_NS);
    assert_allowed_docker_image_hashes!(&setup, FIRST_ENTRY_TIME_NS, &[(old_hash, None)]);
    setup.vote_with_all_participants(successor_hash, SECOND_ENTRY_TIME_NS);
    assert_allowed_docker_image_hashes!(
        &setup,
        SECOND_ENTRY_TIME_NS,
        &[(old_hash, old_hash_expiry), (successor_hash, None)]
    );

    assert_allowed_docker_image_hashes!(
        &setup,
        SECOND_ENTRY_TIME_NS + GRACE_PERIOD_NS,
        &[(old_hash, old_hash_expiry), (successor_hash, None)]
    );
    assert_allowed_docker_image_hashes!(
        &setup,
        SECOND_ENTRY_TIME_NS + GRACE_PERIOD_NS + 1,
        &[(successor_hash, None)]
    );
}

/// **Test for equal-timestamp precedence**
///
/// Ensures that when multiple hashes are inserted at the exact same
/// timestamp, the contract treats the *last inserted* hash as authoritative.
/// After the grace period, only this latest hash is allowed.
#[test]
fn latest_inserted_image_hash_takes_precedence_on_equal_time_stamps() {
    const INITIAL_TIME: u64 = 1;
    const GRACE_PERIOD: u64 = 10;

    let mut setup = TestSetupBuilder::new()
        .with_tee_upgrade_grace_period_seconds(GRACE_PERIOD)
        .build();

    let hash_1 = [1; 32];
    let hash_2 = [2; 32];
    let hash_3 = [3; 32];

    let hashes = [hash_1, hash_2, hash_3];

    for hash in hashes {
        setup.vote_with_all_participants(hash, INITIAL_TIME);
    }
    let superseded_expiry = Some((INITIAL_TIME + GRACE_PERIOD * NANOS_IN_SECOND) / NANOS_IN_SECOND);
    assert_allowed_docker_image_hashes!(
        &setup,
        INITIAL_TIME,
        &[
            (hash_1, superseded_expiry),
            (hash_2, superseded_expiry),
            (hash_3, None),
        ]
    );
    // Jump far in future
    assert_allowed_docker_image_hashes!(&setup, u64::MAX, &[(hash_3, None)]);
}

/// **Test for successor-based grace periods**
///
/// Confirms that a hash’s grace period is tied to the insertion time
/// of its immediate successor, not to the latest hash overall.
/// Each hash expires individually once its successor’s grace period ends.
#[test]
fn hash_grace_period_depends_on_successor_entry_time_not_latest() {
    const FIRST_ENTRY_TIME_NS: u64 = NANOS_IN_SECOND;
    const SECOND_ENTRY_TIME_NS: u64 = 4 * NANOS_IN_SECOND;
    const THIRD_ENTRY_TIME_NS: u64 = 7 * NANOS_IN_SECOND;
    const GRACE_PERIOD_TIME_NS: u64 = 10 * NANOS_IN_SECOND;

    let mut test_setup = TestSetupBuilder::new()
        .with_tee_upgrade_grace_period_seconds(GRACE_PERIOD_TIME_NS / NANOS_IN_SECOND)
        .build();

    let first_code_hash = [1; 32];
    let second_code_hash = [2; 32];
    let third_code_hash = [3; 32];

    let first_hash_expiry = Some((SECOND_ENTRY_TIME_NS + GRACE_PERIOD_TIME_NS) / NANOS_IN_SECOND);
    let second_hash_expiry = Some((THIRD_ENTRY_TIME_NS + GRACE_PERIOD_TIME_NS) / NANOS_IN_SECOND);

    test_setup.vote_with_all_participants(first_code_hash, FIRST_ENTRY_TIME_NS);
    assert_allowed_docker_image_hashes!(
        &test_setup,
        FIRST_ENTRY_TIME_NS,
        &[(first_code_hash, None)]
    );

    test_setup.vote_with_all_participants(second_code_hash, SECOND_ENTRY_TIME_NS);
    assert_allowed_docker_image_hashes!(
        &test_setup,
        SECOND_ENTRY_TIME_NS,
        &[
            (first_code_hash, first_hash_expiry),
            (second_code_hash, None)
        ]
    );

    test_setup.vote_with_all_participants(third_code_hash, THIRD_ENTRY_TIME_NS);
    assert_allowed_docker_image_hashes!(
        &test_setup,
        THIRD_ENTRY_TIME_NS,
        &[
            (first_code_hash, first_hash_expiry),
            (second_code_hash, second_hash_expiry),
            (third_code_hash, None),
        ]
    );

    assert_allowed_docker_image_hashes!(
        &test_setup,
        SECOND_ENTRY_TIME_NS + GRACE_PERIOD_TIME_NS + 1,
        &[
            (second_code_hash, second_hash_expiry),
            (third_code_hash, None),
        ]
    );

    let expiration_second_hash = THIRD_ENTRY_TIME_NS + GRACE_PERIOD_TIME_NS;
    assert_allowed_docker_image_hashes!(
        &test_setup,
        expiration_second_hash,
        &[
            (second_code_hash, second_hash_expiry),
            (third_code_hash, None),
        ]
    );

    assert_allowed_docker_image_hashes!(
        &test_setup,
        expiration_second_hash + 1,
        &[(third_code_hash, None)]
    );
}

/// **Test for indefinite validity of the latest hash**
///
/// Ensures that if no successor hash is ever inserted, the most recent
/// image hash remains valid indefinitely, regardless of how far
/// blockchain time advances.
#[test]
fn latest_image_never_expires_if_its_not_superseded() {
    const START_TIME_SECONDS: u64 = 1;
    const GRACE_PERIOD_SECONDS: u64 = 10;

    let mut test_setup = TestSetupBuilder::new()
        .with_tee_upgrade_grace_period_seconds(GRACE_PERIOD_SECONDS)
        .build();

    let only_image_code_hash = [123; 32];
    test_setup
        .vote_with_all_participants(only_image_code_hash, START_TIME_SECONDS * NANOS_IN_SECOND);

    // Even far in the future, latest remains allowed

    assert_allowed_docker_image_hashes!(&test_setup, u64::MAX, &[(only_image_code_hash, None)]);
}

/// **Test for nodes starting with old but valid image hashes during grace period**
///
/// This test simulates the scenario where new nodes join the network running
/// older Docker image versions that are still within their grace period.
/// It verifies that:
/// 1. Multiple image versions can coexist during their grace periods
/// 2. New nodes can successfully submit attestations with older but valid hashes
/// 3. Nodes running older images remain valid until their specific grace period expires
/// 4. The contract accepts attestations from nodes with any currently allowed hash
///
/// This validates the scenario where nodes may start up with slightly
/// older images after new ones have been voted in, as long as they're still
/// within the tee_upgrade_deadline_duration.
///
/// **Timeline Visualization (Grace Period = 15s):**
/// ```
/// Time:     T=1s   T=4s   T=7s          T=10s         T=19s   T=20s         T=22s   T=23s
///           │      │      │             │             │       │             │       │
/// v1 hash:  ●─────────────────────────────────────────────────X (expires)
/// v2 hash:         ●────────────────────────────────────────────────────────────────X (expires)
/// v3 hash:                ●────────────────────────────────────────────────────────────→ (never expires)
///           │      │      │             │             │       │             │       │
/// Events:   │      │      │             │             │       │             │       │
///          v1     v2     v3          Test all      v1 exp   Check v1      v2 exp   Check v2
///         vote   vote   vote         3 versions    @ T=19s   expired      @ T=22s   expired
///                                    still valid             only v2,v3             only v3
///
/// Grace Period Rules:
/// - v1 expires at: T=4s + 15s + 1s = T=20s
/// - v2 expires at: T=7s + 15s + 1s = T=23s
/// - v3 never expires (no successor hash)
///
/// Note: The +1s ensures we test *after* the grace period deadline has passed.
/// Without it, the hash would still be valid exactly at the deadline timestamp.
/// ```
#[test]
fn nodes_can_start_with_old_valid_hashes_during_grace_period() {
    const INITIAL_TIME_NANOS: u64 = NANOS_IN_SECOND;
    const GRACE_PERIOD_SECONDS: u64 = 15;
    const GRACE_PERIOD_NANOS: u64 = GRACE_PERIOD_SECONDS * NANOS_IN_SECOND;
    const HASH_DEPLOYMENT_INTERVAL_NANOS: u64 = 3 * NANOS_IN_SECOND;

    let mut test_setup = TestSetupBuilder::new()
        .with_tee_upgrade_grace_period_seconds(GRACE_PERIOD_SECONDS)
        .build();

    let hash_v1 = [1; 32]; // Original version
    let hash_v2 = [2; 32]; // Updated version
    let hash_v3 = [3; 32]; // Latest version

    // Deploy three hash versions at 3-second intervals (T=1s, T=4s, T=7s)
    let hashes = [hash_v1, hash_v2, hash_v3];
    let mut deployment_times = Vec::new();
    let mut deployment_time = INITIAL_TIME_NANOS;

    for &hash in hashes.iter() {
        test_setup.vote_with_all_participants(hash, deployment_time);
        deployment_times.push(deployment_time);
        deployment_time += HASH_DEPLOYMENT_INTERVAL_NANOS;
    }

    let v1_expiry = Some((deployment_times[1] + GRACE_PERIOD_NANOS) / NANOS_IN_SECOND);
    let v2_expiry = Some((deployment_times[2] + GRACE_PERIOD_NANOS) / NANOS_IN_SECOND);

    // At T=10s: All three versions should be allowed (within grace periods)
    let test_time_1 = deployment_times[0] + GRACE_PERIOD_NANOS;
    assert_allowed_docker_image_hashes!(
        &test_setup,
        test_time_1,
        &[(hash_v1, v1_expiry), (hash_v2, v2_expiry), (hash_v3, None)]
    );

    // Use existing participant nodes for testing different hash versions
    let node_ids = test_setup.get_participant_node_ids();

    // Test that nodes can submit attestations with all hash versions at T=10s
    // All attestations should succeed during grace period (current time: T=10s)
    for (node, &hash) in node_ids.iter().zip(hashes.iter()) {
        let attestation = TestSetup::create_attestation_with_hash_constraint(hash);
        test_setup.submit_attestation_for_node(node, attestation);
    }

    // Advance to T=19s: hash_v1 should expire (v2 deployed at T=4s + 15s grace = T=19s)
    // Note: v1 expires when its successor's (v2) grace period ends, not when v1's own grace period ends
    let v1_expiry_time = deployment_times[1] + GRACE_PERIOD_NANOS;

    // +1s ensures we're testing *after* expiration occurs - at T=19s the hash is still valid,
    // but at T=20s it has expired and should be filtered out by allowed_docker_image_hashes()
    // T=20s: hash_v1 is expired. Verify that only hash_v2 and hash_v3 are allowed.
    let expected_after_v1_expiry = [hash_v2, hash_v3];
    assert_allowed_docker_image_hashes!(
        &test_setup,
        v1_expiry_time + 1,
        &[(hash_v2, v2_expiry), (hash_v3, None)]
    );

    // Verify that submitting attestation with expired hash_v1 now fails
    let expired_attestation = TestSetup::create_attestation_with_hash_constraint(hash_v1);
    let result = test_setup.try_submit_attestation_for_node(&node_ids[0], expired_attestation);
    assert!(
        result.is_err(),
        "Attestation with expired hash_v1 should fail"
    );

    // Test late-joining nodes at current time T=20s (after hash_v1 expired)
    // Only hash_v2 and hash_v3 should be valid for new nodes
    // Reuse existing node_ids (nodes 2 and 3 since hash_v1 expired)
    for (node, hash) in node_ids[1..].iter().zip(expected_after_v1_expiry.iter()) {
        let late_attestation = TestSetup::create_attestation_with_hash_constraint(*hash);
        test_setup.submit_attestation_for_node(node, late_attestation);
    }

    // Advance to T=22s: hash_v2 should expire (v3 deployed at T=7s + 15s grace = T=22s)
    let v2_expiry_time = deployment_times[2] + GRACE_PERIOD_NANOS;
    assert_allowed_docker_image_hashes!(&test_setup, v2_expiry_time + 1, &[(hash_v3, None)]);

    // Verify that only the latest hash is now accepted
    // Reuse the third node (index 2) for final validation
    let final_attestation = TestSetup::create_attestation_with_hash_constraint(hash_v3);
    // This should succeed since hash_v3 is the only remaining valid hash
    test_setup.submit_attestation_for_node(&node_ids[2], final_attestation);
}

#[rstest]
#[case(ContractProtocolState::Running)]
#[case(ContractProtocolState::Initializing)]
#[case(ContractProtocolState::Resharing)]
fn vote_code_hash_works_in_contract_protocol_states(#[case] state: ContractProtocolState) {
    let mut setup = TestSetupBuilder::new()
        .with_contract_protocol_state(state)
        .build();

    let code_hash = [1; 32];

    setup.vote_with_all_participants(code_hash, 100);
    assert_allowed_docker_image_hashes!(&setup, 100, &[(code_hash, None)]);
}