use codec::{Encode, Decode};

use sp_std::{vec::Vec, collections::btree_map::BTreeMap};
use sp_runtime::{
    RuntimeDebug, RuntimeAppPublic,
    traits::{AtLeast32BitUnsigned, Saturating},
    offchain::storage::StorageValueRef,
};

use super::{DkgPackage, ExodusPackage, ExodusOk, PackageContext};
use crate::{AuthIndex, ExodusSession, ExodusHash, ExodusError};

fn create_storage_key(prefix: &[u8], index: &[u8], network: &[u8]) -> Vec<u8> {
    let mut key = crate::DB_STORAGE_PREFIX.to_vec();
    key.extend(prefix);
    key.extend(index);
    key.extend(network);
    key
}

#[derive(Default, Encode, Decode, Clone, PartialEq, Eq, RuntimeDebug)]
pub struct PersistentData<BlockNumber, DataType> {
    pub release_block: BlockNumber,
    pub data: DataType,
}

#[derive(Encode, Decode, Clone, PartialEq, Eq, RuntimeDebug)]
pub enum DkgStorageEnum {
    Round0 { package_hash: ExodusHash },
    Round1 {
        secret_package: Vec<u8>,
        public_package: Vec<u8>,
    },
    Round2 {
        secret_package: Vec<u8>,
        public_packages: BTreeMap<AuthIndex, Vec<u8>>,
    },
    Round3 { decrypted: BTreeMap<AuthIndex, Vec<u8>> },
    Round4,
    Round5 {
        secret_share: Vec<u8>,
        merkle_proof: Vec<ExodusHash>,
    },
    Round6,
}

#[derive(Default, Encode, Decode, Clone, PartialEq, Eq, RuntimeDebug)]
pub struct ExodusNonce(Vec<u8>);

#[derive(Default, Encode, Decode, Clone, PartialEq, Eq, RuntimeDebug)]
pub struct ExodusGroup {
    pub self_binding_factor: Vec<u8>,
    pub binding_factors_proof: Vec<ExodusHash>,
}

#[derive(Default, Encode, Decode, Clone, PartialEq, Eq, RuntimeDebug)]
pub struct ExodusShare;

#[derive(Default, Encode, Decode, Clone, PartialEq, Eq, RuntimeDebug)]
pub struct ExodusExile;

type ExodusNonces<BlockNumber> = BTreeMap<ExodusSession, PersistentData<BlockNumber, ExodusNonce>>;
type ExodusGroups<BlockNumber> = BTreeMap<ExodusSession, PersistentData<BlockNumber, ExodusGroup>>;
type ExodusShares<BlockNumber> = BTreeMap<ExodusSession, PersistentData<BlockNumber, ExodusShare>>;
type ExodusExiles<BlockNumber> = BTreeMap<ExodusSession, PersistentData<BlockNumber, ExodusExile>>;

pub struct ExodusStorageWrapper<ExodusInnerType> {
    pub inner: ExodusInnerType,
    pub data_changed: bool,
}

impl<ExodusInnerType> ExodusStorageWrapper<ExodusInnerType> {
    pub fn new(inner: ExodusInnerType) -> Self {
        Self { inner, data_changed: false }
    }
}

#[derive(Default)]
pub struct ExodusStorage<BlockNumber, NetworkCurve>
where
	BlockNumber: Default,
	NetworkCurve: Default,
{
    current_block: BlockNumber,
    block_longevity: BlockNumber,
    pub network_curve: NetworkCurve,
    pub authority_index: AuthIndex,
}

