// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Copyright (c) DUSK NETWORK. All rights reserved.

//! Types used by Dusk's stake contract.

use alloc::string::String;
use alloc::vec::Vec;

use bytecheck::CheckBytes;
use dusk_bytes::Serializable;
use piecrust_uplink::CONTRACT_ID_BYTES;
use rkyv::{Archive, Deserialize, Serialize};
#[cfg(feature = "serde")]
use serde_with::hex::Hex;
#[cfg(feature = "serde")]
use serde_with::{As, DisplayFromStr};

use crate::abi::ContractId;
use crate::signatures::bls::{
    PublicKey as BlsPublicKey, SecretKey as BlsSecretKey,
    Signature as BlsSignature,
};
use crate::transfer::withdraw::Withdraw as TransferWithdraw;
use crate::{Dusk, dusk};

/// ID of the genesis stake contract
pub const STAKE_CONTRACT: ContractId = crate::reserved(0x2);

/// Epoch used for stake operations
pub const EPOCH: u64 = 2160;

/// Default number of warnings before being penalized
pub const DEFAULT_STAKE_WARNINGS: u8 = 1;

/// The minimum amount of Dusk one can stake.
#[deprecated(
    since = "0.1.0",
    note = "please use `DEFAULT_MINIMUM_STAKE` instead"
)]
pub const MINIMUM_STAKE: Dusk = DEFAULT_MINIMUM_STAKE;

/// The default minimum amount of Dusk one can stake.
pub const DEFAULT_MINIMUM_STAKE: Dusk = dusk(1_000.0);

/// Configuration for the stake contract
#[derive(Debug, Clone, Archive, Serialize, Deserialize)]
#[archive_attr(derive(CheckBytes))]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct StakeConfig {
    /// Number of warnings before being penalized
    pub warnings: u8,
    /// Minimum amount of Dusk that can be staked
    #[cfg_attr(feature = "serde", serde(with = "As::<DisplayFromStr>"))]
    pub minimum_stake: Dusk,
}

impl StakeConfig {
    /// Create a new default stake configuration.
    #[must_use]
    pub const fn new() -> Self {
        Self {
            warnings: DEFAULT_STAKE_WARNINGS,
            minimum_stake: DEFAULT_MINIMUM_STAKE,
        }
    }
}

impl Default for StakeConfig {
    fn default() -> Self {
        Self::new()
    }
}

/// Calculate the block height at which the next epoch takes effect.
#[must_use]
pub const fn next_epoch(block_height: u64) -> u64 {
    let to_next_epoch = EPOCH - (block_height % EPOCH);
    block_height + to_next_epoch
}

/// Stake a value on the stake contract.
#[derive(Debug, Clone, PartialEq, Eq, Archive, Serialize, Deserialize)]
#[archive_attr(derive(CheckBytes))]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct Stake {
    chain_id: u8,
    keys: StakeKeys,
    #[cfg_attr(feature = "serde", serde(with = "As::<DisplayFromStr>"))]
    value: u64,
    signature: DoubleSignature,
}

impl Stake {
    const MESSAGE_SIZE: usize =
        1 + BlsPublicKey::SIZE + BlsPublicKey::SIZE + u64::SIZE;

    /// Create a new stake specifying the owner.
    #[must_use]
    pub fn new(
        account_sk: &BlsSecretKey,
        owner_sk: &BlsSecretKey,
        value: u64,
        chain_id: u8,
    ) -> Self {
        let account = BlsPublicKey::from(account_sk);
        let owner = BlsPublicKey::from(owner_sk);

        let keys = StakeKeys::new(account, owner);

        let mut stake = Stake {
            chain_id,
            keys,
            value,
            signature: DoubleSignature::default(),
        };

        let msg = stake.signature_message();

        stake.signature = DoubleSignature {
            account: account_sk.sign(&msg),
            owner: owner_sk.sign(&msg),
        };

        stake
    }

    /// Create a new stake from a contract.
    #[must_use]
    pub fn new_from_contract(
        sk: &BlsSecretKey,
        contract: ContractId,
        value: u64,
        chain_id: u8,
    ) -> Self {
        let key = BlsPublicKey::from(sk);

        let keys = StakeKeys::new(key, contract);

        let mut stake = Stake {
            chain_id,
            keys,
            value,
            signature: DoubleSignature::default(),
        };

        let msg = stake.signature_message();

        stake.signature = DoubleSignature {
            account: sk.sign(&msg),
            owner: sk.sign(&msg),
        };

        stake
    }

