//! A future Solana message format.
//!
//! This crate defines two versions of `Message` in their own modules:
//! [`legacy`] and [`v0`]. `legacy` is the current version as of Solana 1.10.0.
//! `v0` is a [future message format] that encodes more account keys into a
//! transaction than the legacy format.
//!
//! [`legacy`]: crate::legacy
//! [`v0`]: crate::v0
//! [future message format]: https://docs.solanalabs.com/proposals/versioned-transactions

use std::collections::HashSet;

pub use loaded::*;
#[cfg(feature = "frozen-abi")]
use rialo_frozen_abi_macro::AbiExample;
use rialo_s_instruction::Instruction;
use rialo_s_pubkey::Pubkey;
use rialo_s_sdk_ids::bpf_loader_upgradeable;
use rialo_sanitize::SanitizeError;
#[cfg(feature = "serde")]
use serde_derive::{Deserialize, Serialize};

use crate::{
    compiled_instruction::CompiledInstruction,
    compiled_keys::{CompileError, CompiledKeys},
    AccountKeys, ConfigHashPrefix, MessageHeader,
};

mod loaded;

/// The maximum number of static account keys allowed in a v0 message.
pub const MAX_NUM_STATIC_ACCOUNT_KEYS: usize = 32;

/// A Solana transaction message (v0).
///
/// This message format supports succinct account loading with
/// on-chain address lookup tables.
///
/// See the crate documentation for further description.
///
#[cfg_attr(feature = "frozen-abi", derive(AbiExample))]
#[cfg_attr(
    feature = "serde",
    derive(Deserialize, Serialize),
    serde(rename_all = "camelCase")
)]
#[derive(Default, Debug, PartialEq, Eq, Clone)]
pub struct Message {
    /// The message header, identifying signed and read-only `account_keys`.
    /// Header values only describe static `account_keys`, they do not describe
    /// any additional account keys loaded via address table lookups.
    pub header: MessageHeader,

    /// List of accounts loaded by this transaction.
    #[cfg_attr(feature = "serde", serde(with = "rialo_s_short_vec"))]
    pub account_keys: Vec<Pubkey>,

    /// The transaction valid from, as specified by the user.
    pub valid_from: i64,

    /// The config hash prefix for replay protection across chains.
    pub config_hash_prefix: ConfigHashPrefix,

    /// Should the transaction be executed with the OCC scheduler.
    pub occ: bool,

    /// Instructions that invoke a designated program, are executed in sequence,
    /// and committed in one atomic transaction if all succeed.
    ///
    /// # Notes
    ///
    /// Program indexes must index into the list of message `account_keys` because
    /// program id's cannot be dynamically loaded from a lookup table.
    ///
    /// Account indexes must index into the list of addresses
    /// constructed from the concatenation of three key lists:
    ///   1) message `account_keys`
    ///   2) ordered list of keys loaded from `writable` lookup table indexes
    ///   3) ordered list of keys loaded from `readable` lookup table indexes
    #[cfg_attr(feature = "serde", serde(with = "rialo_s_short_vec"))]
    pub instructions: Vec<CompiledInstruction>,
}

impl Message {
    /// Sanitize message fields and compiled instruction indexes
    pub fn sanitize(&self) -> Result<(), SanitizeError> {
        let num_static_account_keys = self.account_keys.len();
        if usize::from(self.header.num_required_signatures)
            .saturating_add(usize::from(self.header.num_readonly_unsigned_accounts))
            > num_static_account_keys
        {
            return Err(SanitizeError::IndexOutOfBounds);
        }

        // there should be at least 1 RW fee-payer account.
        if self.header.num_readonly_signed_accounts >= self.header.num_required_signatures {
            return Err(SanitizeError::InvalidValue);
        }

        // this is redundant with the above sanitization checks which require that:
        // 1) the header describes at least 1 RW account
        // 2) the header doesn't describe more account keys than the number of account keys
        if num_static_account_keys == 0 {
            return Err(SanitizeError::InvalidValue);
        }

        // the number of static account keys must be <= MAX_NUM_STATIC_ACCOUNT_KEYS
        if num_static_account_keys > MAX_NUM_STATIC_ACCOUNT_KEYS {
            return Err(SanitizeError::IndexOutOfBounds);
        }

        // `expect` is safe because of earlier check that
        // `num_static_account_keys` is non-zero
        let max_account_ix = num_static_account_keys
            .checked_sub(1)
            .expect("message doesn't contain any account keys");

        // reject program ids loaded from lookup tables so that
        // static analysis on program instructions can be performed
        // without loading on-chain data from a bank
        let max_program_id_ix =
            // `expect` is safe because of earlier check that
            // `num_static_account_keys` is non-zero
            num_static_account_keys
                .checked_sub(1)
                .expect("message doesn't contain any static account keys");

        for ci in &self.instructions {
            if usize::from(ci.program_id_index) > max_program_id_ix {
                return Err(SanitizeError::IndexOutOfBounds);
            }
            // A program cannot be a payer.
            if ci.program_id_index == 0 {
                return Err(SanitizeError::IndexOutOfBounds);
            }
            for ai in &ci.accounts {
                if usize::from(*ai) > max_account_ix {
                    return Err(SanitizeError::IndexOutOfBounds);
                }
            }
        }

        Ok(())
    }
}

