#![cfg_attr(not(feature = "std"), no_std, no_main)]

extern crate alloc;

use alloc::vec::Vec;

use polymesh_ink::*;

#[ink::contract(env = PolymeshEnvironment)]
mod settlements {
    use alloc::vec;
    use ink::storage::Mapping;

    use crate::*;

    pub const UNIT: Balance = 1_000_000u128;

    /// A contract that uses the settlements pallet.
    #[ink(storage)]
    pub struct Settlements {
        /// The `AccountId` of a privileged account that override the
        /// code hash for `PolymeshInk`.
        ///
        /// This address is set to the account that instantiated this contract.
        admin: AccountId,
        /// Upgradable Polymesh Ink API.
        api: PolymeshInk,

        /// Ticker pair.
        ticker1: Ticker,
        ticker2: Ticker,
        initialized: bool,
        /// Venue for settlements.
        venue: VenueId,
        /// Contract's identity.
        did: IdentityId,
        /// Custodial portfolios.
        portfolios: Mapping<IdentityId, PortfolioId>,
    }

    /// The contract error types.
    #[derive(Debug, scale::Encode, scale::Decode)]
    #[cfg_attr(feature = "std", derive(scale_info::TypeInfo))]
    pub enum Error {
        /// PolymeshInk errors.
        PolymeshInk(PolymeshError),
        /// Caller needs to pay the contract for the protocol fee.
        /// (Amount needed)
        InsufficientTransferValue(Balance),
        /// Scale decode failed.
        ScaleError,
        /// Missing Identity.  MultiSig's are not supported.
        MissingIdentity,
        /// Contract hasn't been initialized.
        NotInitialized,
        /// Contract has already been initialized.
        AlreadyInitialized,
        /// Invalid portfolio authorization.
        InvalidPortfolioAuthorization,
        /// The caller has already initialized a portfolio.
        AlreadyHavePortfolio,
        /// The caller doesn't have a portfolio yet.
        NoPortfolio,
        /// Invalid ticker.
        InvalidTicker,
    }

    impl From<PolymeshError> for Error {
        fn from(err: PolymeshError) -> Self {
            Self::PolymeshInk(err)
        }
    }

    /// The contract result type.
    pub type Result<T> = core::result::Result<T, Error>;

    impl Settlements {
        /// Creates a new contract.
        #[ink(constructor)]
        pub fn new(ticker1: Ticker, ticker2: Ticker) -> Result<Self> {
            Ok(Self {
                ticker1,
                ticker2,
                admin: Self::env().caller(),
                portfolios: Default::default(),
                did: PolymeshInk::get_our_did()?,
                venue: Default::default(),
                initialized: false,
                api: PolymeshInk::new()?,
            })
        }

        /// Update the code hash of the polymesh runtime API.
        ///
        /// Only the `admin` is allowed to call this.
        #[ink(message)]
        pub fn update_code_hash(&mut self, hash: Hash) {
            assert_eq!(
                self.env().caller(),
                self.admin,
                "caller {:?} does not have sufficient permissions, only {:?} does",
                self.env().caller(),
                self.admin,
            );
            self.api.update_code_hash(hash);
        }

        /// Update the `polymesh-ink` API using the tracker.
        ///
        /// Anyone can pay the gas fees to do the update using the tracker.
        #[ink(message)]
        pub fn update_polymesh_ink(&mut self) -> Result<()> {
            self.api.check_for_upgrade()?;
            Ok(())
        }

        fn create_asset(&mut self, ticker: Ticker) -> Result<()> {
            self.api.asset_create_and_issue(
                AssetName(b"".to_vec()),
                ticker,
                AssetType::EquityCommon,
                true, // Divisible token.
                Some(1_000_000 * UNIT),
            )?;
            Ok(())
        }

        fn ensure_ticker(&self, ticker: Ticker) -> Result<()> {
            if self.ticker1 != ticker && self.ticker2 != ticker {
                Err(Error::InvalidTicker)
            } else {
                Ok(())
            }
        }

