use sp_runtime::{RuntimeDebug, traits::UniqueSaturatedInto};
use codec::{Encode, Decode};
use scale_info::TypeInfo;
use strum::EnumCount;

use frame_support::{BoundedVec, traits::ConstU32};

use ghost_helpers::networks::NetworkType;
use ghost_traits::bounded_bitmap::{
    BoundedBitmapGenerator, BoundedBitmapReader, BoundedBitmapWriter,
};

use crate::{ELEMENT_MAX_BYTES, ExodusSession, DkgIndex};

type RotationSessions = BoundedVec<RotationSession, ConstU32<{ NetworkType::COUNT as u32 }>>;

#[derive(Default, Encode, Decode, Clone, Copy, PartialEq, Eq, RuntimeDebug, TypeInfo)]
pub struct ReadyDkgAuthorities<AuthoritiesBitmap> {
    authorities_indexes: AuthoritiesBitmap,
    dkg_index: DkgIndex,
}

impl<AuthoritiesBitmap> ReadyDkgAuthorities<AuthoritiesBitmap>
where
	AuthoritiesBitmap: BoundedBitmapReader,
{
    pub fn contains_index<N>(&self, index: N) -> bool
	where N: UniqueSaturatedInto<usize>,
	{
        self.authorities_indexes.contains(index)
    }

    pub fn get_indexes(&self) -> &AuthoritiesBitmap {
        &self.authorities_indexes
    }

    pub fn get_dkg_index(&self) -> DkgIndex {
        self.dkg_index
    }
}

impl<AuthoritiesBitmap, Block>
    From<&PendingDkgAuthorities<AuthoritiesBitmap, Block>>
    for ReadyDkgAuthorities<AuthoritiesBitmap>
where
    Block: Copy,
    AuthoritiesBitmap: BoundedBitmapWriter + BoundedBitmapReader + BoundedBitmapGenerator + Clone,
    for<'a> AuthoritiesBitmap: core::ops::BitAndAssign<&'a AuthoritiesBitmap>,
{
    fn from(pending: &PendingDkgAuthorities<AuthoritiesBitmap, Block>) -> Self {
        Self {
            authorities_indexes: pending.get_indexes().clone(),
            dkg_index: pending.get_dkg_index(),
        }
    }
}

#[derive(Default, Encode, Decode, Clone, Copy, PartialEq, Eq, RuntimeDebug, TypeInfo)]
pub struct RotationSession {
    pub exodus_session: ExodusSession,
    pub network_type: NetworkType
}

#[derive(Default, Encode, Decode, Clone, PartialEq, Eq, RuntimeDebug, TypeInfo)]
pub struct PendingDkgAuthorities<AuthoritiesBitmap, Block> {
    verifying_key: BoundedVec<u8, ConstU32<ELEMENT_MAX_BYTES>>,
    authorities_indexes: AuthoritiesBitmap,
    rotation_sessions: RotationSessions,
    dkg_state: DkgState<Block>,
}

