//! Contains structures and functions related to transaction creation.

use alloc::boxed::Box;
use alloc::collections::{BTreeMap, BTreeSet};
use alloc::string::{String, ToString};
use alloc::vec::Vec;

use miden_protocol::Word;
use miden_protocol::account::AccountId;
use miden_protocol::asset::{Asset, NonFungibleAsset};
use miden_protocol::crypto::merkle::MerkleError;
use miden_protocol::crypto::merkle::store::MerkleStore;
use miden_protocol::errors::{
    AccountError,
    AssetError,
    AssetVaultError,
    NoteError,
    StorageMapError,
    TransactionInputError,
    TransactionScriptError,
};
use miden_protocol::note::{
    Note,
    NoteDetails,
    NoteDetailsCommitment,
    NoteId,
    NoteRecipient,
    NoteScript,
    NoteTag,
    PartialNote,
};
use miden_protocol::transaction::{InputNote, InputNotes, TransactionArgs, TransactionScript};
use miden_protocol::vm::AdviceMap;
use miden_standards::account::interface::{AccountInterface, AccountInterfaceError};
use miden_standards::errors::CodeBuilderError;
use miden_tx::utils::serde::{
    ByteReader,
    ByteWriter,
    Deserializable,
    DeserializationError,
    Serializable,
};
use thiserror::Error;

mod builder;
pub use builder::{
    PaymentNoteDescription,
    PswapTransactionData,
    SwapTransactionData,
    TransactionRequestBuilder,
};

mod foreign;
pub use foreign::ForeignAccount;
pub(crate) use foreign::account_proof_into_inputs;

use crate::store::InputNoteRecord;

// TRANSACTION REQUEST
// ================================================================================================

pub type NoteArgs = Word;

/// Specifies a transaction script to be executed in a transaction.
///
/// A transaction script is a program which is executed after scripts of all input notes have been
/// executed.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum TransactionScriptTemplate {
    /// Specifies the exact transaction script to be executed in a transaction.
    CustomScript(TransactionScript),
    /// Specifies that the transaction script must create the specified output notes.
    ///
    /// It is up to the client to determine how the output notes will be created and this will
    /// depend on the capabilities of the account the transaction request will be applied to.
    /// For example, for Basic Wallets, this may involve invoking `create_note` procedure.
    SendNotes(Vec<PartialNote>),
}

/// Specifies a transaction request that can be executed by an account.
///
/// A request contains information about input notes to be consumed by the transaction (if any),
/// description of the transaction script to be executed (if any), and a set of notes expected
/// to be generated by the transaction or by consuming notes generated by the transaction.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TransactionRequest {
    /// Notes to be consumed by the transaction.
    /// includes both authenticated and unauthenticated notes.
    /// Notes which ID is present in the store are considered authenticated,
    /// the ones which ID is does not exist are considered unauthenticated.
    input_notes: Vec<Note>,
    /// Optional arguments of the input notes to be consumed by the transaction. This
    /// includes both authenticated and unauthenticated notes.
    input_notes_args: Vec<(NoteId, Option<NoteArgs>)>,
    /// Template for the creation of the transaction script.
    script_template: Option<TransactionScriptTemplate>,
    /// A map of recipients of the output notes expected to be generated by the transaction.
    expected_output_recipients: BTreeMap<Word, NoteRecipient>,
    /// A map of details and tags of notes we expect to be created as part of future transactions
    /// with their respective tags.
    ///
    /// For example, after a swap note is consumed, a payback note is expected to be created.
    expected_future_notes: BTreeMap<NoteDetailsCommitment, (NoteDetails, NoteTag)>,
    /// Initial state of the `AdviceMap` that provides data during runtime.
    advice_map: AdviceMap,
    /// Initial state of the `MerkleStore` that provides data during runtime.
    merkle_store: MerkleStore,
    /// Foreign account data requirements keyed by account ID. At execution time, account data
    /// will be retrieved from the network, and injected as advice inputs. Additionally, the
    /// account's code will be added to the executor and prover.
    foreign_accounts: BTreeMap<AccountId, ForeignAccount>,
    /// The number of blocks in relation to the transaction's reference block after which the
    /// transaction will expire. If `None`, the transaction will not expire.
    expiration_delta: Option<u16>,
    /// Indicates whether to **silently** ignore invalid input notes when executing the
    /// transaction. This will allow the transaction to be executed even if some input notes
    /// are invalid.
    ignore_invalid_input_notes: bool,
    /// Optional [`Word`] that will be pushed to the operand stack before the transaction script
    /// execution.
    script_arg: Option<Word>,
    /// Optional [`Word`] that will be pushed to the stack for the authentication procedure
    /// during transaction execution.
    auth_arg: Option<Word>,
    /// Note scripts that the node's NTX builder will need in its script registry.
    ///
    /// See [`TransactionRequestBuilder::expected_ntx_scripts`] for details.
    expected_ntx_scripts: Vec<NoteScript>,
}