        fn ensure_has_portfolio(&self, did: IdentityId) -> Result<PortfolioId> {
            self.portfolios.get(did).ok_or(Error::NoPortfolio)
        }

        fn ensure_no_portfolio(&self, did: IdentityId) -> Result<()> {
            if self.portfolios.get(did).is_some() {
                return Err(Error::AlreadyHavePortfolio);
            }
            Ok(())
        }

        fn ensure_initialized(&self) -> Result<()> {
            if !self.initialized {
                return Err(Error::NotInitialized);
            }
            Ok(())
        }

        fn init_venue(&mut self, name: Vec<u8>) -> Result<()> {
            if self.initialized {
                return Err(Error::AlreadyInitialized);
            }
            // Update our identity id.
            self.did = PolymeshInk::get_our_did()?;
            // Create venue.
            self.venue = self
                .api
                .create_venue(VenueDetails(b"Contract Venue".to_vec()), VenueType::Other)?;

            // Create a portfolio for the contract.
            let portfolio = self.api.create_portfolio(name)?;
            self.portfolios.insert(self.did, &portfolio);

            // Create tickers.
            self.create_asset(self.ticker1)?;
            self.create_asset(self.ticker2)?;

            self.initialized = true;
            Ok(())
        }

        fn transfer_assets(&self, legs: Vec<Leg>, portfolios: Vec<PortfolioId>) -> Result<()> {
            self.api.settlement_execute(self.venue, legs, portfolios)?;
            Ok(())
        }

        #[ink(message, payable)]
        pub fn init(&mut self, name: Vec<u8>) -> Result<()> {
            self.init_venue(name)
        }

        #[ink(message)]
        pub fn venue(&self) -> Result<VenueId> {
            self.ensure_initialized()?;
            Ok(self.venue)
        }

        #[ink(message)]
        pub fn contract_did(&self) -> Result<IdentityId> {
            self.ensure_initialized()?;
            Ok(self.did)
        }

        fn fund_caller(&self) -> Result<()> {
            // Get the caller's identity.
            let caller_did = PolymeshInk::get_caller_did()?;

            // Ensure the caller has a portfolio.
            let caller_portfolio = self.ensure_has_portfolio(caller_did)?;

            // Transfer some tokens to the caller's portfolio.
            let our_portfolio = PortfolioId {
                did: self.did,
                kind: PortfolioKind::Default,
            };
            self.transfer_assets(
                vec![
                    Leg::Fungible {
                        sender: our_portfolio,
                        receiver: caller_portfolio,
                        ticker: self.ticker1,
                        amount: 10 * UNIT,
                    },
                    Leg::Fungible {
                        sender: our_portfolio,
                        receiver: caller_portfolio,
                        ticker: self.ticker2,
                        amount: 20 * UNIT,
                    },
                ],
                vec![our_portfolio, caller_portfolio],
            )?;

            Ok(())
        }

        #[ink(message)]
        /// Accept custody of a portfolio and give the caller some tokens.
        pub fn add_portfolio(&mut self, auth_id: u64, portfolio: PortfolioKind) -> Result<()> {
            self.ensure_initialized()?;
            let portfolio = self.api.accept_portfolio_custody(auth_id, portfolio)?;
            let caller_did = portfolio.did;
            // Ensure the caller doesn't have a portfolio.
            self.ensure_no_portfolio(caller_did)?;
            // Save the caller's portfolio.
            self.portfolios.insert(caller_did, &portfolio);

            // Give the caller some funds.
            self.fund_caller()?;
            Ok(())
        }

