light-token 0.4.0

use light_compressed_account::compressed_account::PackedMerkleContext;
use light_program_profiler::profile;
use light_sdk::{
    error::LightSdkError,
    instruction::{AccountMetasVec, PackedAccounts, PackedStateTreeInfo, SystemAccountMetaConfig},
};
use light_token_interface::instructions::{
    extensions::ExtensionInstructionData,
    transfer2::{CompressedCpiContext, MultiInputTokenDataWithContext},
};
use solana_account_info::AccountInfo;
use solana_instruction::{AccountMeta, Instruction};
use solana_pubkey::Pubkey;

use super::{
    account2::CTokenAccount2,
    transfer2::{
        account_metas::Transfer2AccountsMetaConfig, create_transfer2_instruction, Transfer2Config,
        Transfer2Inputs,
    },
};
use crate::{
    compat::TokenData, error::TokenSdkError, utils::TokenDefaultAccounts, AnchorDeserialize,
    AnchorSerialize, ValidityProof,
};

/// Struct to hold all the data needed for DecompressFull operation
/// Contains the complete compressed account data and destination index
#[derive(Debug, Clone, AnchorDeserialize, AnchorSerialize)]
pub struct DecompressFullIndices {
    pub source: MultiInputTokenDataWithContext, // Complete compressed account data with merkle context
    pub destination_index: u8,                  // Destination ctoken Solana account (must exist)
    /// TLV extensions for this compressed account (e.g., CompressedOnly extension).
    /// Used to transfer extension state during decompress.
    pub tlv: Option<Vec<ExtensionInstructionData>>,
    /// Whether this is an ATA decompression. For ATAs, the source.owner is the ATA address
    /// (not the wallet), so it should NOT be marked as a signer - the wallet signs the tx instead.
    pub is_ata: bool,
}

/// Decompress full balance from compressed token accounts with pre-computed indices
///
/// # Arguments
/// * `fee_payer` - The fee payer pubkey
/// * `validity_proof` - Validity proof for the compressed accounts (zkp or index)
/// * `cpi_context_pubkey` - Optional CPI context account for optimized multi-program transactions
/// * `indices` - Slice of source/destination pairs for decompress operations
/// * `packed_accounts` - Slice of all accounts that will be used in the instruction
///
/// # Returns
/// An instruction that decompresses the full balance of all provided token accounts
#[profile]
pub fn decompress_full_token_accounts_with_indices<'info>(
    fee_payer: Pubkey,
    validity_proof: ValidityProof,
    cpi_context_pubkey: Option<Pubkey>,
    indices: &[DecompressFullIndices],
    packed_accounts: &[AccountInfo<'info>],
) -> Result<Instruction, TokenSdkError> {
    if indices.is_empty() {
        return Err(TokenSdkError::InvalidAccountData);
    }

    // Process each set of indices
    let mut token_accounts = Vec::with_capacity(indices.len());
    let mut in_tlv_data: Vec<Vec<ExtensionInstructionData>> = Vec::with_capacity(indices.len());
    let mut has_any_tlv = false;

    // Convert packed_accounts to AccountMetas
    // TODO: we may have to add conditional delegate signers for delegate
    // support via CPI.
    // Build signer flags in O(n) instead of scanning on every meta push
    let mut signer_flags = vec![false; packed_accounts.len()];

    for idx in indices.iter() {
        // Create CTokenAccount2 with the source data
        // For decompress_full, we don't have an output tree since everything goes to the destination
        let mut token_account = CTokenAccount2::new(vec![idx.source])?;

        // Set up decompress_full - decompress entire balance to destination ctoken account
        token_account.decompress(idx.source.amount, idx.destination_index)?;
        token_accounts.push(token_account);

        // Collect TLV data for this input
        if let Some(tlv) = &idx.tlv {
            has_any_tlv = true;
            in_tlv_data.push(tlv.clone());
        } else {
            in_tlv_data.push(Vec::new());
        }

        let owner_idx = idx.source.owner as usize;
        if owner_idx >= signer_flags.len() {
            return Err(TokenSdkError::InvalidAccountData);
        }
        // For ATAs, the owner is the ATA address (a PDA that can't sign).
        // The wallet signs the transaction instead, so don't mark the owner as signer.
        if !idx.is_ata {
            signer_flags[owner_idx] = true;
        }
    }

    let mut packed_account_metas = Vec::with_capacity(packed_accounts.len());

    for (i, info) in packed_accounts.iter().enumerate() {
        packed_account_metas.push(AccountMeta {
            pubkey: *info.key,
            is_signer: info.is_signer || signer_flags[i],
            is_writable: info.is_writable,
        });
    }
    let (meta_config, transfer_config) = if let Some(cpi_context) = cpi_context_pubkey {
        let cpi_context_config = CompressedCpiContext {
            set_context: false,
            first_set_context: false,
        };

        (
            Transfer2AccountsMetaConfig {
                fee_payer: Some(fee_payer),
                cpi_context: Some(cpi_context),
                decompressed_accounts_only: false,
                sol_pool_pda: None,
                sol_decompression_recipient: None,
                with_sol_pool: false,
                packed_accounts: Some(packed_account_metas),
            },
            Transfer2Config::default()
                .filter_zero_amount_outputs()
                .with_cpi_context(cpi_context_config),
        )
    } else {
        (
            Transfer2AccountsMetaConfig::new(fee_payer, packed_account_metas),
            Transfer2Config::default().filter_zero_amount_outputs(),
        )
    };

    // Create the transfer2 instruction with all decompress operations
    let inputs = Transfer2Inputs {
        meta_config,
        token_accounts,
        transfer_config,
        validity_proof,
        in_tlv: if has_any_tlv { Some(in_tlv_data) } else { None },
        ..Default::default()
    };

    create_transfer2_instruction(inputs)
}