impl<B, NC> ExodusStorage<B, NC>
where
	B: AtLeast32BitUnsigned + Copy,
	B: Encode + Decode + Default,
	NC: Encode + Decode + Default + Copy,
{
    pub fn with_network_curve(mut self, network_curve: NC) -> Self {
        self.network_curve = network_curve;
        self
    }

    pub fn with_current_block(mut self, current_block: B) -> Self {
        self.current_block = current_block;
        self
    }

    pub fn with_block_longevity(mut self, block_longevity: B) -> Self {
        self.block_longevity = block_longevity;
        self
    }

    pub fn with_authority_index(mut self, authority_index: AuthIndex) -> Self {
        self.authority_index = authority_index;
        self
    }

    pub fn read_local_storage(&self) -> ExodusRecords<B> {
        let encoded_index = self.authority_index.encode();
        let encoded_network = self.network_curve.encode();

        let nonces_key = create_storage_key(crate::EXODUS_PREFIX_NONCES, &encoded_index, &encoded_network);
        let exodus_nonces = StorageValueRef::persistent(&nonces_key)
            .get::<ExodusNonces<B>>()
            .ok()
            .flatten()
            .unwrap_or_default();

        let shares_key = create_storage_key(crate::EXODUS_PREFIX_SHARES, &encoded_index, &encoded_network);
        let exodus_shares = StorageValueRef::persistent(&shares_key)
            .get::<ExodusShares<B>>()
            .ok()
            .flatten()
            .unwrap_or_default();

        let groups_key = create_storage_key(crate::EXODUS_PREFIX_BLENDS, &encoded_index, &encoded_network);
        let exodus_groups = StorageValueRef::persistent(&groups_key)
            .get::<ExodusGroups<B>>()
            .ok()
            .flatten()
            .unwrap_or_default();

        let exiles_key = create_storage_key(crate::EXODUS_PREFIX_EXILES, &encoded_index, &encoded_network);
        let exodus_exiles = StorageValueRef::persistent(&exiles_key)
            .get::<ExodusExiles<B>>()
            .ok()
            .flatten()
            .unwrap_or_default();

        ExodusRecords {
            exodus_nonces: ExodusStorageWrapper::new(exodus_nonces),
            exodus_shares: ExodusStorageWrapper::new(exodus_shares),
            exodus_groups: ExodusStorageWrapper::new(exodus_groups),
            exodus_exiles: ExodusStorageWrapper::new(exodus_exiles),
        }
    }

    pub fn prepare_signed_call_with_context<T, A, F1, F2, F3>(
        &self,
        authority_key: &A,
        exodus_records: &mut ExodusRecords<B>,
        make_data: F1,
        submit_call: F2,
        insert_data: F3,
    ) -> Result<(), ExodusError>
	where
		T: crate::pallet::Config,
		A: RuntimeAppPublic,
		F1: FnOnce(PackageContext<NC>) -> Result<ExodusPackage<NC, T>, ExodusError>,
		F2: FnOnce(ExodusPackage<NC, T>, A::Signature) -> Result<(), ExodusError>,
		F3: FnOnce(&mut ExodusRecords<B>, B),
	{
        let context = PackageContext::default()
            .with_authority_index(self.authority_index)
            .with_network_curve(self.network_curve);

        let package = make_data(context)?;

        let signature = authority_key
            .sign(&package.encode())
            .ok_or(ExodusError::BadSignature(self.authority_index))?;

        submit_call(package, signature)?;

        let next_release_block = self.current_block.saturating_add(self.block_longevity);
        insert_data(exodus_records, next_release_block);

        self.try_write_local_storage(exodus_records);

        Ok(())
    }

    pub fn force_clear_local_storage(&self) {
        let encoded_index = self.authority_index.encode();
        let encoded_network = self.network_curve.encode();

        let nonces_key = create_storage_key(crate::EXODUS_PREFIX_NONCES, &encoded_index, &encoded_network);
        let shares_key = create_storage_key(crate::EXODUS_PREFIX_SHARES, &encoded_index, &encoded_network);
        let groups_key = create_storage_key(crate::EXODUS_PREFIX_BLENDS, &encoded_index, &encoded_network);
        let exiles_key = create_storage_key(crate::EXODUS_PREFIX_EXILES, &encoded_index, &encoded_network);

        StorageValueRef::persistent(&nonces_key).clear();
        StorageValueRef::persistent(&shares_key).clear();
        StorageValueRef::persistent(&groups_key).clear();
        StorageValueRef::persistent(&exiles_key).clear();
    }

    pub fn try_write_local_storage(&self, records: &ExodusRecords<B>) {
        let encoded_index = self.authority_index.encode();
        let encoded_network = self.network_curve.encode();

        if records.exodus_nonces.data_changed {
            let nonces_key = create_storage_key(crate::EXODUS_PREFIX_NONCES, &encoded_index, &encoded_network);
            StorageValueRef::persistent(&nonces_key).set(&records.exodus_nonces.inner);
        }

        if records.exodus_shares.data_changed {
            let shares_key = create_storage_key(crate::EXODUS_PREFIX_SHARES, &encoded_index, &encoded_network);
            StorageValueRef::persistent(&shares_key).set(&records.exodus_shares.inner);
        }

        if records.exodus_groups.data_changed {
            let groups_key = create_storage_key(crate::EXODUS_PREFIX_BLENDS, &encoded_index, &encoded_network);
            StorageValueRef::persistent(&groups_key).set(&records.exodus_groups.inner);
        }

        if records.exodus_exiles.data_changed {
            let exiles_key = create_storage_key(crate::EXODUS_PREFIX_EXILES, &encoded_index, &encoded_network);
            StorageValueRef::persistent(&exiles_key).set(&records.exodus_exiles.inner);
        }
    }
}

