use std::collections::{BTreeMap, HashSet};

use super::key_event::KeyEvent;
use super::running::RunningContractState;
use crate::errors::{Error, InvalidParameters};
use crate::primitives::key_state::{
    AuthenticatedAccountId, EpochId, KeyEventId, KeyForDomain, Keyset,
};
use crate::primitives::thresholds::ProposedThresholdParameters;
use near_account_id::AccountId;
use near_mpc_contract_interface::types::{DomainId, ReconstructionThreshold};
use near_sdk::near;

/// In this state, we reshare the key of every domain onto a new set of participants and threshold.
/// Similar to key generation, we reshare the key of one domain at a time; when we finish resharing
/// for one domain, we move on to the next or transition to the Running state.
///
/// This state is reached by calling vote_new_parameters from the Running state.
///
/// This state keeps the previous running state because:
///  - The previous running state's ThresholdParameters are needed in order to facilitate the
///    possible re-proposal of a new ThresholdParameters, in case the currently proposed set of
///    participants are no longer all online. For tracking the votes we also use the same
///    tracking structure in the running state.
///  - The previous running state's keys are needed to copy the public keys.
///  - We use the previous running state's DomainRegistry.
#[near(serializers=[borsh, json])]
#[derive(Debug)]
#[cfg_attr(feature = "dev-utils", derive(Clone, PartialEq))]
pub struct ResharingContractState {
    pub previous_running_state: RunningContractState,
    pub reshared_keys: Vec<KeyForDomain>,
    pub resharing_key: KeyEvent,
    pub cancellation_requests: HashSet<AuthenticatedAccountId>,
    /// Per-domain `ReconstructionThreshold` updates carried from the accepted
    /// proposal. Applied to the [`DomainRegistry`](crate::primitives::domain::DomainRegistry)
    /// when resharing completes; empty means "keep current per-domain thresholds".
    pub per_domain_thresholds: BTreeMap<DomainId, ReconstructionThreshold>,
}

impl ResharingContractState {
    pub fn previous_keyset(&self) -> &Keyset {
        &self.previous_running_state.keyset
    }

    /// Returns the epoch ID that we would transition into if resharing were completed successfully.
    /// This would increment if we end up voting for a re-proposal.
    pub fn prospective_epoch_id(&self) -> EpochId {
        self.resharing_key.epoch_id()
    }

    /// Casts a vote for a re-proposal. Requires the signer to be a participant of the prospective epoch.
    /// Returns a new [`ResharingContractState`] if all participants of the re-proposal voted for the re-proposal.
    /// Note that transitioning to a new state implicitly requires `threshold` number of votes from participants of the
    /// previous running state.
    pub fn vote_new_parameters(
        &mut self,
        prospective_epoch_id: EpochId,
        proposal: &ProposedThresholdParameters,
    ) -> Result<Option<ResharingContractState>, Error> {
        let expected_prospective_epoch_id = self.prospective_epoch_id().next();
        if prospective_epoch_id != expected_prospective_epoch_id {
            return Err(InvalidParameters::EpochMismatch {
                expected: expected_prospective_epoch_id,
                provided: prospective_epoch_id,
            }
            .into());
        }
        if self
            .previous_running_state
            .process_new_parameters_proposal(proposal)?
        {
            return Ok(Some(ResharingContractState {
                previous_running_state: RunningContractState::new(
                    self.previous_running_state.domains.clone(),
                    self.previous_running_state.keyset.clone(),
                    self.previous_running_state.parameters.clone(),
                    self.previous_running_state.add_domains_votes.clone(),
                ),
                reshared_keys: Vec::new(),
                resharing_key: KeyEvent::new(
                    self.prospective_epoch_id().next(),
                    self.previous_running_state
                        .domains
                        .get_domain_by_index(0)
                        .unwrap()
                        .clone(),
                    proposal.parameters().clone(),
                ),
                cancellation_requests: HashSet::new(),
                per_domain_thresholds: proposal.per_domain_thresholds().clone(),
            }));
        }
        Ok(None)
    }

