// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

// Copyright 2021 Conflux Foundation. All rights reserved.
// Conflux is free software and distributed under GNU General Public License.
// See http://www.gnu.org/licenses/

use crate::{
    access_path::Path,
    account_address::AccountAddress,
    account_config::{
        type_tag_for_currency_code, AccountResource, AccountRole,
        BalanceResource, ChainIdResource, ChildVASP, Credential,
        CurrencyInfoResource, DesignatedDealer, DesignatedDealerPreburns,
        FreezingBit, ParentVASP, PreburnQueueResource, PreburnResource,
    },
    block_metadata::DiemBlockResource,
    diem_timestamp::DiemTimestampResource,
    on_chain_config::{
        ConfigurationResource, DiemVersion, OnChainConfig,
        RegisteredCurrencies, VMPublishingOption, ValidatorSet,
    },
    validator_config::{
        ValidatorConfigResource, ValidatorOperatorConfigResource,
    },
};
use anyhow::{format_err, Error, Result};
use move_core_types::{identifier::Identifier, move_resource::MoveResource};
use serde::{de::DeserializeOwned, Deserialize, Serialize};
use std::{collections::btree_map::BTreeMap, convert::TryFrom, fmt};

#[derive(Default, Deserialize, PartialEq, Serialize)]
pub struct AccountState(BTreeMap<Vec<u8>, Vec<u8>>);

impl AccountState {
    // By design and do not remove
    pub fn get_account_address(&self) -> Result<Option<AccountAddress>> {
        self.get_account_resource().map(|opt_ar| {
            opt_ar.map(|ar| ar.sent_events().key().get_creator_address())
        })
    }

    pub fn get_account_resource(&self) -> Result<Option<AccountResource>> {
        self.get_resource::<AccountResource>()
    }

    pub fn get_balance_resources(
        &self, currency_codes: &[Identifier],
    ) -> Result<BTreeMap<Identifier, BalanceResource>> {
        currency_codes
            .iter()
            .filter_map(|currency_code| {
                let currency_type_tag =
                    type_tag_for_currency_code(currency_code.to_owned());
                // TODO: update this to use BalanceResource::resource_path once
                // that takes type parameters
                self.get_resource_impl(&BalanceResource::access_path_for(
                    currency_type_tag,
                ))
                .transpose()
                .map(|balance| balance.map(|b| (currency_code.to_owned(), b)))
            })
            .collect()
    }

    pub fn get_preburn_balances(
        &self, currency_codes: &[Identifier],
    ) -> Result<BTreeMap<Identifier, PreburnResource>> {
        currency_codes
            .iter()
            .filter_map(|currency_code| {
                let currency_type_tag =
                    type_tag_for_currency_code(currency_code.to_owned());
                // TODO: update this to use PreburnResource::resource_path once
                // that takes type parameters
                self.get_resource_impl(&PreburnResource::access_path_for(
                    currency_type_tag,
                ))
                .transpose()
                .map(|preburn_balance| {
                    preburn_balance.map(|b| (currency_code.to_owned(), b))
                })
            })
            .collect()
    }

    pub fn get_preburn_queue_balances(
        &self, currency_codes: &[Identifier],
    ) -> Result<BTreeMap<Identifier, PreburnQueueResource>> {
        currency_codes
            .iter()
            .filter_map(|currency_code| {
                let currency_type_tag =
                    type_tag_for_currency_code(currency_code.to_owned());
                // TODO: update this to use PreburnQueueResource::resource_path
                // once that takes type parameters
                self.get_resource_impl(&PreburnQueueResource::access_path_for(
                    currency_type_tag,
                ))
                .transpose()
                .map(|preburn_balance| {
                    preburn_balance.map(|b| (currency_code.to_owned(), b))
                })
            })
            .collect()
    }

    pub fn get_chain_id_resource(&self) -> Result<Option<ChainIdResource>> {
        self.get_resource::<ChainIdResource>()
    }