/// Helper function to pack compressed token accounts into DecompressFullIndices
/// Used in tests to build indices for multiple compressed accounts to decompress
///
/// # Arguments
/// * `token_data` - Slice of TokenData from compressed accounts
/// * `tree_infos` - Packed tree info for each compressed account
/// * `destination_indices` - Destination account indices for each decompression
/// * `packed_accounts` - PackedAccounts that will be used to insert/get indices
/// * `tlv` - Optional TLV extensions for the compressed account
/// * `version` - TokenDataVersion (1=V1, 2=V2, 3=ShaFlat) for hash computation
///
/// # Returns
/// Vec of DecompressFullIndices ready to use with decompress_full_token_accounts_with_indices
#[profile]
pub fn pack_for_decompress_full(
    token: &TokenData,
    tree_info: &PackedStateTreeInfo,
    destination: Pubkey,
    packed_accounts: &mut PackedAccounts,
    tlv: Option<Vec<ExtensionInstructionData>>,
    version: u8,
) -> DecompressFullIndices {
    let source = MultiInputTokenDataWithContext {
        owner: packed_accounts.insert_or_get_config(token.owner, true, false),
        amount: token.amount,
        has_delegate: token.delegate.is_some(),
        delegate: token
            .delegate
            .map(|d| packed_accounts.insert_or_get(d))
            .unwrap_or(0),
        mint: packed_accounts.insert_or_get(token.mint),
        version,
        merkle_context: PackedMerkleContext {
            merkle_tree_pubkey_index: tree_info.merkle_tree_pubkey_index,
            queue_pubkey_index: tree_info.queue_pubkey_index,
            prove_by_index: tree_info.prove_by_index,
            leaf_index: tree_info.leaf_index,
        },
        root_index: tree_info.root_index,
    };

    DecompressFullIndices {
        source,
        destination_index: packed_accounts.insert_or_get(destination),
        tlv,
        is_ata: false, // Non-ATA: owner is a signer
    }
}