    /// Starts a new attempt to reshare the key for the current domain.
    /// Returns an Error if the signer is not the leader (the participant with the lowest ID).
    pub fn start(
        &mut self,
        key_event_id: KeyEventId,
        key_event_timeout_blocks: u64,
    ) -> Result<(), Error> {
        self.resharing_key
            .start(key_event_id, key_event_timeout_blocks)
    }

    /// Casts a successfully-reshared vote for the attempt identified by `key_event_id`.
    /// Upon success (a return of Ok(...)), the effect of this method is one of the following:
    ///  - A vote has been collected but we don't have enough votes yet.
    ///  - Everyone has now voted; the state transitions into resharing the key for the next domain.
    ///    (This returns Ok(None) still).
    ///  - Same as the last case, except that all domains' keys have been reshared now, and we
    ///    return Ok(Some(running state)) that the caller should now transition into.
    ///
    /// Fails in the following cases:
    ///  - There is no active key resharing attempt (including if the attempt timed out).
    ///  - The key_event_id corresponds to a different domain, different epoch, or different attempt
    ///    from the current key resharing attempt.
    ///  - The signer is not a participant in the *proposed* set of participants.
    pub fn vote_reshared(
        &mut self,
        key_event_id: KeyEventId,
    ) -> Result<Option<RunningContractState>, Error> {
        let previous_key = self.previous_keyset().domains[self.reshared_keys.len()].clone();
        if self
            .resharing_key
            .vote_success(&key_event_id, previous_key.key.clone())?
        {
            let new_key = KeyForDomain {
                domain_id: key_event_id.domain_id,
                attempt: key_event_id.attempt_id,
                key: previous_key.key,
            };
            self.reshared_keys.push(new_key);
            if let Some(next_domain) = self
                .previous_running_state
                .domains
                .get_domain_by_index(self.reshared_keys.len())
            {
                self.resharing_key = KeyEvent::new(
                    self.prospective_epoch_id(),
                    next_domain.clone(),
                    self.resharing_key.proposed_parameters().clone(),
                );
            } else {
                // Resharing complete: fold the per-domain threshold updates into
                // the registry and store the proposed parameters. The updates live
                // only on this resharing state, so they are structurally dropped
                // here rather than scrubbed off the stored parameters.
                let new_domains = self
                    .previous_running_state
                    .domains
                    .with_threshold_updates(&self.per_domain_thresholds)?;
                return Ok(Some(RunningContractState::new(
                    new_domains,
                    Keyset::new(self.prospective_epoch_id(), self.reshared_keys.clone()),
                    self.resharing_key.proposed_parameters().clone(),
                    self.previous_running_state.add_domains_votes.clone(),
                )));
            }
        }
        Ok(None)
    }

    /// Casts a vote to abort the current key resharing attempt.
    /// After aborting, another call to start() is necessary to start a new attempt.
    /// Returns error if there is no active attempt, or if the signer is not a proposed participant.
    pub fn vote_abort(&mut self, key_event_id: KeyEventId) -> Result<(), Error> {
        self.resharing_key.vote_abort(key_event_id)
    }

    pub fn vote_cancel_resharing(&mut self) -> Result<Option<RunningContractState>, Error> {
        let previous_running_participants = self.previous_running_state.parameters.participants();
        let authenticated_candidate = AuthenticatedAccountId::new(previous_running_participants)?;
        self.cancellation_requests.insert(authenticated_candidate);

        let cancellation_votes_count = self.cancellation_requests.len() as u64;
        let previous_running_threshold = self.previous_running_state.parameters.threshold();

        let threshold_cancellation_votes_reached: bool =
            cancellation_votes_count >= previous_running_threshold.value();

        let running_state = if threshold_cancellation_votes_reached {
            let mut previous_running_state = self.previous_running_state.clone();
            let prospective_epoch_id = self.prospective_epoch_id();
            previous_running_state.previously_cancelled_resharing_epoch_id =
                Some(prospective_epoch_id);

            Some(previous_running_state)
        } else {
            None
        };

        Ok(running_state)
    }