pub struct ExodusRecords<BlockNumber>
where
	BlockNumber: AtLeast32BitUnsigned,
{
    exodus_nonces: ExodusStorageWrapper<ExodusNonces<BlockNumber>>,
    exodus_shares: ExodusStorageWrapper<ExodusShares<BlockNumber>>,
    exodus_groups: ExodusStorageWrapper<ExodusGroups<BlockNumber>>,
    exodus_exiles: ExodusStorageWrapper<ExodusExiles<BlockNumber>>,
}

impl<BlockNumber> ExodusRecords<BlockNumber>
where
	BlockNumber: AtLeast32BitUnsigned + Default,
{
    pub fn get_secret_nonce(&self, exodus_session: &ExodusSession) -> Option<&[u8]> {
        self.exodus_nonces.inner
            .get(exodus_session)
            .map(|v| v.data.0.as_slice())
    }

    pub fn get_group_record(
        &self,
        exodus_session: &ExodusSession,
    ) -> Option<ExodusGroup> {
        self.exodus_groups.inner
            .get(exodus_session)
            .cloned()
            .map(|storage| storage.data)
    }

    pub fn is_exodus_session_banned(
        &self,
        exodus_session: &ExodusSession,
    ) -> bool {
        self.exodus_exiles.inner.contains_key(exodus_session)
    }

    pub fn is_group_registered(
        &self,
        exodus_session: &ExodusSession,
    ) -> bool {
        self.exodus_groups.inner.contains_key(exodus_session)
    }

    pub fn is_waiting_period_for_nonce(
        &self,
        exodus_session: &ExodusSession,
        current_block: &BlockNumber
    ) -> bool {
        self.exodus_nonces.inner
            .get(exodus_session)
            .map(|data| data.release_block.gt(current_block))
            .unwrap_or(false)
    }

    pub fn is_waiting_period_for_share(
        &self,
        exodus_session: &ExodusSession,
        current_block: &BlockNumber,
    ) -> bool {
        self.exodus_shares.inner
            .get(exodus_session)
            .map(|data| data.release_block.gt(&current_block))
            .unwrap_or(false)
    }

    pub fn is_waiting_period_for_group(
        &self,
        exodus_session: &ExodusSession,
        current_block: &BlockNumber,
    ) -> bool {
        self.exodus_groups.inner
            .get(exodus_session)
            .map(|data| data.release_block.gt(&current_block))
            .unwrap_or(false)
    }

    pub fn exclude_exodus_session(
        &mut self,
        exodus_session: ExodusSession,
    ) {
        let exodus_exile = PersistentData {
            data: ExodusExile,
            release_block: Default::default(),
        };

        self.exodus_exiles.inner.insert(exodus_session, exodus_exile);
        self.exodus_exiles.data_changed = true;
    }

    pub fn insert_exodus_nonce(
        &mut self,
        exodus_session: ExodusSession,
        next_release_block: BlockNumber,
        nonce_commitment: Vec<u8>,
    ) {
        let exodus_nonce = PersistentData {
            data: ExodusNonce(nonce_commitment),
            release_block: next_release_block,
        };

        self.exodus_nonces.inner.insert(exodus_session, exodus_nonce);
        self.exodus_nonces.data_changed = true;
    }

    pub fn insert_exodus_share(
        &mut self,
        exodus_session: ExodusSession,
        next_release_block: BlockNumber,
    ) {
        let exodus_share = PersistentData {
            data: ExodusShare,
            release_block: next_release_block,
        };

        self.exodus_shares.inner.insert(exodus_session, exodus_share);
        self.exodus_shares.data_changed = true;
    }

    pub fn insert_exodus_group(
        &mut self,
        exodus_session: ExodusSession,
        next_release_block: BlockNumber,
        self_binding_factor: Vec<u8>,
        binding_factors_proof: Vec<ExodusHash>,
    ) {
        let exodus_group = PersistentData {
            data: ExodusGroup { self_binding_factor, binding_factors_proof },
            release_block: next_release_block,
        };

        self.exodus_groups.inner.insert(exodus_session, exodus_group);
        self.exodus_groups.data_changed = true;
    }
}