    /// Account to which the stake will belong.
    #[must_use]
    pub fn keys(&self) -> &StakeKeys {
        &self.keys
    }

    /// Account to which the stake will belong.
    #[must_use]
    pub fn account(&self) -> &BlsPublicKey {
        &self.keys.account
    }

    /// Value to stake.
    #[must_use]
    pub fn value(&self) -> u64 {
        self.value
    }

    /// Returns the chain ID of the stake.
    #[must_use]
    pub fn chain_id(&self) -> u8 {
        self.chain_id
    }

    /// Signature of the stake.
    #[must_use]
    pub fn signature(&self) -> &DoubleSignature {
        &self.signature
    }

    /// Return the message that is used as the input to the signature.
    #[must_use]
    pub fn signature_message(&self) -> [u8; Self::MESSAGE_SIZE] {
        let mut bytes = [0u8; Self::MESSAGE_SIZE];

        bytes[0] = self.chain_id;
        let mut offset = 1;

        bytes[offset..offset + BlsPublicKey::SIZE]
            .copy_from_slice(&self.keys.account.to_bytes());
        offset += BlsPublicKey::SIZE;

        match &self.keys.owner {
            StakeFundOwner::Account(key) => bytes
                [offset..offset + BlsPublicKey::SIZE]
                .copy_from_slice(&key.to_bytes()),
            StakeFundOwner::Contract(contract_id) => bytes
                [offset..offset + CONTRACT_ID_BYTES]
                .copy_from_slice(&contract_id.to_bytes()),
        }

        offset += BlsPublicKey::SIZE;

        bytes[offset..offset + u64::SIZE]
            .copy_from_slice(&self.value.to_bytes());

        bytes
    }
}

/// Withdraw some value from the stake contract to a smart contract
///
/// This struct contains the information necessary to perform the withdrawal,
/// including the account initiating the withdrawal, the amount being withdrawn,
/// and the name of the function to invoke on the target contract.
#[derive(Debug, Clone, PartialEq, Archive, Serialize, Deserialize)]
#[archive_attr(derive(CheckBytes))]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct WithdrawToContract {
    account: BlsPublicKey,
    #[cfg_attr(feature = "serde", serde(with = "As::<DisplayFromStr>"))]
    value: u64,
    fn_name: String,
    #[cfg_attr(feature = "serde", serde(with = "As::<Hex>"))]
    data: Vec<u8>,
}

impl WithdrawToContract {
    /// Creates a new `WithdrawToContract` instance.
    ///
    /// # Parameters
    /// - `account`: The BLS public key of the account initiating the
    ///   withdrawal.
    /// - `value`: The amount of the withdrawal in LUX
    /// - `fn_name`: The name of the function to invoke on the target contract
    ///   as callback for the `contract_to_contract` method.
    ///
    /// # Returns
    /// A new `WithdrawToContract` instance with the specified account, value,
    /// and function name.
    pub fn new(
        account: BlsPublicKey,
        value: u64,
        fn_name: impl Into<String>,
    ) -> Self {
        Self {
            account,
            value,
            fn_name: fn_name.into(),
            data: Vec::new(),
        }
    }

    /// Returns the account initiating the withdrawal.
    #[must_use]
    pub fn account(&self) -> &BlsPublicKey {
        &self.account
    }

    /// Returns the value (amount) of the withdrawal.
    #[must_use]
    pub fn value(&self) -> u64 {
        self.value
    }

    /// Returns the name of the callback function to invoke on the target
    /// contract.
    #[must_use]
    pub fn fn_name(&self) -> &str {
        &self.fn_name
    }

    /// Returns the data to be passed to the target contract.
    #[must_use]
    pub fn data(&self) -> &[u8] {
        &self.data
    }

    /// Set the data to be passed to the target contract.
    #[must_use]
    pub fn with_data(mut self, data: Vec<u8>) -> Self {
        self.data = data;
        self
    }
}