impl TransactionRequest {
    // PUBLIC ACCESSORS
    // --------------------------------------------------------------------------------------------

    /// Returns a reference to the transaction request's input note list.
    pub fn input_notes(&self) -> &[Note] {
        &self.input_notes
    }

    /// Returns a list of all input note IDs.
    pub fn input_note_ids(&self) -> impl Iterator<Item = NoteId> {
        self.input_notes.iter().map(Note::id)
    }

    /// Returns the assets held by the transaction's input notes.
    pub fn incoming_assets(&self) -> (BTreeMap<AccountId, u64>, Vec<NonFungibleAsset>) {
        collect_assets(self.input_notes.iter().flat_map(|note| note.assets().iter()))
    }

    /// Returns a map of note IDs to their respective [`NoteArgs`]. The result will include
    /// exclusively note IDs for notes for which [`NoteArgs`] have been defined.
    pub fn get_note_args(&self) -> BTreeMap<NoteId, NoteArgs> {
        self.input_notes_args
            .iter()
            .filter_map(|(note, args)| args.map(|a| (*note, a)))
            .collect()
    }

    /// Returns the expected output own notes of the transaction.
    ///
    /// In this context "own notes" refers to notes that are expected to be created directly by the
    /// transaction script, rather than notes that are created as a result of consuming other
    /// notes.
    pub fn expected_output_own_notes(&self) -> Vec<Note> {
        match &self.script_template {
            Some(TransactionScriptTemplate::SendNotes(notes)) => notes
                .iter()
                .map(|partial| {
                    Note::with_attachments(
                        partial.assets().clone(),
                        *partial.partial_metadata(),
                        self.expected_output_recipients
                            .get(&partial.recipient_digest())
                            .expect("Recipient should be included if it's an own note")
                            .clone(),
                        partial.attachments().clone(),
                    )
                })
                .collect(),
            _ => vec![],
        }
    }

    /// Returns an iterator over the expected output notes.
    pub fn expected_output_recipients(&self) -> impl Iterator<Item = &NoteRecipient> {
        self.expected_output_recipients.values()
    }

    /// Returns an iterator over expected future notes.
    pub fn expected_future_notes(&self) -> impl Iterator<Item = &(NoteDetails, NoteTag)> {
        self.expected_future_notes.values()
    }

    /// Returns the [`TransactionScriptTemplate`].
    pub fn script_template(&self) -> &Option<TransactionScriptTemplate> {
        &self.script_template
    }

    /// Returns the [`AdviceMap`] for the transaction request.
    pub fn advice_map(&self) -> &AdviceMap {
        &self.advice_map
    }

    /// Returns a mutable reference to the [`AdviceMap`] for the transaction request.
    pub fn advice_map_mut(&mut self) -> &mut AdviceMap {
        &mut self.advice_map
    }

    /// Returns the [`MerkleStore`] for the transaction request.
    pub fn merkle_store(&self) -> &MerkleStore {
        &self.merkle_store
    }

    /// Returns the required foreign accounts keyed by account ID.
    pub fn foreign_accounts(&self) -> &BTreeMap<AccountId, ForeignAccount> {
        &self.foreign_accounts
    }

    /// Returns whether to ignore invalid input notes or not.
    pub fn ignore_invalid_input_notes(&self) -> bool {
        self.ignore_invalid_input_notes
    }