#[derive(Default, Clone, Copy)]
pub struct DkgStorage<BlockNumber, NetworkCurve>
where
	BlockNumber: Copy + Clone,
	NetworkCurve: Copy + Clone,
{
    block_longevity: BlockNumber,
    current_block: BlockNumber,
    pub network_curve: NetworkCurve,
    pub authority_index: AuthIndex,
}

impl<B, NC> DkgStorage<B, NC>
where
	B: Saturating + Encode + Decode + Copy + Default + PartialOrd,
	NC: Encode + Copy + Default,
{
    pub fn with_network_curve(mut self, network_curve: NC) -> Self {
        self.network_curve = network_curve;
        self
    }

    pub fn with_current_block(mut self, current_block: B) -> Self {
        self.current_block = current_block;
        self
    }

    pub fn with_block_longevity(mut self, block_longevity: B) -> Self {
        self.block_longevity = block_longevity;
        self
    }

    pub fn with_authority_index(mut self, authority_index: AuthIndex) -> Self {
        self.authority_index = authority_index;
        self
    }

    pub fn get_release_block(&self) -> B {
        self.current_block.saturating_add(self.block_longevity)
    }

    pub fn get_dkg_result(
        &self,
        dkg_record: DkgRecord<B>,
    ) -> Result<ExodusOk<NC>, ExodusError> {
        match dkg_record.prefix_slice {
            crate::DKG_ROUND0_PREFIX => return Ok(ExodusOk::DkgRoundRegistered(0u8, self.authority_index, self.network_curve)),
            crate::DKG_ROUND1_PREFIX => return Ok(ExodusOk::DkgRoundRegistered(1u8, self.authority_index, self.network_curve)),
            crate::DKG_ROUND2_PREFIX => return Ok(ExodusOk::DkgRoundRegistered(2u8, self.authority_index, self.network_curve)),
            crate::DKG_ROUND3_PREFIX => return Ok(ExodusOk::DkgRoundRegistered(3u8, self.authority_index, self.network_curve)),
            crate::DKG_ROUND4_PREFIX => return Ok(ExodusOk::DkgRoundRegistered(4u8, self.authority_index, self.network_curve)),
            _ => {}
        }

        if !dkg_record.prefix_slice.starts_with(crate::DKG_SECRET_PREFIX) {
            return Err(ExodusError::UnknownPrefix);
        }

        match &dkg_record.maybe_stored_data {
            Some(PersistentData { data: DkgStorageEnum::Round5 { .. }, .. }) => {
                Ok(ExodusOk::DkgRoundRegistered(5u8, self.authority_index, self.network_curve))
            },
            Some(PersistentData { data: DkgStorageEnum::Round6 { .. }, .. }) => {
                Ok(ExodusOk::DkgRoundRegistered(6u8, self.authority_index, self.network_curve))
            },
            _ => Err(ExodusError::UnknownPrefix),
        }
    }

    pub fn read_local_storage<'a>(
        &self,
        prefix_slice: &'a [u8],
        authority_index: AuthIndex,
    ) -> DkgRecord<'a, B> {
        let encoded_index = authority_index.encode();
        let encoded_network = self.network_curve.encode();
        let key = create_storage_key(prefix_slice, &encoded_index, &encoded_network);

        let maybe_stored_data = StorageValueRef::persistent(&key)
            .get::<PersistentData<B, DkgStorageEnum>>()
            .ok()
            .flatten();

        DkgRecord { prefix_slice, maybe_stored_data, current_block: self.current_block }
    }

    pub fn prepare_signed_call_with_context<T, A, F1, F2>(
        &self,
        authority_key: &A,
        dkg_record: &mut DkgRecord<B>,
        make_data: F1,
        submit_call: F2,
    ) -> Result<(), ExodusError>
	where
        T: crate::pallet::Config,
        A: RuntimeAppPublic,
		F1: FnOnce(PackageContext<NC>) -> Result<DkgPackage<NC, T>, ExodusError>,
		F2: FnOnce(DkgPackage<NC, T>, A::Signature) -> Result<(), ExodusError>,
    {
        let context = PackageContext::default()
            .with_authority_index(self.authority_index)
            .with_network_curve(self.network_curve);

        let package = make_data(context)?;

        let signature = authority_key
            .sign(&package.encode())
            .ok_or(ExodusError::BadSignature(self.authority_index))?;

        submit_call(package, signature)?;

        self.try_write_local_storage(dkg_record);

        Ok(())
    }

    pub fn try_write_local_storage(
        &self,
        dkg_record: &DkgRecord<B>,
    ) {
        let encoded_index = self.authority_index.encode();
        let encoded_network = self.network_curve.encode();
        let key = create_storage_key(
            dkg_record.prefix_slice,
            &encoded_index,
            &encoded_network,
        );

        match &dkg_record.maybe_stored_data {
            Some(to_store) => StorageValueRef::persistent(&key).set(&to_store),
            None => StorageValueRef::persistent(&key).clear(),
        };
    }

    pub fn clear_local_storage<'a>(
        &self,
        prefix: &'a [u8],
    ) {
        let encoded_index = self.authority_index.encode();
        let encoded_network = self.network_curve.encode();
        let key = create_storage_key(prefix, &encoded_index, &encoded_network);

        StorageValueRef::persistent(&key).clear();
    }
}