/// Withdraw some value from the stake contract.
///
/// This is used in both `unstake` and `withdraw`.
#[derive(Debug, Clone, PartialEq, Archive, Serialize, Deserialize)]
#[archive_attr(derive(CheckBytes))]
// We need to keep field name `withdraw` for serde serialization compatibility
#[allow(clippy::struct_field_names)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct Withdraw {
    account: BlsPublicKey,
    withdraw: TransferWithdraw,
    signature: DoubleSignature,
}

impl Withdraw {
    /// Create a new withdraw call specifying the owner.
    #[must_use]
    pub fn new(
        account_sk: &BlsSecretKey,
        owner_sk: &BlsSecretKey,
        withdraw: TransferWithdraw,
    ) -> Self {
        let account = BlsPublicKey::from(account_sk);

        let mut stake_withdraw = Withdraw {
            account,
            withdraw,
            signature: DoubleSignature::default(),
        };

        let msg = stake_withdraw.signature_message();

        stake_withdraw.signature = DoubleSignature {
            account: account_sk.sign(&msg),
            owner: owner_sk.sign(&msg),
        };

        stake_withdraw
    }

    /// Create a new withdraw call using the same account as the owner.
    #[must_use]
    pub fn with_single_key(
        sk: &BlsSecretKey,
        withdraw: TransferWithdraw,
    ) -> Self {
        Self::new(sk, sk, withdraw)
    }

    /// The public key to withdraw from.
    #[must_use]
    pub fn account(&self) -> &BlsPublicKey {
        &self.account
    }

    /// The inner withdrawal to pass to the transfer contract.
    #[must_use]
    pub fn transfer_withdraw(&self) -> &TransferWithdraw {
        &self.withdraw
    }

    /// Signature of the withdraw.
    #[must_use]
    pub fn signature(&self) -> &DoubleSignature {
        &self.signature
    }

    /// Signature message used for [`Withdraw`].
    #[must_use]
    pub fn signature_message(&self) -> Vec<u8> {
        let mut bytes = Vec::new();

        bytes.extend(self.account.to_bytes());
        bytes.extend(self.withdraw.wrapped_signature_message());

        bytes
    }
}

/// Event emitted after a stake contract operation is performed.
#[derive(Debug, Clone, Archive, Deserialize, Serialize, PartialEq)]
#[archive_attr(derive(CheckBytes))]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct StakeEvent {
    /// Keys associated to the event.
    pub keys: StakeKeys,
    /// Effective value of the relevant operation, be it `stake`,
    /// `unstake`,`withdraw`
    #[cfg_attr(feature = "serde", serde(with = "As::<DisplayFromStr>"))]
    pub value: u64,
    /// The locked amount involved in the operation (e.g., for `stake` or
    /// `unstake`). Defaults to zero for operations that do not involve
    /// locking.
    #[cfg_attr(feature = "serde", serde(with = "As::<DisplayFromStr>"))]
    pub locked: u64,
}

impl StakeEvent {
    /// Creates a new `StakeEvent` with the specified keys and value.
    ///
    /// ### Parameters
    /// - `keys`: The keys associated with the stake event.
    /// - `value`: The effective value of the operation (e.g., `stake`,
    ///   `unstake`, `withdraw`).
    ///
    /// The `locked` amount is initialized to zero by default.
    #[must_use]
    pub fn new(keys: StakeKeys, value: u64) -> Self {
        Self {
            keys,
            value,
            locked: 0,
        }
    }
    /// Sets the locked amount for the `StakeEvent`.
    ///
    /// ### Parameters
    /// - `locked`: The locked amount associated with the operation.
    #[must_use]
    pub fn locked(mut self, locked: u64) -> Self {
        self.locked = locked;
        self
    }
}

/// Event emitted after a slash operation is performed.
#[derive(Debug, Clone, Archive, Deserialize, Serialize, PartialEq)]
#[archive_attr(derive(CheckBytes))]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct SlashEvent {
    /// Account slashed.
    pub account: BlsPublicKey,
    /// Slashed amount
    #[cfg_attr(feature = "serde", serde(with = "As::<DisplayFromStr>"))]
    pub value: u64,
    /// New eligibility for the slashed account
    #[cfg_attr(feature = "serde", serde(with = "As::<DisplayFromStr>"))]
    pub next_eligibility: u64,
}