impl Message {
    /// Create a signable transaction message from a `payer` public key,
    /// `recent_blockhash`, list of `instructions`.
    ///
    /// # Examples
    ///
    /// This example uses the [`solana_rpc_client`], [`rialo_s_sdk`], and [`anyhow`] crates.
    ///
    /// [`solana_rpc_client`]: https://docs.rs/solana-rpc-client
    /// [`rialo_s_sdk`]: https://docs.rs/solana-sdk
    /// [`anyhow`]: https://docs.rs/anyhow
    ///
    /// ```
    /// # use rialo_s_program::example_mocks::{
    /// #     solana_rpc_client,
    /// #     rialo_s_sdk,
    /// # };
    /// # use std::borrow::Cow;
    /// # use rialo_s_sdk::account::Account;
    /// use std::time::{SystemTime, UNIX_EPOCH};
    /// use anyhow::Result;
    /// use rialo_s_instruction::{AccountMeta, Instruction};
    /// use rialo_s_message::{VersionedMessage, v0};
    /// use rialo_s_pubkey::Pubkey;
    /// use solana_rpc_client::rpc_client::RpcClient;
    /// use rialo_s_sdk::{
    ///      signature::{Keypair, Signer},
    ///      transaction::VersionedTransaction,
    /// };
    ///
    /// fn create_tx(
    ///     client: &RpcClient,
    ///     instruction: Instruction,
    ///     payer: &Keypair,
    /// ) -> Result<VersionedTransaction> {
    ///     let valid_from = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs();
    ///     // In production, fetch from network: client.get_recent_validator_config_hash()
    ///     // Using EXAMPLE_CONFIG_HASH for documentation purposes only
    ///     let config_hash_prefix = rialo_s_message::EXAMPLE_CONFIG_HASH;
    ///     let tx = VersionedTransaction::try_new(
    ///         VersionedMessage::V0(v0::Message::try_compile(
    ///             &payer.pubkey(),
    ///             &[instruction],
    ///             valid_from as i64,
    ///             config_hash_prefix,
    ///             false, // use OCC?
    ///         )?),
    ///         &[payer],
    ///     )?;
    ///
    ///     Ok(tx)
    /// }
    /// #
    /// # let client = RpcClient::new(String::new());
    /// # let payer = Keypair::new();
    /// # let instruction = Instruction::new_with_bincode(Pubkey::new_unique(), &(), vec![
    /// #   AccountMeta::new(Pubkey::new_unique(), false),
    /// # ]);
    /// # create_tx(&client, instruction, &payer)?;
    /// # Ok::<(), anyhow::Error>(())
    /// ```
    pub fn try_compile(
        payer: &Pubkey,
        instructions: &[Instruction],
        valid_from: i64,
        config_hash_prefix: ConfigHashPrefix,
        occ: bool,
    ) -> Result<Self, CompileError> {
        let compiled_keys = CompiledKeys::compile(instructions, Some(*payer));

        let (header, static_keys) = compiled_keys.try_into_message_components()?;
        let account_keys = AccountKeys::new(&static_keys, None);
        let instructions = account_keys.try_compile_instructions(instructions)?;

        Ok(Self {
            header,
            account_keys: static_keys,
            valid_from,
            config_hash_prefix,
            instructions,
            occ,
        })
    }

    #[cfg(feature = "bincode")]
    /// Serialize this message with a version #0 prefix using bincode encoding.
    pub fn serialize(&self) -> Vec<u8> {
        bincode::serialize(&(crate::MESSAGE_VERSION_PREFIX, self)).unwrap()
    }

    /// Returns the config hash prefix for replay protection.
    pub fn config_hash_prefix(&self) -> ConfigHashPrefix {
        self.config_hash_prefix
    }

    /// Returns true if the account at the specified index is called as a program by an instruction
    pub fn is_key_called_as_program(&self, key_index: usize) -> bool {
        if let Ok(key_index) = u8::try_from(key_index) {
            self.instructions
                .iter()
                .any(|ix| ix.program_id_index == key_index)
        } else {
            false
        }
    }