impl<AuthoritiesBitmap, Block> PendingDkgAuthorities<AuthoritiesBitmap, Block>
where
	Block: Copy,
	AuthoritiesBitmap: BoundedBitmapWriter + BoundedBitmapReader + BoundedBitmapGenerator + Clone,
	for<'a> AuthoritiesBitmap: core::ops::BitAndAssign<&'a AuthoritiesBitmap>,
{
    pub fn new_dkg(&mut self, authorities_len: usize) {
        self.authorities_indexes = AuthoritiesBitmap::empty(authorities_len);
        self.dkg_state.restart();
        self.rotation_sessions = Default::default();
    }

    pub fn set_block(&mut self, block: Block) {
        self.dkg_state.set_block(block);
    }

    pub fn set_verifying_key(
        &mut self,
        verifying_key: BoundedVec<u8, ConstU32<ELEMENT_MAX_BYTES>>,
    ) {
        self.verifying_key = verifying_key;
    }

    pub fn nullify_indexes(&mut self) {
        self.verifying_key = Default::default();
        self.authorities_indexes.nullify();
    }

    pub fn restart_dkg(&mut self) {
        self.verifying_key = Default::default();
        self.authorities_indexes.nullify();
        self.dkg_state.restart();
        self.rotation_sessions = Default::default();
    }

    pub fn next_phase(&mut self) {
        self.dkg_state.next_phase()
    }

    #[inline(always)]
    pub fn intersect_indexes(&mut self, mask: &AuthoritiesBitmap) {
        self.authorities_indexes &= mask;
    }

    pub fn is_dkg_pending(&self) -> bool {
        self.dkg_state.is_dkg_pending()
    }

    pub fn is_dkg_vacant(&self) -> bool {
        self.dkg_state.is_dkg_vacant()
    }

    pub fn get_dkg_index(&self) -> DkgIndex {
        self.dkg_state.get_dkg_index()
    }

    pub fn get_block(&self) -> Block {
        self.dkg_state.get_block()
    }

    pub fn get_phase(&self) -> DkgPhase {
        self.dkg_state.get_phase()
    }

    pub fn get_verifying_key(
        &self
    ) -> BoundedVec<u8, ConstU32<ELEMENT_MAX_BYTES>> {
        self.verifying_key.clone()
    }

    pub fn is_zero_phase(&self) -> bool {
        self.dkg_state.get_phase() == DkgPhase::Round0
    }

    pub fn is_first_phase(&self) -> bool {
        self.dkg_state.get_phase() == DkgPhase::Round1
    }

    pub fn is_second_phase(&self) -> bool {
        self.dkg_state.get_phase() == DkgPhase::Round2
    }

    pub fn is_third_phase(&self) -> bool {
        self.dkg_state.get_phase() == DkgPhase::Round3
    }

    pub fn is_fourth_phase(&self) -> bool {
        self.dkg_state.get_phase() == DkgPhase::Round4
    }

    pub fn is_fifth_phase(&self) -> bool {
        self.dkg_state.get_phase() == DkgPhase::Round5
    }

    pub fn is_sixth_phase(&self) -> bool {
        self.dkg_state.get_phase() == DkgPhase::Round6
    }

    pub fn is_dkg_finalized(&self) -> bool {
        self.dkg_state.get_phase() == DkgPhase::Finalized
    }

    pub fn contains_index<N>(&self, index: N) -> bool
	where N: UniqueSaturatedInto<usize>,
	{
        self.authorities_indexes.contains(index)
    }

    pub fn insert_index<N>(&mut self, i: N) -> Option<N>
	where N: Copy + TryInto<usize>
	{
        self.authorities_indexes.insert(i)
    }

    pub fn try_push_rotation_session(
        &mut self,
        exodus_session: ExodusSession,
        network_type: NetworkType,
    ) -> Option<()> {
        let rotation_session = RotationSession {
            exodus_session,
            network_type,
        };

        self.rotation_sessions.try_push(rotation_session).ok()
    }

    pub fn count_ones<N>(&self) -> N
	where u32: UniqueSaturatedInto<N>,
	{
        self.authorities_indexes.count_ones()
    }

    pub fn highest_bit<N>(&self) -> Option<N>
	where usize: UniqueSaturatedInto<N>,
	{
        self.authorities_indexes.highest_bit()
    }

    pub fn get_indexes(&self) -> &AuthoritiesBitmap {
        &self.authorities_indexes
    }

    pub fn iter_rotation_sessions(&self) -> impl Iterator<Item = &RotationSession> {
        self.rotation_sessions.iter()
    }

}

#[derive(Default, Encode, Decode, Clone, Copy, PartialEq, Eq, RuntimeDebug, TypeInfo)]
pub struct DkgState<Block> {
    block: Block,
    index: DkgIndex,
    phase: DkgPhase,
}

impl<Block: Copy> DkgState<Block> {
    fn is_dkg_pending(&self) -> bool { self.phase.is_dkg_pending() }
    fn is_dkg_vacant(&self) -> bool { self.phase.is_dkg_vacant() }

    fn get_dkg_index(&self) -> DkgIndex { self.index }
    fn get_phase(&self) -> DkgPhase { self.phase }
    fn get_block(&self) -> Block { self.block }

    fn restart(&mut self) {
        self.phase = Default::default();
        self.index = self.index.checked_add(1)
            .unwrap_or_default();
    }

    fn next_phase(&mut self) {
        self.phase.next_phase()
    }

    fn set_block(&mut self, block: Block) {
        self.block = block;
    }
}

#[derive(Default, Encode, Decode, Clone, Copy, PartialEq, Eq, RuntimeDebug, TypeInfo)]
pub enum DkgPhase {
    #[default]
    Vacant,
    Round0,
    Round1,
    Round2,
    Round3,
    Round4,
    Round5,
    Round6,
    Round7,
    Finalized,
}

impl DkgPhase {
    fn is_dkg_pending(&self) -> bool {
        match &self {
            Self::Finalized | Self::Round7 => false,
            _ => true
        }
    }

    fn is_dkg_vacant(&self) -> bool {
        match &self {
            Self::Vacant => true,
            _ => false,
        }
    }

    fn next_phase(&mut self) {
        *self = match *self {
            Self::Vacant => Self::Round0,
            Self::Round0 => Self::Round1,
            Self::Round1 => Self::Round2,
            Self::Round2 => Self::Round3,
            Self::Round3 => Self::Round4,
            Self::Round4 => Self::Round5,
            Self::Round5 => Self::Round6,
            Self::Round6 => Self::Round7,
            Self::Round7 => Self::Finalized,
            Self::Finalized => Self::Finalized,
        };
    }
}
