autonomi 0.9.0

// Copyright 2025 MaidSafe.net limited.
//
// This SAFE Network Software is licensed to you under The General Public License (GPL), version 3.
// Unless required by applicable law or agreed to in writing, the SAFE Network Software distributed
// under the GPL Licence is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. Please review the Licences for the specific language governing
// permissions and limitations relating to use of the SAFE Network Software.

use super::payments::{MerklePaymentError, MerklePaymentReceipt};
use super::upload::MerklePutError;
use crate::Client;
use crate::client::config::CHUNK_UPLOAD_BATCH_SIZE;
use crate::client::data_types::chunk::DataMapChunk;
use crate::client::files::Metadata;
use crate::client::{ClientEvent, UploadSummary};
use crate::self_encryption::{EncryptionStream, MAX_CHUNK_SIZE, encrypt_directory_files};
use ant_evm::merkle_payments::MAX_LEAVES;
use ant_evm::{AttoTokens, EvmWallet};
use ant_protocol::NetworkAddress;
use ant_protocol::storage::{ChunkAddress, DataTypes};
use std::collections::{HashMap, HashSet};
use std::path::PathBuf;
use thiserror::Error;
use xor_name::XorName;

/// Payment option for Merkle batch uploads
#[derive(Clone)]
pub enum MerklePaymentOption<'a> {
    /// Fresh upload - pays for all chunks
    Wallet(&'a EvmWallet),
    /// Upload with External Payment Flow - assumes all proofs present, fails if not
    Receipt(MerklePaymentReceipt),
    /// Continue/Retry upload with partial payment receipt - uses existing proofs, pays for any missing chunks
    ContinueWithReceipt(&'a EvmWallet, MerklePaymentReceipt),
}

/// Error with optional receipt attached
/// Receipt is only present if payments were made before the error occurred
#[derive(Debug, Error)]
#[error("{error}")]
pub struct MerkleUploadErrorWithReceipt {
    /// Receipt if any payments were made before failure (None = no payment happened)
    pub receipt: Option<MerklePaymentReceipt>,
    /// The actual error details
    #[source]
    pub error: MerkleUploadError,
}

impl MerkleUploadErrorWithReceipt {
    /// Create error, only including receipt if it contains actual payments
    fn new(receipt: MerklePaymentReceipt, kind: MerkleUploadError) -> Self {
        let receipt = if receipt.proofs.is_empty() {
            None // No payments made
        } else {
            Some(receipt) // Real payments - include proof
        };
        Self {
            receipt,
            error: kind,
        }
    }

    fn encryption(receipt: MerklePaymentReceipt, msg: String) -> Self {
        Self::new(receipt, MerkleUploadError::Encryption(msg))
    }

    fn payment(receipt: MerklePaymentReceipt, err: MerklePaymentError) -> Self {
        Self::new(receipt, MerkleUploadError::Payment(err))
    }

    fn upload(receipt: MerklePaymentReceipt, err: MerklePutError) -> Self {
        Self::new(receipt, MerkleUploadError::Upload(err))
    }
}

#[derive(Debug, Error)]
pub enum MerkleUploadError {
    #[error("Encryption error: {0}")]
    Encryption(String),
    #[error("Payment error: {0}")]
    Payment(MerklePaymentError),
    #[error("Upload error: {0}")]
    Upload(MerklePutError),
}

impl Client {
    /// Estimate the cost of uploading a directory of files with Merkle batch payments
    ///
    /// This calls the smart contract's view function (0 gas) which runs the exact same
    /// pricing logic as the actual payment, ensuring accurate cost estimation.
    ///
    /// # Arguments
    /// * `path` - The path to the directory
    /// * `is_public` - Whether the files will be uploaded as public
    /// * `wallet` - The EVM wallet (needed for network checking)
    ///
    /// # Returns
    /// * `AttoTokens` - Estimated total cost
    pub async fn file_cost_merkle(
        &self,
        path: PathBuf,
        is_public: bool,
        wallet: &EvmWallet,
    ) -> Result<AttoTokens, MerkleUploadError> {
        debug!(
            "merkle payment: file_cost_merkle starting for path: {path:?}, is_public: {is_public}"
        );

        // Check if the wallet uses the same network as the client
        if wallet.network() != self.evm_network() {
            return Err(MerkleUploadError::Payment(
                MerklePaymentError::EvmWalletNetworkMismatch,
            ));
        }

        #[cfg(feature = "loud")]
        println!("Encrypting files to calculate cost...");

        // Collect all XorNames
        let (all_xor_names, _file_chunk_counts, _file_results) = self
            .collect_xornames_from_dir(path, is_public)
            .await
            .map_err(MerkleUploadError::Encryption)?;

        let total_chunks = all_xor_names.len();
        debug!("merkle payment: file_cost_merkle total chunks: {total_chunks}");

        #[cfg(feature = "loud")]
        println!("Encrypted into {total_chunks} chunks");

        // Split into batches of MAX_LEAVES
        let batches: Vec<Vec<XorName>> = all_xor_names
            .chunks(MAX_LEAVES)
            .map(|c| c.to_vec())
            .collect();
        let num_batches = batches.len();

        #[cfg(feature = "loud")]
        println!("Estimating cost for {num_batches} batch(es)...");

        // Estimate cost for each batch and sum
        let mut total_cost = ant_evm::U256::ZERO;

        for (batch_idx, batch_xornames) in batches.into_iter().enumerate() {
            let batch_num = batch_idx + 1;
            debug!("Estimating batch {batch_num}/{num_batches}");

            // Prepare batch (build tree, query pools)
            let (tree, _candidate_pools, pool_commitments, merkle_payment_timestamp) = self
                .prepare_merkle_batch(DataTypes::Chunk, batch_xornames, MAX_CHUNK_SIZE)
                .await
                .map_err(MerkleUploadError::Payment)?;

            // Estimate cost for this batch
            let batch_cost = wallet
                .estimate_merkle_payment_cost(
                    tree.depth(),
                    &pool_commitments,
                    merkle_payment_timestamp,
                )
                .await
                .map_err(|e| MerkleUploadError::Payment(MerklePaymentError::EvmWalletError(e)))?;

            total_cost = total_cost.saturating_add(batch_cost);
        }

        let estimated_cost = AttoTokens::from_atto(total_cost);
        debug!("merkle payment: file_cost_merkle estimated total cost: {estimated_cost}");

        #[cfg(feature = "loud")]
        println!("Total estimated cost: {estimated_cost}");

        Ok(estimated_cost)
    }

    /// Helper function to pay for a directory of files with Merkle batch payments
    async fn files_put_with_merkle_payment_internal(
        &self,
        path: PathBuf,
        is_public: bool,
        wallet: Option<&EvmWallet>,
        mut receipt: MerklePaymentReceipt,
    ) -> Result<(AttoTokens, Vec<(PathBuf, DataMapChunk, Metadata)>), MerkleUploadErrorWithReceipt>
    {
        debug!("merkle payment: starting for path: {path:?}, is_public: {is_public}");

        // Check wallet network (if wallet provided)
        if let Some(w) = wallet
            && w.network() != self.evm_network()
        {
            return Err(MerkleUploadErrorWithReceipt::payment(
                receipt,
                MerklePaymentError::EvmWalletNetworkMismatch,
            ));
        }

        // Encrypt files to collect ALL XorNames
        #[cfg(feature = "loud")]
        println!("Encrypting files a first time to create the Merkle Tree(s)...");
        let (all_xor_names, file_chunk_counts, first_pass_results) = self
            .collect_xornames_from_dir(path.clone(), is_public)
            .await
            .map_err(|e| MerkleUploadErrorWithReceipt::encryption(receipt.clone(), e))?;
        receipt.file_chunk_counts = file_chunk_counts;
        let total_files = receipt.file_chunk_counts.len();

        let total_chunks = all_xor_names.len();
        info!("Collected {total_chunks} XorNames from {total_files} files");

        // Check which chunks already exist on the network and split into existing/new
        let (mut already_exist, mut xor_to_pay_ordered) =
            self.split_existing_and_new_chunks(all_xor_names).await;
        let to_pay_len = xor_to_pay_ordered.len();
        let already_paid_count = already_exist.len();

        // Early return if all chunks already exist - no need for second encryption
        if to_pay_len == 0 {
            #[cfg(feature = "loud")]
            println!("✓ All {total_chunks} chunks already exist on the network, nothing to upload");

            // Send upload completion event
            self.send_upload_complete_event(&receipt, already_paid_count)
                .await;

            return Ok((receipt.amount_paid, first_pass_results));
        } else {
            std::mem::drop(first_pass_results);
        }

        // Split into batches of MAX_LEAVES - using drain to avoid cloning
        let num_batches = to_pay_len.div_ceil(MAX_LEAVES);
        let mut batches: Vec<Vec<XorName>> = Vec::with_capacity(num_batches);
        while !xor_to_pay_ordered.is_empty() {
            let drain_count = std::cmp::min(MAX_LEAVES, xor_to_pay_ordered.len());
            batches.push(xor_to_pay_ordered.drain(..drain_count).collect());
        }
        info!("Split into {num_batches} Merkle Tree(s) of up to {MAX_LEAVES} chunks each");

        // Start upload streams (second encryption pass - needed to get actual chunk data)
        #[cfg(feature = "loud")]
        println!("🚀 Starting upload of {to_pay_len} chunks in {num_batches} Merkle Tree(s)...");
        let mut streams: Vec<EncryptionStream> = encrypt_directory_files(path, is_public)
            .await
            .map_err(|e| MerkleUploadErrorWithReceipt::encryption(receipt.clone(), e.to_string()))?
            .into_iter()
            .map(|stream| {
                stream.map_err(|e| MerkleUploadErrorWithReceipt::encryption(receipt.clone(), e))
            })
            .collect::<Result<Vec<EncryptionStream>, MerkleUploadErrorWithReceipt>>()?;

        let mut results: Vec<(PathBuf, DataMapChunk, Metadata)> = Vec::new();

        // Interleaved pay/upload for each batch
        for (batch_idx, batch_xornames) in batches.into_iter().enumerate() {
            let batch_num = batch_idx + 1;
            let batch_size = batch_xornames.len();
            info!("Processing batch {batch_num}/{num_batches} ({batch_size} chunks)");

            // Pay for this batch if needed
            let needs_payment = batch_xornames
                .iter()
                .any(|xn| !receipt.proofs.contains_key(xn));
            if needs_payment {
                receipt = self
                    .pay_for_merkle_tree_batch(
                        wallet,
                        batch_xornames,
                        receipt.clone(),
                        batch_num,
                        num_batches,
                    )
                    .await
                    .map_err(|kind| MerkleUploadErrorWithReceipt::new(receipt.clone(), kind))?;
            }

            #[cfg(feature = "loud")]
            println!("🌳 Merkle Tree {batch_num}/{num_batches}: Uploading {batch_size} chunks...");

            // Upload this batch's chunks (skip chunks that already exist)
            let upload_result = self
                .upload_batch_with_merkle(streams, &receipt, &mut already_exist, batch_size)
                .await
                .map_err(|err| MerkleUploadErrorWithReceipt::upload(receipt.clone(), err))?;

            streams = upload_result.streams;
            results.extend(upload_result.completed_files);

            // Retry failed chunks if any
            if !upload_result.failed_chunks.is_empty() {
                const MAX_RETRIES: usize = 3;
                const RETRY_PAUSE_SECS: u64 = 60;

                let remaining_failures = self
                    .retry_failed_merkle_chunks(
                        upload_result.failed_chunks,
                        &receipt,
                        &mut already_exist,
                        MAX_RETRIES,
                        RETRY_PAUSE_SECS,
                    )
                    .await
                    .map_err(|err| MerkleUploadErrorWithReceipt::upload(receipt.clone(), err))?;

                if !remaining_failures.is_empty() {
                    let failed_count = remaining_failures.len();
                    error!("{failed_count} chunks failed after {MAX_RETRIES} retries");
                    return Err(MerkleUploadErrorWithReceipt::upload(
                        receipt,
                        MerklePutError::Batch(super::upload::MerkleBatchUploadState {
                            failed: remaining_failures,
                        }),
                    ));
                }
            }

            info!(
                "Batch {batch_num}/{num_batches} complete, {} files finished so far",
                results.len()
            );
        }

        // Handle any remaining streams
        for mut stream in streams {
            let datamap = if let Some(datamap) = stream.data_map_chunk() {
                datamap
            } else {
                // flush the stream of remaining duplicate chunks to force transition to StreamDone state and try again
                while stream.next_batch(16).is_some() {}

                stream
                    .data_map_chunk()
                    .ok_or(MerkleUploadErrorWithReceipt::upload(
                        receipt.clone(),
                        MerklePutError::StreamShouldHaveDatamap,
                    ))?
            };

            // add the datamap to the results
            results.push((
                stream.relative_path.clone(),
                datamap,
                stream.metadata.clone(),
            ));

            // report progress
            if let Some(public_addr) = stream.data_address() {
                let path = &stream.relative_path;
                let f = results.len();
                debug!("[File {f}/{total_files}] ({path:?}) is now available at: {public_addr:?}");
                #[cfg(feature = "loud")]
                println!(
                    "[File {f}/{total_files}] ({path:?}) is now available at: {public_addr:?}"
                );
            }
        }

        debug!(
            "merkle payment: {total_chunks} chunks uploaded for {total_files} files successfully"
        );
        #[cfg(feature = "loud")]
        println!("✓ All {total_chunks} chunks uploaded successfully!");

        // Send upload completion event
        self.send_upload_complete_event(&receipt, already_paid_count)
            .await;

        Ok((receipt.amount_paid, results))
    }

    /// Upload a directory of files with Merkle batch payments
    /// It is very important that the files are not changed while they are being uploaded as it could invalidate the Merkle payment.
    ///
    /// # Arguments
    /// * `path` - The path to the directory to upload
    /// * `is_public` - Whether the files are uploaded as public
    /// * `payment` - The payment option (wallet or cached receipt)
    ///
    /// # Returns
    /// * Tuple of (amount_paid, results) where:
    ///   - `amount_paid` - Total amount paid for the Merkle batch (in AttoTokens)
    ///   - `results` - Vector of (relative_path, datamap, metadata) tuples for each uploaded file
    ///
    /// # Errors
    /// On error, check `error.receipt` for any payments made before the failure.
    /// If `Some(receipt)`, payments were made and can be reused via [`MerklePaymentOption::ContinueWithReceipt`].
    pub async fn files_put_with_merkle_payment(
        &self,
        path: PathBuf,
        is_public: bool,
        payment: MerklePaymentOption<'_>,
    ) -> Result<(AttoTokens, Vec<(PathBuf, DataMapChunk, Metadata)>), MerkleUploadErrorWithReceipt>
    {
        debug!(
            "merkle payment: files_put starting upload for path: {path:?}, is_public: {is_public}"
        );

        match payment {
            MerklePaymentOption::Wallet(wallet) => {
                self.files_put_with_merkle_payment_internal(
                    path,
                    is_public,
                    Some(wallet),
                    MerklePaymentReceipt::default(),
                )
                .await
            }
            MerklePaymentOption::Receipt(receipt) => {
                self.files_put_with_merkle_payment_internal(path, is_public, None, receipt)
                    .await
            }
            MerklePaymentOption::ContinueWithReceipt(wallet, receipt) => {
                self.files_put_with_merkle_payment_internal(path, is_public, Some(wallet), receipt)
                    .await
            }
        }
    }

    /// Collect all XorNames from a directory, returning (all_xornames, file_chunk_counts, file_results)
    /// The file_results contain (relative_path, datamap, metadata) for each file, useful when all chunks already exist.
    async fn collect_xornames_from_dir(
        &self,
        path: PathBuf,
        is_public: bool,
    ) -> Result<
        (
            Vec<XorName>,
            HashMap<String, usize>,
            Vec<(PathBuf, DataMapChunk, Metadata)>,
        ),
        String,
    > {
        let streams: Vec<EncryptionStream> = encrypt_directory_files(path, is_public)
            .await
            .map_err(|e| e.to_string())?
            .into_iter()
            .collect::<Result<Vec<EncryptionStream>, String>>()?;

        let mut all_xor_names = Vec::new();
        let mut file_chunk_counts = HashMap::new();
        let mut file_results = Vec::new();

        for stream in streams {
            let file_path = stream.file_path.clone();
            let (xor_names, relative_path, datamap, metadata) =
                collect_xor_names_from_stream(stream)?;
            file_chunk_counts.insert(file_path, xor_names.len());
            all_xor_names.extend(xor_names);
            file_results.push((relative_path, datamap, metadata));
        }

        Ok((all_xor_names, file_chunk_counts, file_results))
    }

    /// Split chunks into existing and new, checking the network for existence.
    ///
    /// Performs a fast parallel existence check and splits the input into:
    /// - set of existing chunks (already on network, no payment needed)
    /// - vec of new chunks (need payment), in original order with duplicates removed
    ///
    /// Takes ownership of the input to avoid unnecessary copies.
    async fn split_existing_and_new_chunks(
        &self,
        xornames: Vec<XorName>,
    ) -> (HashSet<XorName>, Vec<XorName>) {
        #[cfg(feature = "loud")]
        println!("Checking for existing chunks on the network...");
        debug!("Checking for existing chunks on the network...");

        // Dedupe while preserving order (keep first occurrence only)
        let mut seen: HashSet<XorName> = HashSet::new();
        let unique_ordered: Vec<XorName> =
            xornames.into_iter().filter(|xn| seen.insert(*xn)).collect();
        let total = unique_ordered.len();

        // Check existence on network
        let addresses: Vec<NetworkAddress> = unique_ordered
            .iter()
            .map(|xn| NetworkAddress::from(ChunkAddress::new(*xn)))
            .collect();
        let batch_size = std::cmp::max(16, *CHUNK_UPLOAD_BATCH_SIZE);
        let existing_addrs = self.check_records_exist_batch(&addresses, batch_size).await;

        // Convert to XorName set
        let existing_set: HashSet<XorName> = existing_addrs
            .into_iter()
            .filter_map(|addr| addr.xorname())
            .collect();

        let existing_count = existing_set.len();
        info!("Found {existing_count}/{total} unique chunks already exist on the network");
        #[cfg(feature = "loud")]
        println!("Found {existing_count}/{total} unique chunks already exist on the network");

        // Filter out existing, keeping original order
        let new_chunks: Vec<XorName> = unique_ordered
            .into_iter()
            .filter(|xn| !existing_set.contains(xn))
            .collect();

        (existing_set, new_chunks)
    }

    /// Pay for a Merkle Tree batch, returning the merged receipt
    async fn pay_for_merkle_tree_batch(
        &self,
        wallet: Option<&EvmWallet>,
        batch_xornames: Vec<XorName>,
        mut receipt: MerklePaymentReceipt,
        batch_num: usize,
        num_batches: usize,
    ) -> Result<MerklePaymentReceipt, MerkleUploadError> {
        // Need wallet to pay - error if Receipt variant (wallet is None)
        let w = wallet.ok_or_else(|| {
            let missing_xn = batch_xornames
                .iter()
                .find(|xn| !receipt.proofs.contains_key(xn))
                .copied()
                .unwrap_or_default();
            MerkleUploadError::Upload(MerklePutError::MissingPaymentProofFor(missing_xn))
        })?;

        let batch_size = batch_xornames.len();
        debug!("Merkle Tree {batch_num}/{num_batches}: Paying for {batch_size} chunks...");
        #[cfg(feature = "loud")]
        println!("💸 Merkle Tree {batch_num}/{num_batches}: Paying for {batch_size} chunks...");

        let batch_receipt = self
            .pay_for_single_merkle_batch(DataTypes::Chunk, batch_xornames, MAX_CHUNK_SIZE, w)
            .await
            .map_err(MerkleUploadError::Payment)?;

        receipt.merge(batch_receipt);
        Ok(receipt)
    }

    /// Send upload completion event
    async fn send_upload_complete_event(
        &self,
        receipt: &MerklePaymentReceipt,
        records_already_paid: usize,
    ) {
        // Use proofs count for accurate records_paid (handles duplicates correctly)
        let records_paid = receipt.proofs.len();
        if let Some(sender) = &self.client_event_sender {
            let summary = UploadSummary {
                records_paid,
                records_already_paid,
                tokens_spent: receipt.amount_paid.as_atto(),
            };

            if let Err(err) = sender.send(ClientEvent::UploadComplete(summary)).await {
                error!("Failed to send upload completion event: {err:?}");
            }
        }
    }
}

/// Collect all XorNames from a stream, returning (xornames, relative_path, datamap, metadata)
fn collect_xor_names_from_stream(
    mut encryption_stream: EncryptionStream,
) -> Result<(Vec<XorName>, PathBuf, DataMapChunk, Metadata), String> {
    let mut xor_names: Vec<XorName> = Vec::new();
    let xorname_collection_batch_size: usize = std::cmp::max(32, *CHUNK_UPLOAD_BATCH_SIZE);
    let mut total = 0;
    let estimated_total = encryption_stream.total_chunks();
    let file_path = encryption_stream.file_path.clone();
    let start = std::time::Instant::now();
    #[cfg(feature = "loud")]
    println!("Begin encrypting ~{estimated_total} chunks from {file_path}...");
    debug!("Begin encrypting ~{estimated_total} chunks from {file_path}...");
    while let Some(batch) = encryption_stream.next_batch(xorname_collection_batch_size) {
        let batch_len = batch.len();
        total += batch_len;
        for chunk in batch {
            xor_names.push(*chunk.name());
        }
        #[cfg(feature = "loud")]
        println!(
            "Encrypted {total}/{estimated_total} chunks in {:?}",
            start.elapsed()
        );
        debug!(
            "Encrypted {total}/{estimated_total} chunks in {:?}",
            start.elapsed()
        );
    }

    // Extract the datamap now that the stream is drained (in StreamDone or InMemory state)
    let datamap = encryption_stream
        .data_map_chunk()
        .ok_or_else(|| format!("No datamap available for {file_path}"))?;

    Ok((
        xor_names,
        encryption_stream.relative_path,
        datamap,
        encryption_stream.metadata,
    ))
}