    /// Returns the script argument for the transaction request.
    pub fn script_arg(&self) -> &Option<Word> {
        &self.script_arg
    }

    /// Returns the auth argument for the transaction request.
    pub fn auth_arg(&self) -> &Option<Word> {
        &self.auth_arg
    }

    /// Returns the expected NTX scripts that the node's NTX builder will need in its registry.
    pub fn expected_ntx_scripts(&self) -> &[NoteScript] {
        &self.expected_ntx_scripts
    }

    /// Builds the [`InputNotes`] needed for the transaction execution.
    ///
    /// Authenticated input notes are provided by the caller (typically fetched from the store).
    /// Any requested notes not present in that authenticated set are treated as unauthenticated.
    /// The transaction input notes will include both authenticated and unauthenticated notes in
    /// the order they were provided in the transaction request.
    pub(crate) fn build_input_notes(
        &self,
        authenticated_note_records: Vec<InputNoteRecord>,
    ) -> Result<InputNotes<InputNote>, TransactionRequestError> {
        let mut input_notes: BTreeMap<NoteId, InputNote> = BTreeMap::new();

        // Add provided authenticated input notes to the input notes map.
        for authenticated_note_record in authenticated_note_records {
            // Authenticated note records always carry metadata (their inclusion proof
            // injected it), so `id()` is `Some`.
            let authenticated_note_id = authenticated_note_record
                .id()
                .expect("authenticated note record carries metadata so id() is Some");

            if !authenticated_note_record.is_authenticated() {
                return Err(TransactionRequestError::InputNoteNotAuthenticated(
                    authenticated_note_id,
                ));
            }

            if authenticated_note_record.is_consumed() {
                return Err(TransactionRequestError::InputNoteAlreadyConsumed(
                    authenticated_note_id,
                ));
            }

            input_notes.insert(
                authenticated_note_id,
                authenticated_note_record
                    .try_into()
                    .expect("Authenticated note record should be convertible to InputNote"),
            );
        }

        // Add unauthenticated input notes to the input notes map.
        let authenticated_note_ids: BTreeSet<NoteId> = input_notes.keys().copied().collect();
        for note in self.input_notes().iter().filter(|n| !authenticated_note_ids.contains(&n.id()))
        {
            input_notes.insert(note.id(), InputNote::Unauthenticated { note: note.clone() });
        }

        Ok(InputNotes::new(
            self.input_note_ids()
                .map(|note_id| {
                    input_notes
                        .remove(&note_id)
                        .expect("The input note map was checked to contain all input notes")
                })
                .collect(),
        )?)
    }

    /// Converts the [`TransactionRequest`] into [`TransactionArgs`] in order to be executed by a
    /// Miden host.
    pub(crate) fn into_transaction_args(
        self,
        tx_script: Option<TransactionScript>,
    ) -> TransactionArgs {
        let note_args = self.get_note_args();
        let TransactionRequest {
            expected_output_recipients,
            advice_map,
            merkle_store,
            ..
        } = self;

        let mut tx_args = TransactionArgs::new(advice_map).with_note_args(note_args);

        // A script argument without a script has nothing to bind to, so it is only applied when a
        // transaction script is present. With no argument the default empty word is used, which is
        // equivalent to setting no argument at all.
        if let Some(tx_script) = tx_script {
            tx_args =
                tx_args.with_tx_script_and_args(tx_script, self.script_arg.unwrap_or_default());
        }

        if let Some(auth_argument) = self.auth_arg {
            tx_args = tx_args.with_auth_args(auth_argument);
        }

        tx_args
            .extend_output_note_recipients(expected_output_recipients.into_values().map(Box::new));
        tx_args.extend_merkle_store(merkle_store.inner_nodes());

        tx_args
    }