/// The representation of a public key's stake.
///
/// A user can stake for a particular `amount` larger in value than the
/// `MINIMUM_STAKE` value and is `reward`ed for participating in the consensus.
/// A stake is valid only after a particular block height - called the
/// eligibility.
///
/// To keep track of the number of interactions a public key has had with the
/// contract a `counter` is used to prevent repeat attacks - where the same
/// signature could be used to prove ownership of the secret key in two
/// different transactions.
#[derive(
    Debug, Default, Clone, Copy, PartialEq, Eq, Archive, Deserialize, Serialize,
)]
#[archive_attr(derive(CheckBytes))]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct StakeData {
    /// Amount staked and eligibility.
    pub amount: Option<StakeAmount>,
    /// The reward for participating in consensus.
    #[cfg_attr(feature = "serde", serde(with = "As::<DisplayFromStr>"))]
    pub reward: u64,
    /// Faults
    pub faults: u8,
    /// Hard Faults
    pub hard_faults: u8,
}

/// Keys that identify a stake
#[derive(
    Debug, Default, Clone, Copy, PartialEq, Eq, Archive, Deserialize, Serialize,
)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[archive_attr(derive(CheckBytes))]
pub struct StakeKeys {
    /// Key used for consensus operations, such as voting or producing blocks.
    pub account: BlsPublicKey,

    /// Key used to manage funds for operations like staking, unstaking, or
    /// withdrawing.
    ///
    /// This field can represent ownership either through an individual account
    /// (`StakeFundOwner::Account`) or through a smart contract
    /// (`StakeFundOwner::Contract`).
    pub owner: StakeFundOwner,
}

impl StakeKeys {
    /// Creates a new `StakeKeys` instance where both the consensus key and the
    /// owner key are set to the same account key.
    #[must_use]
    pub fn single_key(account: BlsPublicKey) -> Self {
        Self::new(account, account)
    }

    /// Creates a new `StakeKeys` instance with the specified account and funds
    /// owner.
    ///
    /// # Parameters
    /// - `account`: The BLS public key used for consensus operations, such as
    ///   voting or producing blocks.
    /// - `owner`: The owner of the funds, which can be either an account or a
    ///   contract. This parameter is any type that implements
    ///   `Into<StakeFundOwner>`, allowing flexibility in how the funds owner is
    ///   specified.
    #[must_use]
    pub fn new<F: Into<StakeFundOwner>>(
        account: BlsPublicKey,
        owner: F,
    ) -> Self {
        let owner = owner.into();
        Self { account, owner }
    }
}

/// Data used to check ownership of funds
#[derive(
    Debug, Clone, Copy, PartialEq, Eq, Archive, Deserialize, Serialize,
)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[archive_attr(derive(CheckBytes))]
pub enum StakeFundOwner {
    /// Represents an account-based owner identified by a BLS public key.
    ///
    /// This variant is used when the stake funds are associated directly with
    /// an individual account.
    Account(BlsPublicKey),

    /// Represents a contract-based owner identified by a contract ID.
    ///
    /// This variant is used when the stake funds are managed or associated
    /// with a specific smart contract.
    Contract(ContractId),
}

impl Default for StakeFundOwner {
    fn default() -> Self {
        BlsPublicKey::default().into()
    }
}

impl From<BlsPublicKey> for StakeFundOwner {
    fn from(value: BlsPublicKey) -> Self {
        Self::Account(value)
    }
}

impl From<ContractId> for StakeFundOwner {
    fn from(value: ContractId) -> Self {
        Self::Contract(value)
    }
}

/// Signature used to handle funds (stake/unstake/withdraw)
#[derive(
    Debug, Default, Clone, Copy, PartialEq, Eq, Archive, Deserialize, Serialize,
)]
#[archive_attr(derive(CheckBytes))]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct DoubleSignature {
    /// Signature created with the account key
    pub account: BlsSignature,
    /// Signature created with the owner key
    pub owner: BlsSignature,
}

impl StakeData {
    /// An empty stake.
    pub const EMPTY: Self = Self {
        amount: None,
        reward: 0,
        faults: 0,
        hard_faults: 0,
    };

    /// Create a new stake given its initial `value` and `reward`, together with
    /// the `block_height` of its creation.
    #[must_use]
    pub const fn new(value: u64, reward: u64, block_height: u64) -> Self {
        let eligibility = Self::eligibility_from_height(block_height);
        Self::with_eligibility(value, reward, eligibility)
    }