    pub fn is_participant_or_prospective_participant(&self, account_id: &AccountId) -> bool {
        self.previous_running_state
            .is_participant_given_account_id(account_id)
            || self
                .resharing_key
                .proposed_parameters()
                .participants()
                .is_participant_given_account_id(account_id)
    }
}
#[cfg(test)]
pub mod tests {
    use crate::primitives::test_utils::NUM_PROTOCOLS;
    use crate::state::{
        key_event::tests::{Environment, find_leader},
        running::RunningContractState,
    };
    use crate::{
        primitives::{
            domain::AddDomainsVotes,
            key_state::{AttemptId, KeyEventId},
            test_utils::gen_account_id,
            threshold_votes::ThresholdParametersVotes,
            thresholds::{ProposedThresholdParameters, Threshold, ThresholdParameters},
        },
        state::test_utils::{gen_resharing_state, gen_running_state_with_params},
    };
    use near_account_id::AccountId;
    use near_mpc_contract_interface::types::{DomainId, ReconstructionThreshold};
    use rstest::rstest;
    use std::collections::{BTreeMap, BTreeSet};

    fn test_resharing_contract_state_for(num_domains: usize) {
        println!("Testing with {} domains", num_domains);
        let (mut env, mut state) = gen_resharing_state(num_domains);
        let candidates: BTreeSet<AccountId> = state
            .resharing_key
            .proposed_parameters()
            .participants()
            .participants()
            .iter()
            .map(|(aid, _, _)| aid.clone())
            .collect();

        let mut resulting_running_state: Option<RunningContractState> = None;
        for i in 0..num_domains {
            println!("Testing domain {}", i);
            assert!(!state.resharing_key.is_active());
            let first_key_event_id = KeyEventId {
                attempt_id: AttemptId::new(),
                domain_id: state
                    .previous_running_state
                    .domains
                    .get_domain_by_index(i)
                    .unwrap()
                    .id,
                epoch_id: state.prospective_epoch_id(),
            };
            let leader = find_leader(&state.resharing_key);
            for c in &candidates {
                env.set_signer(c);
                // verify that no votes can be cast before the resharing started.
                let _ = state.vote_reshared(first_key_event_id).unwrap_err();
                let _ = state.vote_abort(first_key_event_id).unwrap_err();
                if *c != leader.0 {
                    let _ = state.start(first_key_event_id, 1).unwrap_err();
                } else {
                    // Also check that starting with the wrong KeyEventId fails.
                    let _ = state
                        .start(first_key_event_id.next_attempt(), 1)
                        .unwrap_err();
                }
            }
            // start the resharing; verify that the resharing is for the right epoch and domain ID.
            env.set_signer(&leader.0);
            state.start(first_key_event_id, 0).unwrap();
            let key_event = state.resharing_key.current_key_event_id().unwrap();
            assert_eq!(key_event, first_key_event_id);

            // check that randos can't vote.
            for _ in 0..20 {
                env.set_signer(&gen_account_id());
                let _ = state.vote_reshared(key_event).unwrap_err();
                let _ = state.vote_abort(key_event).unwrap_err();
            }

            // check that timing out will abort the instance
            env.advance_block_height(1);
            assert!(!state.resharing_key.is_active());
            for c in &candidates {
                env.set_signer(c);
                let _ = state.vote_reshared(key_event).unwrap_err();
                let _ = state.vote_abort(key_event).unwrap_err();
                assert!(!state.resharing_key.is_active());
            }

            // assert that votes for a different resharings do not count
            env.set_signer(&leader.0);
            state.start(first_key_event_id.next_attempt(), 0).unwrap();
            let key_event = state.resharing_key.current_key_event_id().unwrap();
            let bad_key_events = [
                KeyEventId::new(
                    key_event.epoch_id,
                    key_event.domain_id,
                    key_event.attempt_id.next(),
                ),
                KeyEventId::new(
                    key_event.epoch_id,
                    DomainId(key_event.domain_id.0 + 1),
                    key_event.attempt_id,
                ),
                KeyEventId::new(
                    key_event.epoch_id.next(),
                    key_event.domain_id,
                    key_event.attempt_id,
                ),
            ];
            for bad_key_event in bad_key_events {
                for c in &candidates {
                    env.set_signer(c);
                    let _ = state.vote_reshared(bad_key_event).unwrap_err();
                    let _ = state.vote_abort(bad_key_event).unwrap_err();
                }
            }
            assert_eq!(state.resharing_key.num_completed(), 0);

            // check that vote_abort immediately causes failure.
            env.advance_block_height(1);
            env.set_signer(&leader.0);
            state.start(key_event.next_attempt(), 0).unwrap();
            let key_event = state.resharing_key.current_key_event_id().unwrap();
            env.set_signer(candidates.iter().next().unwrap());
            state.vote_abort(key_event).unwrap();
            assert!(!state.resharing_key.is_active());

            // assert that valid votes get counted correctly
            env.set_signer(&leader.0);
            state.start(key_event.next_attempt(), 0).unwrap();
            let key_event = state.resharing_key.current_key_event_id().unwrap();
            for (i, c) in candidates.clone().into_iter().enumerate() {
                env.set_signer(&c);
                assert!(resulting_running_state.is_none());
                assert_eq!(state.resharing_key.num_completed(), i);
                resulting_running_state = state.vote_reshared(key_event).unwrap();
                let _ = state.vote_abort(key_event).unwrap_err();
            }
        }

        // assert that the final running state is correct
        let running_state = resulting_running_state.unwrap();
        assert_eq!(
            &running_state.parameters,
            state.resharing_key.proposed_parameters(),
        );
        assert_eq!(running_state.keyset.epoch_id, state.prospective_epoch_id());
        assert_eq!(running_state.keyset.domains, state.reshared_keys);
        assert_eq!(running_state.keyset.domains.len(), num_domains);
        assert_eq!(running_state.domains, state.previous_running_state.domains);
        assert_eq!(
            running_state.parameters_votes,
            ThresholdParametersVotes::default()
        );
        assert_eq!(running_state.add_domains_votes, AddDomainsVotes::default());
    }