    /// Builds the transaction script based on the account capabilities and the transaction request.
    ///
    /// Returns `None` when the request carries no script template, producing a transaction with no
    /// transaction script (a zero script root).
    ///
    /// Scripts supplied by the caller via [`TransactionScriptTemplate::CustomScript`] are expected
    /// to have already been compiled against the client's source manager (e.g. via
    /// [`Client::code_builder`](crate::Client::code_builder)).
    pub(crate) fn build_transaction_script(
        &self,
        account_interface: &AccountInterface,
    ) -> Result<Option<TransactionScript>, TransactionRequestError> {
        match &self.script_template {
            Some(TransactionScriptTemplate::CustomScript(script)) => Ok(Some(script.clone())),
            Some(TransactionScriptTemplate::SendNotes(notes)) => {
                // TODO: We could pass `SourceManager` to this call, but it needs to be supported
                // in the protocol struct (however, this should also not fail to build often)
                Ok(Some(account_interface.build_send_notes_script(notes, self.expiration_delta)?))
            },
            None => Ok(None),
        }
    }
}

// SERIALIZATION
// ================================================================================================

impl Serializable for TransactionRequest {
    fn write_into<W: ByteWriter>(&self, target: &mut W) {
        self.input_notes.write_into(target);
        self.input_notes_args.write_into(target);
        match &self.script_template {
            None => target.write_u8(0),
            Some(TransactionScriptTemplate::CustomScript(script)) => {
                target.write_u8(1);
                script.write_into(target);
            },
            Some(TransactionScriptTemplate::SendNotes(notes)) => {
                target.write_u8(2);
                notes.write_into(target);
            },
        }
        self.expected_output_recipients.write_into(target);
        self.expected_future_notes.write_into(target);
        self.advice_map.write_into(target);
        self.merkle_store.write_into(target);
        let foreign_accounts: Vec<_> = self.foreign_accounts.values().cloned().collect();
        foreign_accounts.write_into(target);
        self.expiration_delta.write_into(target);
        target.write_u8(u8::from(self.ignore_invalid_input_notes));
        self.script_arg.write_into(target);
        self.auth_arg.write_into(target);
        self.expected_ntx_scripts.write_into(target);
    }
}

impl Deserializable for TransactionRequest {
    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
        let input_notes = Vec::<Note>::read_from(source)?;
        let input_notes_args = Vec::<(NoteId, Option<NoteArgs>)>::read_from(source)?;

        let script_template = match source.read_u8()? {
            0 => None,
            1 => {
                let transaction_script = TransactionScript::read_from(source)?;
                Some(TransactionScriptTemplate::CustomScript(transaction_script))
            },
            2 => {
                let notes = Vec::<PartialNote>::read_from(source)?;
                Some(TransactionScriptTemplate::SendNotes(notes))
            },
            _ => {
                return Err(DeserializationError::InvalidValue(
                    "Invalid script template type".to_string(),
                ));
            },
        };

        let expected_output_recipients = BTreeMap::<Word, NoteRecipient>::read_from(source)?;
        let expected_future_notes =
            BTreeMap::<NoteDetailsCommitment, (NoteDetails, NoteTag)>::read_from(source)?;

        let advice_map = AdviceMap::read_from(source)?;
        let merkle_store = MerkleStore::read_from(source)?;
        let mut foreign_accounts = BTreeMap::new();
        for foreign_account in Vec::<ForeignAccount>::read_from(source)? {
            foreign_accounts.entry(foreign_account.account_id()).or_insert(foreign_account);
        }
        let expiration_delta = Option::<u16>::read_from(source)?;
        let ignore_invalid_input_notes = source.read_u8()? == 1;
        let script_arg = Option::<Word>::read_from(source)?;
        let auth_arg = Option::<Word>::read_from(source)?;
        let expected_ntx_scripts = Vec::<NoteScript>::read_from(source)?;

        Ok(TransactionRequest {
            input_notes,
            input_notes_args,
            script_template,
            expected_output_recipients,
            expected_future_notes,
            advice_map,
            merkle_store,
            foreign_accounts,
            expiration_delta,
            ignore_invalid_input_notes,
            script_arg,
            auth_arg,
            expected_ntx_scripts,
        })
    }
}

// HELPERS
// ================================================================================================