    pub fn get_configuration_resource(
        &self,
    ) -> Result<Option<ConfigurationResource>> {
        self.get_resource::<ConfigurationResource>()
    }

    pub fn get_diem_timestamp_resource(
        &self,
    ) -> Result<Option<DiemTimestampResource>> {
        self.get_resource::<DiemTimestampResource>()
    }

    pub fn get_validator_config_resource(
        &self,
    ) -> Result<Option<ValidatorConfigResource>> {
        self.get_resource::<ValidatorConfigResource>()
    }

    pub fn get_validator_operator_config_resource(
        &self,
    ) -> Result<Option<ValidatorOperatorConfigResource>> {
        self.get_resource::<ValidatorOperatorConfigResource>()
    }

    pub fn get_freezing_bit(&self) -> Result<Option<FreezingBit>> {
        self.get_resource::<FreezingBit>()
    }

    pub fn get_account_role(
        &self, currency_codes: &[Identifier],
    ) -> Result<Option<AccountRole>> {
        if self.0.contains_key(&ParentVASP::resource_path()) {
            match (
                self.get_resource::<ParentVASP>(),
                self.get_resource::<Credential>(),
            ) {
                (Ok(Some(vasp)), Ok(Some(credential))) => {
                    Ok(Some(AccountRole::ParentVASP { vasp, credential }))
                }
                _ => Ok(None),
            }
        } else if self.0.contains_key(&ChildVASP::resource_path()) {
            self.get_resource::<ChildVASP>()
                .map(|r_opt| r_opt.map(AccountRole::ChildVASP))
        } else if self.0.contains_key(&DesignatedDealer::resource_path()) {
            match (
                self.get_resource::<Credential>(),
                self.get_preburn_balances(&currency_codes),
                self.get_preburn_queue_balances(&currency_codes),
                self.get_resource::<DesignatedDealer>(),
            ) {
                (
                    Ok(Some(dd_credential)),
                    Ok(preburn_balances),
                    Ok(preburn_queues),
                    Ok(Some(designated_dealer)),
                ) => {
                    let preburn_balances = if preburn_balances.is_empty()
                        && !preburn_queues.is_empty()
                    {
                        DesignatedDealerPreburns::PreburnQueue(preburn_queues)
                    } else if !preburn_balances.is_empty()
                        && preburn_queues.is_empty()
                    {
                        DesignatedDealerPreburns::Preburn(preburn_balances)
                    } else {
                        return Ok(None);
                    };
                    Ok(Some(AccountRole::DesignatedDealer {
                        dd_credential,
                        preburn_balances,
                        designated_dealer,
                    }))
                }
                _ => Ok(None),
            }
        } else {
            // TODO: add role_id to Unknown
            Ok(Some(AccountRole::Unknown))
        }
    }

    pub fn get_validator_set(&self) -> Result<Option<ValidatorSet>> {
        self.get_config::<ValidatorSet>()
    }

    pub fn get_diem_version(&self) -> Result<Option<DiemVersion>> {
        self.get_config::<DiemVersion>()
    }

    pub fn get_vm_publishing_option(
        &self,
    ) -> Result<Option<VMPublishingOption>> {
        self.0
            .get(&VMPublishingOption::CONFIG_ID.access_path().path)
            .map(|bytes| VMPublishingOption::deserialize_into_config(bytes))
            .transpose()
            .map_err(Into::into)
    }

    pub fn get_registered_currency_info_resources(
        &self,
    ) -> Result<Vec<CurrencyInfoResource>> {
        let currencies: Option<RegisteredCurrencies> = self.get_config()?;
        match currencies {
            Some(currencies) => {
                let codes = currencies.currency_codes();
                let mut resources = vec![];
                for code in codes {
                    let access_path =
                        CurrencyInfoResource::resource_path_for(code.clone());
                    let info: CurrencyInfoResource = self
                        .get_resource_impl(&access_path.path)?
                        .ok_or_else(|| {
                            format_err!(
                                "currency info resource not found: {}",
                                code
                            )
                        })?;
                    resources.push(info);
                }
                Ok(resources)
            }
            None => Ok(vec![]),
        }
    }