    /// Returns true if the account at the specified index was requested to be
    /// writable.  This method should not be used directly.
    fn is_writable_index(&self, key_index: usize) -> bool {
        let header = &self.header;
        let num_account_keys = self.account_keys.len();
        let num_signed_accounts = usize::from(header.num_required_signatures);
        if key_index >= num_account_keys {
            // TODO: Remove this `if` condition in the future.
            //       Linear: https://linear.app/subzero-labs/issue/SUB-1292/remove-indexing-of-keys-when-index-is-greater-than-static-account-keys
            // This case existed before when Address Lookup Tables existed.
            // In normal runs, this case should not happen. However, for now,
            // we keep this behavior to avoid any breaking changes.
            false
        } else if key_index >= num_signed_accounts {
            let num_unsigned_accounts = num_account_keys.saturating_sub(num_signed_accounts);
            let num_writable_unsigned_accounts = num_unsigned_accounts
                .saturating_sub(usize::from(header.num_readonly_unsigned_accounts));
            let unsigned_account_index = key_index.saturating_sub(num_signed_accounts);
            unsigned_account_index < num_writable_unsigned_accounts
        } else {
            let num_writable_signed_accounts = num_signed_accounts
                .saturating_sub(usize::from(header.num_readonly_signed_accounts));
            key_index < num_writable_signed_accounts
        }
    }

    /// Returns true if any static account key is the bpf upgradeable loader
    fn is_upgradeable_loader_in_static_keys(&self) -> bool {
        self.account_keys
            .iter()
            .any(|&key| key == bpf_loader_upgradeable::id())
    }

    /// Returns true if the account at the specified index was requested as
    /// writable. Before loading addresses, we can't demote write locks properly
    /// so this should not be used by the runtime. The `reserved_account_keys`
    /// param is optional to allow clients to approximate writability without
    /// requiring fetching the latest set of reserved account keys.
    pub fn is_maybe_writable(
        &self,
        key_index: usize,
        reserved_account_keys: Option<&HashSet<Pubkey>>,
    ) -> bool {
        self.is_writable_index(key_index)
            && !self.is_account_maybe_reserved(key_index, reserved_account_keys)
            && !{
                // demote program ids
                self.is_key_called_as_program(key_index)
                    && !self.is_upgradeable_loader_in_static_keys()
            }
    }