        #[ink(message)]
        /// Allow the caller to withdrawal funds from the contract controlled portfolio.
        pub fn withdrawal(
            &mut self,
            ticker: Ticker,
            amount: Balance,
            dest: PortfolioKind,
        ) -> Result<()> {
            self.ensure_initialized()?;
            self.ensure_ticker(ticker)?;

            // Get the caller's identity.
            let caller_did = PolymeshInk::get_caller_did()?;
            let dest = PortfolioId {
                did: caller_did,
                kind: dest,
            };

            // Ensure the caller has a portfolio.
            let caller_portfolio = self.ensure_has_portfolio(caller_did)?;

            self.api.move_portfolio_funds(
                caller_portfolio, // Contract controlled portfolio.
                dest,             // Caller controlled portfolio.
                vec![Fund {
                    description: FundDescription::Fungible {
                        ticker: ticker,
                        amount,
                    },
                    memo: None,
                }],
            )?;
            Ok(())
        }

        #[ink(message)]
        /// Return the caller's portfolio custodianship back to them.
        pub fn withdrawal_all(&mut self) -> Result<()> {
            self.ensure_initialized()?;

            // Get the caller's identity.
            let caller_did = PolymeshInk::get_caller_did()?;

            // Ensure the caller has a portfolio.
            let portfolio = self.ensure_has_portfolio(caller_did)?;

            // Remove our custodianship.
            self.api.quit_portfolio_custody(portfolio)?;
            // Remove the portfolio.
            self.portfolios.remove(caller_did);

            Ok(())
        }

        #[ink(message)]
        /// Trade.
        pub fn trade(
            &mut self,
            sell: Ticker,
            sell_amount: Balance,
            buy: Ticker,
            buy_amount: Balance,
        ) -> Result<()> {
            self.ensure_initialized()?;
            self.ensure_ticker(sell)?;
            self.ensure_ticker(buy)?;

            // Get the caller's identity.
            let caller_did = PolymeshInk::get_caller_did()?;

            // Ensure the caller has a portfolio.
            let caller_portfolio = self.ensure_has_portfolio(caller_did)?;

            // Use settlement to complete the trade.
            let our_portfolio = PortfolioId {
                did: self.did,
                kind: PortfolioKind::Default,
            };
            self.transfer_assets(
                vec![
                    Leg::Fungible {
                        sender: caller_portfolio,
                        receiver: our_portfolio,
                        ticker: sell,
                        amount: sell_amount,
                    },
                    Leg::Fungible {
                        sender: our_portfolio,
                        receiver: caller_portfolio,
                        ticker: buy,
                        amount: buy_amount,
                    },
                ],
                vec![our_portfolio, caller_portfolio],
            )?;

            Ok(())
        }

        /// Get an identity's asset balance.
        #[ink(message)]
        pub fn asset_balance_of(
            &mut self,
            ticker: Ticker,
            did: IdentityId,
        ) -> PolymeshResult<Balance> {
            Ok(self.api.asset_balance_of(ticker, did)?)
        }

        /// Get the `total_supply` of an asset.
        #[ink(message)]
        pub fn asset_total_supply(&mut self, ticker: Ticker) -> PolymeshResult<Balance> {
            Ok(self.api.asset_total_supply(ticker)?)
        }

        /// Get Corporate action distribution summary.
        #[ink(message)]
        pub fn distribution_summary(
            &mut self,
            ca_id: CAId,
        ) -> PolymeshResult<Option<DistributionSummary>> {
            Ok(self.api.distribution_summary(ca_id)?)
        }

        /// Claim dividends.
        #[ink(message)]
        pub fn dividend_claim(&mut self, ca_id: CAId) -> PolymeshResult<()> {
            Ok(self.api.dividend_claim(ca_id)?)
        }

        /// Create a simple dividend distribution.
        #[ink(message)]
        pub fn create_dividend(
            &mut self,
            ticker: Ticker,
            portfolio: Option<PortfolioNumber>,
            currency: Ticker,
            per_share: Balance,
            amount: Balance,
        ) -> PolymeshResult<()> {
            let now = Self::env().block_timestamp();
            let dividend = SimpleDividend {
                ticker,
                decl_date: now,
                record_date: now,
                portfolio,
                currency,
                per_share,
                amount,
                payment_at: now + 1_000,
                expires_at: None,
            };
            Ok(self.api.create_dividend(dividend)?)
        }
    }
}