    #[rstest]
    #[case(1)]
    #[case(2)]
    #[case(3)]
    #[case(NUM_PROTOCOLS)]
    #[case(2*NUM_PROTOCOLS)]
    fn test_resharing_contract_state(#[case] num_domains: usize) {
        test_resharing_contract_state_for(num_domains);
    }

    #[test]
    fn test_resharing_reproposal() {
        let (mut env, mut state) = gen_resharing_state(NUM_PROTOCOLS);

        // Vote for first domain's key.
        let leader = find_leader(&state.resharing_key);
        env.set_signer(&leader.0);
        let first_key_event_id = KeyEventId {
            attempt_id: AttemptId::new(),
            domain_id: state
                .previous_running_state
                .domains
                .get_domain_by_index(0)
                .unwrap()
                .id,
            epoch_id: state.prospective_epoch_id(),
        };
        state.start(first_key_event_id, 0).unwrap();

        let old_participants = state
            .previous_running_state
            .parameters
            .participants()
            .clone();
        {
            let new_participants = state
                .resharing_key
                .proposed_parameters()
                .participants()
                .participants()
                .clone();
            for (account, _, _) in new_participants {
                env.set_signer(&account);
                state.vote_reshared(first_key_event_id).unwrap();
            }
        }
        assert!(state.reshared_keys.len() == 1);

        // Generate two sets of params:
        //  - old params -> new_params_1 is a valid proposal.
        //  - new_params_1 -> new_params_2 is a valid proposal.
        //  - old params -> new_params_2 is NOT a valid proposal.
        //
        // Reproposing with new_params_1 should succeed, but then reproposing with new_params_2
        // should be rejected, since all re-proposals must be valid against the original.
        let mut new_participants_1 = old_participants.clone();
        // Threshold valid for the grown set (1.5x): old_len sits within
        // [ceil(0.6 * 1.5n), 1.5n].
        let new_threshold = Threshold::new(old_participants.len() as u64);
        new_participants_1.add_random_participants_till_n((old_participants.len() * 3).div_ceil(2));
        let new_participants_2 = new_participants_1
            .subset(new_participants_1.len() - old_participants.len()..new_participants_1.len());
        let new_params_1 = ThresholdParameters::new(new_participants_1, new_threshold).unwrap();
        let new_params_2 = ThresholdParameters::new(new_participants_2, new_threshold).unwrap();
        // Proposals carry an empty (no-change) set of per-domain threshold updates.
        let proposed_1 = ProposedThresholdParameters::new(new_params_1.clone(), BTreeMap::new());
        let proposed_2 = ProposedThresholdParameters::new(new_params_2.clone(), BTreeMap::new());
        state
            .previous_running_state
            .parameters
            .validate_incoming_proposal(&new_params_1)
            .unwrap();
        new_params_1
            .validate_incoming_proposal(&new_params_2)
            .unwrap();
        let _ = state
            .previous_running_state
            .parameters
            .validate_incoming_proposal(&new_params_2)
            .unwrap_err();

        // Reproposing with invalid epoch ID should fail.
        {
            env.set_signer(&old_participants.participants()[0].0);
            let _ = state
                .vote_new_parameters(state.prospective_epoch_id(), &proposed_1)
                .unwrap_err();
            let _ = state
                .vote_new_parameters(state.prospective_epoch_id().next().next(), &proposed_1)
                .unwrap_err();
        }

        // Repropose with new_params_1.
        let mut new_state = None;
        for (account, _, _) in new_params_1.participants().participants() {
            env.set_signer(account);
            assert!(new_state.is_none());
            new_state = state
                .vote_new_parameters(state.prospective_epoch_id().next(), &proposed_1)
                .unwrap();
        }
        // We should've gotten a new resharing state.
        assert!(new_state.is_some());
        let mut new_state = new_state.unwrap();
        // New state should start from the beginning, with the epoch ID bumped.
        assert_eq!(new_state.reshared_keys.len(), 0);
        assert_eq!(
            new_state.resharing_key.epoch_id(),
            state.prospective_epoch_id().next()
        );
        assert_eq!(new_state.resharing_key.proposed_parameters(), &new_params_1);
        assert_eq!(
            new_state.resharing_key.domain_id(),
            state
                .previous_running_state
                .domains
                .get_domain_by_index(0)
                .unwrap()
                .id
        );

        // Repropose with new_params_2. That should fail.
        env.set_signer(&old_participants.participants()[0].0);
        let _ = new_state
            .vote_new_parameters(new_state.prospective_epoch_id().next(), &proposed_2)
            .unwrap_err();
    }