    /// Returns true if the account at the specified index is in the reserved
    /// account keys set. Before loading addresses, we can't detect reserved
    /// account keys properly so this shouldn't be used by the runtime.
    fn is_account_maybe_reserved(
        &self,
        key_index: usize,
        reserved_account_keys: Option<&HashSet<Pubkey>>,
    ) -> bool {
        let mut is_maybe_reserved = false;
        if let Some(reserved_account_keys) = reserved_account_keys {
            if let Some(key) = self.account_keys.get(key_index) {
                is_maybe_reserved = reserved_account_keys.contains(key);
            }
        }
        is_maybe_reserved
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::VersionedMessage;

    #[test]
    fn test_sanitize() {
        assert!(Message {
            header: MessageHeader {
                num_required_signatures: 1,
                ..MessageHeader::default()
            },
            account_keys: vec![Pubkey::new_unique()],
            ..Message::default()
        }
        .sanitize()
        .is_ok());
    }

    #[test]
    fn test_sanitize_with_instruction() {
        assert!(Message {
            header: MessageHeader {
                num_required_signatures: 1,
                ..MessageHeader::default()
            },
            account_keys: vec![Pubkey::new_unique(), Pubkey::new_unique()],
            instructions: vec![CompiledInstruction {
                program_id_index: 1,
                accounts: vec![0],
                data: vec![]
            }],
            ..Message::default()
        }
        .sanitize()
        .is_ok());
    }

    #[test]
    fn test_sanitize_without_signer() {
        assert!(Message {
            header: MessageHeader::default(),
            account_keys: vec![Pubkey::new_unique()],
            ..Message::default()
        }
        .sanitize()
        .is_err());
    }

    #[test]
    fn test_sanitize_without_writable_signer() {
        assert!(Message {
            header: MessageHeader {
                num_required_signatures: 1,
                num_readonly_signed_accounts: 1,
                ..MessageHeader::default()
            },
            account_keys: vec![Pubkey::new_unique()],
            ..Message::default()
        }
        .sanitize()
        .is_err());
    }

    #[test]
    #[ignore] // TODO(SUB-1547): Investigate how to re-enable tests after importing Agave crates
    fn test_sanitize_with_max_account_keys() {
        assert!(Message {
            header: MessageHeader {
                num_required_signatures: 1,
                ..MessageHeader::default()
            },
            account_keys: (0..=u8::MAX).map(|_| Pubkey::new_unique()).collect(),
            ..Message::default()
        }
        .sanitize()
        .is_ok());
    }

    #[test]
    fn test_sanitize_with_too_many_account_keys() {
        assert!(Message {
            header: MessageHeader {
                num_required_signatures: 1,
                ..MessageHeader::default()
            },
            account_keys: (0..=256).map(|_| Pubkey::new_unique()).collect(),
            ..Message::default()
        }
        .sanitize()
        .is_err());
    }

    #[test]
    fn test_sanitize_with_max_table_loaded_keys() {
        assert!(Message {
            header: MessageHeader {
                num_required_signatures: 1,
                ..MessageHeader::default()
            },
            account_keys: vec![Pubkey::new_unique()],
            ..Message::default()
        }
        .sanitize()
        .is_ok());
    }

    #[test]
    #[ignore] // TODO(SUB-1547): Investigate how to re-enable tests after importing Agave crates
    fn test_sanitize_with_too_many_table_loaded_keys() {
        assert!(Message {
            header: MessageHeader {
                num_required_signatures: 1,
                ..MessageHeader::default()
            },
            account_keys: vec![Pubkey::new_unique()],
            ..Message::default()
        }
        .sanitize()
        .is_err());
    }

    #[test]
    fn test_sanitize_with_invalid_ix_program_id() {
        let message = Message {
            header: MessageHeader {
                num_required_signatures: 1,
                ..MessageHeader::default()
            },
            account_keys: vec![Pubkey::new_unique()],
            instructions: vec![CompiledInstruction {
                program_id_index: 2,
                accounts: vec![],
                data: vec![],
            }],
            ..Message::default()
        };

        assert!(message.sanitize().is_err());
    }

    #[test]
    fn test_sanitize_with_invalid_ix_account() {
        assert!(Message {
            header: MessageHeader {
                num_required_signatures: 1,
                ..MessageHeader::default()
            },
            account_keys: vec![Pubkey::new_unique(), Pubkey::new_unique()],
            instructions: vec![CompiledInstruction {
                program_id_index: 1,
                accounts: vec![3],
                data: vec![]
            }],
            ..Message::default()
        }
        .sanitize()
        .is_err());
    }

    #[test]
    fn test_serialize() {
        let message = Message::default();
        let versioned_msg = VersionedMessage::V0(message.clone());
        assert_eq!(message.serialize(), versioned_msg.serialize());
    }

    #[test]
    #[ignore] // TODO(SUB-1547): Investigate how to re-enable tests after importing Agave crates
    fn test_is_maybe_writable() {
        let key0 = Pubkey::new_unique();
        let key1 = Pubkey::new_unique();
        let key2 = Pubkey::new_unique();
        let key3 = Pubkey::new_unique();
        let key4 = Pubkey::new_unique();
        let key5 = Pubkey::new_unique();

        let message = Message {
            header: MessageHeader {
                num_required_signatures: 3,
                num_readonly_signed_accounts: 2,
                num_readonly_unsigned_accounts: 1,
            },
            account_keys: vec![key0, key1, key2, key3, key4, key5],
            ..Message::default()
        };

        let reserved_account_keys = HashSet::from([key3]);

        assert!(message.is_maybe_writable(0, Some(&reserved_account_keys)));
        assert!(!message.is_maybe_writable(1, Some(&reserved_account_keys)));
        assert!(!message.is_maybe_writable(2, Some(&reserved_account_keys)));
        assert!(!message.is_maybe_writable(3, Some(&reserved_account_keys)));
        assert!(message.is_maybe_writable(3, None));
        assert!(message.is_maybe_writable(4, Some(&reserved_account_keys)));
        assert!(!message.is_maybe_writable(5, Some(&reserved_account_keys)));
        assert!(message.is_maybe_writable(6, Some(&reserved_account_keys)));
        assert!(!message.is_maybe_writable(7, Some(&reserved_account_keys)));
        assert!(!message.is_maybe_writable(8, Some(&reserved_account_keys)));
    }

    #[test]
    fn test_is_account_maybe_reserved() {
        let key0 = Pubkey::new_unique();
        let key1 = Pubkey::new_unique();

        let message = Message {
            account_keys: vec![key0, key1],
            ..Message::default()
        };

        let reserved_account_keys = HashSet::from([key1]);

        assert!(!message.is_account_maybe_reserved(0, Some(&reserved_account_keys)));
        assert!(message.is_account_maybe_reserved(1, Some(&reserved_account_keys)));
        assert!(!message.is_account_maybe_reserved(2, Some(&reserved_account_keys)));
        assert!(!message.is_account_maybe_reserved(3, Some(&reserved_account_keys)));
        assert!(!message.is_account_maybe_reserved(4, Some(&reserved_account_keys)));
        assert!(!message.is_account_maybe_reserved(0, None));
        assert!(!message.is_account_maybe_reserved(1, None));
        assert!(!message.is_account_maybe_reserved(2, None));
        assert!(!message.is_account_maybe_reserved(3, None));
        assert!(!message.is_account_maybe_reserved(4, None));
    }
}