pub struct DkgRecord<'a, BlockNumber> {
    maybe_stored_data: Option<PersistentData<BlockNumber, DkgStorageEnum>>,
    current_block: BlockNumber,
    prefix_slice: &'a [u8],
}

impl<'a, BlockNumber> DkgRecord<'a, BlockNumber>
where
	BlockNumber: Default + Copy + PartialOrd,
{
    pub fn is_waiting_period(&self) -> bool {
        self.maybe_stored_data
            .as_ref()
            .map(|data| data.release_block.gt(&self.current_block))
            .unwrap_or(false)
    }

    pub fn storage_exists(&self) -> bool {
        self.maybe_stored_data.is_some()
    }

    pub fn into_inner(&self) -> Option<&DkgStorageEnum> {
        self.maybe_stored_data.as_ref().map(|stored| &stored.data)
    }

    pub fn into_inner_mut(&mut self) -> Option<&mut DkgStorageEnum> {
        self.maybe_stored_data.as_mut().map(|stored| &mut stored.data)
    }

    pub fn try_update_max_signers<F>(
        &mut self,
        get_new_secret: F
    ) -> Result<&Vec<u8>, ExodusError>
	where
		F: FnOnce(&[u8]) -> Result<Vec<u8>, ExodusError>,
 	{
        enum RoundType<'a> {
            R1 { public_package: &'a Vec<u8> },
            R2 { public_packages: &'a BTreeMap<AuthIndex, Vec<u8>> },
        }

        let (old_secret, round_type) = self.maybe_stored_data.as_ref()
            .map(|persistent| match &persistent.data {
                DkgStorageEnum::Round1 { secret_package, public_package } => {
                    Ok((secret_package, RoundType::R1 { public_package }))
                },
                DkgStorageEnum::Round2 { secret_package, public_packages } => {
                    Ok((secret_package, RoundType::R2 { public_packages }))
                },
                _ => Err(ExodusError::NothingToWrite),
            })
            .ok_or(ExodusError::NothingToWrite)??;

        let new_secret = get_new_secret(old_secret)?;
        let new_data = match round_type {
            RoundType::R1 { public_package } => DkgStorageEnum::Round1 {
                secret_package: new_secret,
                public_package: public_package.clone(),
            },
            RoundType::R2 { public_packages } => DkgStorageEnum::Round2 {
                secret_package: new_secret,
                public_packages: public_packages.clone(),
            },
        };

        self.update_with_block(None, new_data)?;

        self.maybe_stored_data.as_ref()
            .map(|persistent| match &persistent.data {
                DkgStorageEnum::Round1 { secret_package, .. } |
            	DkgStorageEnum::Round2 { secret_package, .. } => Ok(secret_package),
                _ => Err(ExodusError::NothingToWrite),
            })
            .ok_or(ExodusError::NothingToWrite)?
    }

    pub fn update_with_block(
        &mut self,
        maybe_release_block: Option<BlockNumber>,
        new_data: DkgStorageEnum,
    ) -> Result<(), ExodusError> {
        use core::mem::discriminant as d;

        let release_block = match &self.maybe_stored_data {
            Some(existing) => {
                if d(&existing.data) != d(&new_data) {
                    return Err(ExodusError::NothingToWrite);
                }
                maybe_release_block.unwrap_or(existing.release_block)
            }
            None => maybe_release_block.unwrap_or_default(),
        };

        self.maybe_stored_data = Some(PersistentData { release_block, data: new_data });

        Ok(())
    }
}