    /// On successful resharing transition, the proposal's
    /// `per_domain_thresholds` updates must be applied to the new
    /// `DomainRegistry`. The updates live only on the proposal /
    /// resharing state, so the stored `RunningContractState.parameters`
    /// (a plain `ThresholdParameters`) cannot carry them at all.
    ///
    /// Pins a participant set with GovernanceThreshold >= 3 (unchanged across the
    /// resharing) and moves the domain to that threshold — the max reconstruction the
    /// cross-domain invariant allows, deterministically distinct from the default 2.
    #[expect(non_snake_case)]
    #[test]
    fn vote_reshared__final_transition__should_apply_threshold_updates_to_registry() {
        // Given a CaitSith domain at the default threshold 2 and a key-refresh proposal
        // moving it to the GovernanceThreshold.
        let mut env = Environment::new(Some(100), None, None);
        let mut running = gen_running_state_with_params(1, 5, 4);
        let current_params = running.parameters.clone();
        let domain_id = running.domains.domains()[0].id;
        let original_threshold = running.domains.domains()[0].reconstruction_threshold;
        let new_threshold = ReconstructionThreshold::new(current_params.threshold().value());
        assert_ne!(new_threshold, original_threshold);
        let mut threshold_updates = BTreeMap::new();
        threshold_updates.insert(domain_id, new_threshold);
        let proposal = ProposedThresholdParameters::new(current_params.clone(), threshold_updates);

        // Drive the proposal to acceptance so we transition into Resharing
        // through the real vote path (which also exercises the fail-fast
        // threshold-update validation in `process_new_parameters_proposal`).
        let prospective_epoch_id = running.prospective_epoch_id();
        let mut state = None;
        for (account, _, _) in proposal.participants().participants() {
            env.set_signer(account);
            state = running
                .vote_new_parameters(prospective_epoch_id, &proposal)
                .unwrap();
        }
        let mut state = state.expect("Should've transitioned into resharing");

        // When all candidates vote-reshared for the (single) domain
        let leader = find_leader(&state.resharing_key);
        env.set_signer(&leader.0);
        let key_event_id = KeyEventId {
            attempt_id: AttemptId::new(),
            domain_id,
            epoch_id: state.prospective_epoch_id(),
        };
        state.start(key_event_id, 0).unwrap();
        let mut new_running = None;
        let candidates: Vec<_> = state
            .resharing_key
            .proposed_parameters()
            .participants()
            .participants()
            .iter()
            .map(|(acc, _, _)| acc.clone())
            .collect();
        for account in &candidates {
            env.set_signer(account);
            new_running = state.vote_reshared(key_event_id).unwrap();
        }

        // Then the new running state's registry carries the updated
        // threshold. (The stored parameters are a plain `ThresholdParameters`
        // and structurally cannot carry pending threshold updates.)
        let new_running = new_running.expect("resharing should have transitioned to Running");
        assert_eq!(
            new_running.domains.domains()[0].reconstruction_threshold,
            new_threshold,
        );
    }