/// Accumulates fungible totals and collectable non-fungible assets from an iterator of assets.
pub(crate) fn collect_assets<'a>(
    assets: impl Iterator<Item = &'a Asset>,
) -> (BTreeMap<AccountId, u64>, Vec<NonFungibleAsset>) {
    let mut fungible_balance_map = BTreeMap::new();
    let mut non_fungible_set = Vec::new();

    assets.for_each(|asset| match asset {
        Asset::Fungible(fungible) => {
            let amount = fungible.amount().as_u64();
            fungible_balance_map
                .entry(fungible.faucet_id())
                .and_modify(|balance| *balance += amount)
                .or_insert(amount);
        },
        Asset::NonFungible(non_fungible) => {
            if !non_fungible_set.contains(non_fungible) {
                non_fungible_set.push(*non_fungible);
            }
        },
    });

    (fungible_balance_map, non_fungible_set)
}

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

// TRANSACTION REQUEST ERROR
// ================================================================================================

// Errors related to a [TransactionRequest]
#[derive(Debug, Error)]
pub enum TransactionRequestError {
    #[error("account interface error")]
    AccountInterfaceError(#[from] AccountInterfaceError),
    #[error("account error")]
    AccountError(#[from] AccountError),
    #[error("asset error")]
    AssetError(#[from] AssetError),
    #[error("duplicate input note: note {0} was added more than once to the transaction")]
    DuplicateInputNote(NoteId),
    #[error(
        "the account proof does not contain the required foreign account data; re-fetch the proof and retry"
    )]
    ForeignAccountDataMissing,
    #[error(
        "foreign account {0} has incompatible visibility; use `ForeignAccount::public()` for public accounts and `ForeignAccount::private()` for private accounts"
    )]
    InvalidForeignAccountId(AccountId),
    #[error(
        "note {0} cannot be used as an authenticated input: it does not have a valid inclusion proof"
    )]
    InputNoteNotAuthenticated(NoteId),
    #[error("note {0} has already been consumed")]
    InputNoteAlreadyConsumed(NoteId),
    #[error(
        "output note declares sender {actual} but the transaction is executed by account {expected}"
    )]
    OutputNoteSenderMismatch { expected: AccountId, actual: AccountId },
    #[error("invalid transaction script")]
    InvalidTransactionScript(#[from] TransactionScriptError),
    #[error("merkle proof error")]
    MerkleError(#[from] MerkleError),
    #[error("empty transaction: the request has no input notes and no account state changes")]
    NoInputNotesNorAccountChange,
    #[error("note not found: {0}")]
    NoteNotFound(String),
    #[error("failed to create note")]
    NoteCreationError(#[from] NoteError),
    #[error("note failed validation")]
    NoteValidationError(#[source] NoteError),
    #[error("note execution failed")]
    NoteExecutionError(#[source] NoteError),
    #[error("failed to build note args")]
    NoteArgError(#[source] NoteError),
    #[error("pay-to-ID note must contain at least one asset to transfer")]
    P2IDNoteWithoutAsset,
    #[error(
        "non-fungible asset issued by faucet {0} is not available in the account vault or incoming notes"
    )]
    MissingNonFungibleAsset(AccountId),
    #[error("PSWAP note can only be cancelled by its creator: expected {expected}, got {actual}")]
    PswapCancelCreatorMismatch { expected: AccountId, actual: AccountId },
    #[error("error building script")]
    CodeBuilderError(#[from] CodeBuilderError),
    #[error("transaction script template error: {0}")]
    ScriptTemplateError(String),
    #[error("storage slot {0} not found in account ID {1}")]
    StorageSlotNotFound(u8, AccountId),
    #[error("error while building the input notes")]
    TransactionInputError(#[from] TransactionInputError),
    #[error("account storage map error")]
    StorageMapError(#[from] StorageMapError),
    #[error("asset vault error")]
    AssetVaultError(#[from] AssetVaultError),
    #[error(
        "unsupported authentication scheme ID {0}; supported schemes are: RpoFalcon512 (0) and EcdsaK256Keccak (1)"
    )]
    UnsupportedAuthSchemeId(u8),
}

// TESTS
// ================================================================================================

#[cfg(test)]
mod tests {
    use std::vec::Vec;

    use miden_protocol::account::auth::{AuthScheme, PublicKeyCommitment};
    use miden_protocol::account::{
        AccountBuilder,
        AccountComponent,
        AccountId,
        AccountType,
        StorageMapKey,
        StorageSlotName,
    };
    use miden_protocol::asset::FungibleAsset;
    use miden_protocol::crypto::rand::{FeltRng, RandomCoin};
    use miden_protocol::note::{NoteAttachments, NoteTag, NoteType};
    use miden_protocol::testing::account_id::{
        ACCOUNT_ID_PRIVATE_FUNGIBLE_FAUCET,
        ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_IMMUTABLE_CODE,
        ACCOUNT_ID_SENDER,
    };
    use miden_protocol::{EMPTY_WORD, Felt, Word};
    use miden_standards::account::auth::AuthSingleSig;
    use miden_standards::note::P2idNote;
    use miden_standards::testing::account_component::MockAccountComponent;
    use miden_tx::utils::serde::{Deserializable, Serializable};

    use super::{TransactionRequest, TransactionRequestBuilder};
    use crate::rpc::domain::account::AccountStorageRequirements;
    use crate::transaction::ForeignAccount;

    #[test]
    fn transaction_request_serialization() {
        assert_transaction_request_serialization_with(|| {
            AuthSingleSig::new(
                PublicKeyCommitment::from(EMPTY_WORD),
                AuthScheme::Falcon512Poseidon2,
            )
            .into()
        });
    }

    #[test]
    fn transaction_request_serialization_ecdsa() {
        assert_transaction_request_serialization_with(|| {
            AuthSingleSig::new(PublicKeyCommitment::from(EMPTY_WORD), AuthScheme::EcdsaK256Keccak)
                .into()
        });
    }

    fn assert_transaction_request_serialization_with<F>(auth_component: F)
    where
        F: FnOnce() -> AccountComponent,
    {
        let sender_id = AccountId::try_from(ACCOUNT_ID_SENDER).unwrap();
        let target_id =
            AccountId::try_from(ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_IMMUTABLE_CODE).unwrap();
        let faucet_id = AccountId::try_from(ACCOUNT_ID_PRIVATE_FUNGIBLE_FAUCET).unwrap();
        let mut rng = RandomCoin::new(Word::default());

        let mut notes = vec![];
        for i in 0..6 {
            let note = P2idNote::create(
                sender_id,
                target_id,
                vec![FungibleAsset::new(faucet_id, 100 + i).unwrap().into()],
                NoteType::Private,
                NoteAttachments::empty(),
                &mut rng,
            )
            .unwrap();
            notes.push(note);
        }

        let mut advice_vec: Vec<(Word, Vec<Felt>)> = vec![];
        for i in 0u32..10 {
            advice_vec.push((rng.draw_word(), vec![Felt::from(i)]));
        }

        let account = AccountBuilder::new(Default::default())
            .with_component(MockAccountComponent::with_empty_slots())
            .with_auth_component(auth_component())
            .account_type(AccountType::Private)
            .build_existing()
            .unwrap();

        // This transaction request wouldn't be valid in a real scenario, it's intended for testing
        let tx_request = TransactionRequestBuilder::new()
            .input_notes(vec![(notes.pop().unwrap(), None)])
            .expected_output_recipients(vec![notes.pop().unwrap().recipient().clone()])
            .expected_future_notes(vec![(
                notes.pop().unwrap().into(),
                NoteTag::with_account_target(sender_id),
            )])
            .extend_advice_map(advice_vec)
            .foreign_accounts([
                ForeignAccount::public(
                    target_id,
                    AccountStorageRequirements::new([(
                        StorageSlotName::new("demo::storage_slot").unwrap(),
                        &[StorageMapKey::new(Word::default())],
                    )]),
                )
                .unwrap(),
                ForeignAccount::private(&account).unwrap(),
            ])
            .own_output_notes(vec![notes.pop().unwrap(), notes.pop().unwrap()])
            .script_arg(rng.draw_word())
            .auth_arg(rng.draw_word())
            .expected_ntx_scripts(vec![notes.first().unwrap().recipient().script().clone()])
            .build()
            .unwrap();

        let mut buffer = Vec::new();
        tx_request.write_into(&mut buffer);

        let deserialized_tx_request = TransactionRequest::read_from_bytes(&buffer).unwrap();
        assert_eq!(tx_request, deserialized_tx_request);
    }
}