/// Pack accounts for decompress with ATA support.
///
/// For ATA decompress (is_ata=true):
/// - Owner (ATA pubkey) is added without signer flag (ATA can't sign)
/// - Wallet owner is already added as signer by the caller
///
/// For non-ATA decompress:
/// - Owner is added as signer (normal case)
#[profile]
pub fn pack_for_decompress_full_with_ata(
    token: &TokenData,
    tree_info: &PackedStateTreeInfo,
    destination: Pubkey,
    packed_accounts: &mut PackedAccounts,
    tlv: Option<Vec<ExtensionInstructionData>>,
    version: u8,
    is_ata: bool,
) -> DecompressFullIndices {
    // For ATA: owner (ATA pubkey) is not a signer - wallet owner signs instead
    // For non-ATA: owner is a signer
    let owner_is_signer = !is_ata;

    let source = MultiInputTokenDataWithContext {
        owner: packed_accounts.insert_or_get_config(token.owner, owner_is_signer, false),
        amount: token.amount,
        has_delegate: token.delegate.is_some(),
        delegate: token
            .delegate
            .map(|d| packed_accounts.insert_or_get(d))
            .unwrap_or(0),
        mint: packed_accounts.insert_or_get(token.mint),
        version,
        merkle_context: PackedMerkleContext {
            merkle_tree_pubkey_index: tree_info.merkle_tree_pubkey_index,
            queue_pubkey_index: tree_info.queue_pubkey_index,
            prove_by_index: tree_info.prove_by_index,
            leaf_index: tree_info.leaf_index,
        },
        root_index: tree_info.root_index,
    };

    DecompressFullIndices {
        source,
        destination_index: packed_accounts.insert_or_get(destination),
        tlv,
        is_ata,
    }
}

pub struct DecompressFullAccounts {
    pub compressed_token_program: Pubkey,
    pub cpi_authority_pda: Pubkey,
    pub cpi_context: Option<Pubkey>,
    pub self_program: Option<Pubkey>,
}

impl DecompressFullAccounts {
    pub fn new(cpi_context: Option<Pubkey>) -> Self {
        Self {
            compressed_token_program: TokenDefaultAccounts::default().compressed_token_program,
            cpi_authority_pda: TokenDefaultAccounts::default().cpi_authority_pda,
            cpi_context,
            self_program: None,
        }
    }
    pub fn new_with_cpi_context(cpi_context: Option<Pubkey>, self_program: Option<Pubkey>) -> Self {
        Self {
            compressed_token_program: TokenDefaultAccounts::default().compressed_token_program,
            cpi_authority_pda: TokenDefaultAccounts::default().cpi_authority_pda,
            cpi_context,
            self_program,
        }
    }
}

impl AccountMetasVec for DecompressFullAccounts {
    /// Adds:
    /// 1. system accounts if not set
    /// 2. compressed token program and ctoken cpi authority pda to pre accounts
    fn get_account_metas_vec(&self, accounts: &mut PackedAccounts) -> Result<(), LightSdkError> {
        if !accounts.system_accounts_set() {
            #[cfg(feature = "cpi-context")]
            let config = {
                let mut config = SystemAccountMetaConfig::default();
                config.self_program = self.self_program;
                config.cpi_context = self.cpi_context;
                config
            };
            #[cfg(not(feature = "cpi-context"))]
            let config = {
                let mut config = SystemAccountMetaConfig::default();
                config.self_program = self.self_program;
                config
            };

            accounts.add_system_accounts_v2(config)?;
        }
        // Add both accounts in one operation for better performance
        accounts.pre_accounts.extend_from_slice(&[
            AccountMeta {
                pubkey: self.compressed_token_program,
                is_signer: false,
                is_writable: false,
            },
            AccountMeta {
                pubkey: self.cpi_authority_pda,
                is_signer: false,
                is_writable: false,
            },
        ]);
        Ok(())
    }
}