    pub fn get_diem_block_resource(&self) -> Result<Option<DiemBlockResource>> {
        self.get_resource::<DiemBlockResource>()
    }

    pub fn get(&self, key: &[u8]) -> Option<&Vec<u8>> { self.0.get(key) }

    pub fn get_resource_impl<T: DeserializeOwned>(
        &self, key: &[u8],
    ) -> Result<Option<T>> {
        self.0
            .get(key)
            .map(|bytes| bcs::from_bytes(bytes))
            .transpose()
            .map_err(Into::into)
    }

    pub fn insert(&mut self, key: Vec<u8>, value: Vec<u8>) -> Option<Vec<u8>> {
        self.0.insert(key, value)
    }

    pub fn remove(&mut self, key: &[u8]) -> Option<Vec<u8>> {
        self.0.remove(key)
    }

    pub fn iter(
        &self,
    ) -> impl std::iter::Iterator<Item = (&Vec<u8>, &Vec<u8>)> {
        self.0.iter()
    }

    pub fn get_config<T: OnChainConfig>(&self) -> Result<Option<T>> {
        self.get_resource_impl(&T::CONFIG_ID.access_path().path)
    }

    pub fn get_resource<T: MoveResource + DeserializeOwned>(
        &self,
    ) -> Result<Option<T>> {
        self.get_resource_impl(&T::struct_tag().access_vector())
    }

    /// Return an iterator over the module values stored under this account
    pub fn get_modules(&self) -> impl Iterator<Item = &Vec<u8>> {
        self.0.iter().filter_map(|(k, v)| {
            match Path::try_from(k).expect("Invalid access path") {
                Path::Code(_) => Some(v),
                Path::Resource(_) => None,
            }
        })
    }
}

impl fmt::Debug for AccountState {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // TODO: add support for other types of resources
        let account_resource_str = self
            .get_account_resource()
            .map(|account_resource_opt| format!("{:#?}", account_resource_opt))
            .unwrap_or_else(|e| format!("parse error: {:#?}", e));

        let diem_timestamp_str = self
            .get_diem_timestamp_resource()
            .map(|diem_timestamp_opt| format!("{:#?}", diem_timestamp_opt))
            .unwrap_or_else(|e| format!("parse: {:#?}", e));

        let validator_config_str = self
            .get_validator_config_resource()
            .map(|validator_config_opt| format!("{:#?}", validator_config_opt))
            .unwrap_or_else(|e| format!("parse error: {:#?}", e));

        let validator_set_str = self
            .get_validator_set()
            .map(|validator_set_opt| format!("{:#?}", validator_set_opt))
            .unwrap_or_else(|e| format!("parse error: {:#?}", e));

        write!(
            f,
            "{{ \n \
             AccountResource {{ {} }} \n \
             DiemTimestamp {{ {} }} \n \
             ValidatorConfig {{ {} }} \n \
             ValidatorSet {{ {} }} \n \
             }}",
            account_resource_str,
            diem_timestamp_str,
            validator_config_str,
            validator_set_str,
        )
    }
}

impl TryFrom<(&AccountResource, &BalanceResource)> for AccountState {
    type Error = Error;

    fn try_from(
        (account_resource, balance_resource): (
            &AccountResource,
            &BalanceResource,
        ),
    ) -> Result<Self> {
        let mut btree_map: BTreeMap<Vec<u8>, Vec<u8>> = BTreeMap::new();
        btree_map.insert(
            AccountResource::resource_path(),
            bcs::to_bytes(account_resource)?,
        );
        btree_map.insert(
            BalanceResource::resource_path(),
            bcs::to_bytes(balance_resource)?,
        );

        Ok(Self(btree_map))
    }
}