    /// Create a new stake given its initial `value` and `reward`, together with
    /// the `eligibility`.
    #[must_use]
    pub const fn with_eligibility(
        value: u64,
        reward: u64,
        eligibility: u64,
    ) -> Self {
        let amount = match value {
            0 => None,
            _ => Some(StakeAmount {
                value,
                eligibility,
                locked: 0,
            }),
        };

        Self {
            amount,
            reward,
            faults: 0,
            hard_faults: 0,
        }
    }

    /// Returns true if the stake is valid - meaning there is an `amount` staked
    /// and the given `block_height` is larger or equal to the stake's
    /// eligibility. If there is no `amount` staked this is false.
    #[must_use]
    pub fn is_valid(&self, block_height: u64) -> bool {
        match &self.amount {
            Some(amount) => block_height >= amount.eligibility,
            None => false,
        }
    }

    /// Compute the eligibility of a stake from the starting block height.
    #[must_use]
    pub const fn eligibility_from_height(block_height: u64) -> u64 {
        StakeAmount::eligibility_from_height(block_height)
    }

    /// Check if there is no amount left to withdraw
    ///
    /// Return true if both stake and rewards are 0
    pub fn is_empty(&self) -> bool {
        let stake = self
            .amount
            .as_ref()
            .map(StakeAmount::total_funds)
            .unwrap_or_default();
        self.reward + stake == 0
    }
}

/// Value staked and eligibility.
#[derive(
    Debug, Default, Clone, Copy, PartialEq, Eq, Archive, Deserialize, Serialize,
)]
#[archive_attr(derive(CheckBytes))]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct StakeAmount {
    /// The value staked.
    #[cfg_attr(feature = "serde", serde(with = "As::<DisplayFromStr>"))]
    pub value: u64,
    /// The amount that has been locked (due to a soft slash).
    #[cfg_attr(feature = "serde", serde(with = "As::<DisplayFromStr>"))]
    pub locked: u64,
    /// The eligibility of the stake.
    #[cfg_attr(feature = "serde", serde(with = "As::<DisplayFromStr>"))]
    pub eligibility: u64,
}

impl StakeAmount {
    /// Create a new stake amount.
    #[must_use]
    pub const fn new(value: u64, block_height: u64) -> Self {
        let eligibility = Self::eligibility_from_height(block_height);
        Self::with_eligibility(value, eligibility)
    }

    /// Create a new stake given its initial `value` and `reward`, together with
    /// the `eligibility`.
    #[must_use]
    pub const fn with_eligibility(value: u64, eligibility: u64) -> Self {
        Self {
            value,
            eligibility,
            locked: 0,
        }
    }

    /// Compute the eligibility of a stake from the starting block height.
    #[must_use]
    pub const fn eligibility_from_height(block_height: u64) -> u64 {
        let maturity_blocks = EPOCH;
        next_epoch(block_height) + maturity_blocks
    }

    /// Move `amount` to locked value
    pub fn lock_amount(&mut self, amount: u64) {
        self.value -= amount;
        self.locked += amount;
    }

    /// Get the total funds belonging to the stake (value + locked)
    #[must_use]
    pub fn total_funds(&self) -> u64 {
        self.value + self.locked
    }
}

/// Used in a `reward` call to reward a given account with an amount of Dusk,
/// and emitted as an event, once a reward succeeds.
#[derive(Debug, Clone, PartialEq, Eq, Archive, Serialize, Deserialize)]
#[archive_attr(derive(CheckBytes))]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct Reward {
    /// The account to be rewarded.
    pub account: BlsPublicKey,
    /// The amount to reward.
    #[cfg_attr(feature = "serde", serde(with = "As::<DisplayFromStr>"))]
    pub value: u64,
    /// The reason for the reward.
    pub reason: RewardReason,
}

/// The reason that a reward is issued.
#[derive(Debug, Clone, PartialEq, Eq, Archive, Serialize, Deserialize)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[archive_attr(derive(CheckBytes))]
pub enum RewardReason {
    /// The fixed amount awarded to a generator.
    GeneratorFixed,
    /// Extra amount awarded to a generator.
    GeneratorExtra,
    /// Amount awarded to a voter.
    Voter,
    /// Amount awarded for another reason, such as rewarding Dusk.
    Other,
}