    /// Companion to the test above: two domains end the resharing with different
    /// thresholds. Only Frost is moved (to the GovernanceThreshold); CaitSith keeps 2.
    /// The pinned set keeps GovernanceThreshold >= 3 so the two values differ.
    #[expect(non_snake_case)]
    #[test]
    fn vote_reshared__final_transition__should_apply_distinct_thresholds_per_domain() {
        // Given CaitSith (index 0) and Frost (index 1) at default threshold 2, and a
        // proposal moving only Frost to the GovernanceThreshold.
        let mut env = Environment::new(Some(100), None, None);
        let mut running = gen_running_state_with_params(2, 5, 4);
        let current_params = running.parameters.clone();
        let caitsith_id = running.domains.domains()[0].id;
        let frost_id = running.domains.domains()[1].id;
        let default_threshold = running.domains.domains()[0].reconstruction_threshold;
        let frost_new_threshold = ReconstructionThreshold::new(current_params.threshold().value());
        assert_ne!(frost_new_threshold, default_threshold);
        let mut threshold_updates = BTreeMap::new();
        threshold_updates.insert(frost_id, frost_new_threshold);
        let proposal = ProposedThresholdParameters::new(current_params.clone(), threshold_updates);

        // Drive the proposal to acceptance via the real vote path.
        let prospective_epoch_id = running.prospective_epoch_id();
        let mut state = None;
        for (account, _, _) in proposal.participants().participants() {
            env.set_signer(account);
            state = running
                .vote_new_parameters(prospective_epoch_id, &proposal)
                .unwrap();
        }
        let mut state = state.expect("Should've transitioned into resharing");

        // When every domain is resharing-completed in order.
        let candidates: Vec<_> = state
            .resharing_key
            .proposed_parameters()
            .participants()
            .participants()
            .iter()
            .map(|(acc, _, _)| acc.clone())
            .collect();
        let num_domains = state.previous_running_state.domains.domains().len();
        let mut new_running = None;
        for i in 0..num_domains {
            let domain_id = state
                .previous_running_state
                .domains
                .get_domain_by_index(i)
                .unwrap()
                .id;
            let leader = find_leader(&state.resharing_key);
            env.set_signer(&leader.0);
            let key_event_id = KeyEventId {
                attempt_id: AttemptId::new(),
                domain_id,
                epoch_id: state.prospective_epoch_id(),
            };
            state.start(key_event_id, 0).unwrap();
            for account in &candidates {
                env.set_signer(account);
                new_running = state.vote_reshared(key_event_id).unwrap();
            }
        }

        // Then each domain carries its own threshold: CaitSith keeps 2, Frost holds the update.
        let new_running = new_running.expect("resharing should have transitioned to Running");
        let threshold_for = |id| {
            new_running
                .domains
                .domains()
                .iter()
                .find(|d| d.id == id)
                .unwrap()
                .reconstruction_threshold
        };
        assert_eq!(threshold_for(caitsith_id), default_threshold);
        assert_eq!(threshold_for(frost_id), frost_new_threshold);
    }
}