tenzro-token 0.2.0

//! Network treasury management
//!
//! This module implements the network treasury that accumulates fees from settlements,
//! transactions, and bridge operations in multiple assets (stablecoins + crypto).

use crate::error::{Result, TokenError};
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use std::sync::Arc;
use tenzro_crypto::{PublicKey, Signature};
use tenzro_storage::kv::{CF_METADATA, KvStore};
use tenzro_types::asset::AssetId;
use tenzro_types::primitives::Address;
use tracing::{debug, info, warn};

/// Domain separator for treasury withdrawal-approval signatures.
const WITHDRAWAL_APPROVAL_DOMAIN: &[u8] = b"tenzro/treasury/withdrawal-approval";

/// A pending multisig withdrawal: the payload is pinned by the first approval
/// so every subsequent approval and the final execution are bound to the same
/// (asset, amount) pair.
#[derive(Debug, Clone)]
pub struct PendingWithdrawal {
    /// Asset being withdrawn
    pub asset_id: AssetId,
    /// Amount being withdrawn
    pub amount: u128,
    /// Addresses that have approved so far
    pub approvers: Vec<Address>,
}

/// Canonical signing preimage for a withdrawal approval.
///
/// `domain || withdrawal_id || asset_id || amount_le` — approvers sign these
/// bytes with the key that derives their authorized withdrawer address.
pub fn withdrawal_approval_preimage(
    withdrawal_id: &str,
    asset_id: &AssetId,
    amount: u128,
) -> Vec<u8> {
    let mut preimage = Vec::with_capacity(
        WITHDRAWAL_APPROVAL_DOMAIN.len() + withdrawal_id.len() + asset_id.as_str().len() + 16,
    );
    preimage.extend_from_slice(WITHDRAWAL_APPROVAL_DOMAIN);
    preimage.extend_from_slice(withdrawal_id.as_bytes());
    preimage.extend_from_slice(asset_id.as_str().as_bytes());
    preimage.extend_from_slice(&amount.to_le_bytes());
    preimage
}

/// Checks that a public key derives to the given canonical 32-byte address.
///
/// Convention: the 20-byte address derived from the key is left-aligned in the
/// 32-byte slot with a zero-padded tail.
fn public_key_binds_address(public_key: &PublicKey, address: &Address) -> bool {
    let derived = public_key.to_address();
    let addr_bytes = address.as_bytes();
    addr_bytes[..20] == *derived.as_bytes() && addr_bytes[20..].iter().all(|b| *b == 0)
}

/// Fee distribution configuration
///
/// Defines how network fees are split between treasury, burn, and staker rewards.
/// All values are in basis points (bps) and must sum to 10000 (100%).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FeeDistributionConfig {
    /// Percentage of fees going to treasury (basis points)
    pub treasury_share_bps: u32,
    /// Percentage of fees to burn (basis points)
    pub burn_share_bps: u32,
    /// Percentage of fees going to stakers (basis points)
    pub staker_share_bps: u32,
}

impl Default for FeeDistributionConfig {
    fn default() -> Self {
        Self {
            treasury_share_bps: 4000, // 40%
            burn_share_bps: 3000,     // 30%
            staker_share_bps: 3000,   // 30%
        }
    }
}

impl FeeDistributionConfig {
    /// Validates that the configuration is correct
    pub fn validate(&self) -> Result<()> {
        let total = self
            .treasury_share_bps
            .saturating_add(self.burn_share_bps)
            .saturating_add(self.staker_share_bps);

        if total != 10000 {
            return Err(TokenError::InvalidConfig(format!(
                "Fee distribution must sum to 10000 bps, got {}",
                total
            )));
        }

        Ok(())
    }

    /// Calculates the distribution amounts for a given fee
    pub fn calculate_distribution(&self, total_fee: u128) -> FeeDistribution {
        let treasury_amount = total_fee * (self.treasury_share_bps as u128) / 10000;
        let burn_amount = total_fee * (self.burn_share_bps as u128) / 10000;
        let staker_amount = total_fee * (self.staker_share_bps as u128) / 10000;

        FeeDistribution {
            treasury_amount,
            burn_amount,
            staker_amount,
        }
    }
}

/// Calculated fee distribution
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FeeDistribution {
    /// Amount going to treasury
    pub treasury_amount: u128,
    /// Amount to burn
    pub burn_amount: u128,
    /// Amount for stakers
    pub staker_amount: u128,
}

/// RocksDB-backed storage for treasury state
pub struct TreasuryStorageBackend {
    store: Arc<dyn KvStore>,
}

impl TreasuryStorageBackend {
    /// Creates a new treasury storage backend
    pub fn new(store: Arc<dyn KvStore>) -> Self {
        Self { store }
    }

    fn balance_key(asset_id: &AssetId) -> Vec<u8> {
        let mut key = b"treasury:balance:".to_vec();
        key.extend_from_slice(asset_id.as_str().as_bytes());
        key
    }

    fn collected_key(asset_id: &AssetId) -> Vec<u8> {
        let mut key = b"treasury:collected:".to_vec();
        key.extend_from_slice(asset_id.as_str().as_bytes());
        key
    }

    fn distributed_key(asset_id: &AssetId) -> Vec<u8> {
        let mut key = b"treasury:distributed:".to_vec();
        key.extend_from_slice(asset_id.as_str().as_bytes());
        key
    }

    fn get_u128(&self, key: &[u8]) -> Result<u128> {
        match self.store.get(CF_METADATA, key)? {
            Some(bytes) if bytes.len() == 16 => {
                let array: [u8; 16] = bytes.try_into().unwrap();
                Ok(u128::from_le_bytes(array))
            }
            Some(bytes) => {
                warn!("Invalid treasury bytes length: {}", bytes.len());
                Ok(0)
            }
            None => Ok(0),
        }
    }

    fn set_u128(&self, key: &[u8], value: u128) -> Result<()> {
        self.store.put(CF_METADATA, key, &value.to_le_bytes())?;
        Ok(())
    }

    /// Persists a treasury balance for an asset
    pub fn persist_balance(&self, asset_id: &AssetId, balance: u128) -> Result<()> {
        self.set_u128(&Self::balance_key(asset_id), balance)
    }

    /// Persists total collected for an asset
    pub fn persist_collected(&self, asset_id: &AssetId, collected: u128) -> Result<()> {
        self.set_u128(&Self::collected_key(asset_id), collected)
    }

    /// Persists total distributed for an asset
    pub fn persist_distributed(&self, asset_id: &AssetId, distributed: u128) -> Result<()> {
        self.set_u128(&Self::distributed_key(asset_id), distributed)
    }

    /// Loads a treasury balance for an asset
    pub fn load_balance(&self, asset_id: &AssetId) -> Result<u128> {
        self.get_u128(&Self::balance_key(asset_id))
    }

    /// Loads total collected for an asset
    pub fn load_collected(&self, asset_id: &AssetId) -> Result<u128> {
        self.get_u128(&Self::collected_key(asset_id))
    }

    /// Loads total distributed for an asset
    pub fn load_distributed(&self, asset_id: &AssetId) -> Result<u128> {
        self.get_u128(&Self::distributed_key(asset_id))
    }
}

/// Network treasury
///
/// Manages multi-asset treasury accumulating network fees and supporting
/// treasury-based operations.
pub struct NetworkTreasury {
    /// Treasury address
    pub treasury_address: Address,
    /// Multi-asset balances (AssetId -> Balance)
    balances: RwLock<HashMap<AssetId, u128>>,
    /// Total fees collected per asset
    total_fees_collected: RwLock<HashMap<AssetId, u128>>,
    /// Total distributed per asset
    total_distributed: RwLock<HashMap<AssetId, u128>>,
    /// Fee distribution configuration
    distribution_config: RwLock<FeeDistributionConfig>,
    /// Authorized withdrawers
    authorized_withdrawers: RwLock<Vec<Address>>,
    /// Minimum number of approvals required for withdrawal (multisig threshold)
    withdrawal_threshold: RwLock<usize>,
    /// Pending withdrawal approvals (withdrawal_id -> set of approvers)
    pending_approvals: RwLock<HashMap<String, PendingWithdrawal>>,
    /// Optional storage backend for persistence
    storage: Option<Arc<TreasuryStorageBackend>>,
    /// Per-bridge-adapter sponsorship pool balances. Tracks TNZO held
    /// on behalf of users who sponsored cross-chain bridge fees (per
    /// the Cosmos ICS-29 / Hyperlane IGP pattern adapted for Tenzro).
    /// Keyed by the bridge adapter's canonical string id
    /// (`layerzero`, `ccip`, `wormhole`, `debridge`, `hyperlane`,
    /// `axelar`, `lifi`, `canton`). Operators tune `refill_threshold_bps`
    /// via GovernanceEngine; the off-chain rebalancer reads these
    /// fields to decide when to top up the relayer.
    bridge_sponsorship_pools: RwLock<HashMap<String, BridgeSponsorshipPoolRecord>>,
}

/// Per-bridge-adapter sponsorship pool record persisted in the
/// NetworkTreasury. Mirrors the shape exposed by
/// `tenzro_bridge::fee_sponsor::SponsorshipPool`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BridgeSponsorshipPoolRecord {
    pub adapter: String,
    /// 20-byte deterministic vault address — `SHA-256(
    /// "tenzro/bridge/sponsorship-vault" || adapter)[..20]`.
    pub vault_address: [u8; 20],
    pub tnzo_balance_wei: u128,
    pub refill_threshold_bps: u32,
    pub native_outstanding_smallest_unit: u128,
}

impl std::fmt::Debug for NetworkTreasury {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("NetworkTreasury")
            .field("treasury_address", &self.treasury_address)
            .field("storage", &self.storage.as_ref().map(|_| "Some(...)"))
            .finish()
    }
}

impl NetworkTreasury {
    /// Creates a new network treasury (in-memory only)
    pub fn new(treasury_address: Address) -> Self {
        Self {
            treasury_address,
            balances: RwLock::new(HashMap::new()),
            total_fees_collected: RwLock::new(HashMap::new()),
            total_distributed: RwLock::new(HashMap::new()),
            distribution_config: RwLock::new(FeeDistributionConfig::default()),
            authorized_withdrawers: RwLock::new(Vec::new()),
            withdrawal_threshold: RwLock::new(1), // Default: single sig until configured
            pending_approvals: RwLock::new(HashMap::new()),
            storage: None,
            bridge_sponsorship_pools: RwLock::new(HashMap::new()),
        }
    }

    /// Creates a new network treasury with persistent storage
    pub fn with_storage(treasury_address: Address, storage: Arc<TreasuryStorageBackend>) -> Self {
        Self {
            treasury_address,
            balances: RwLock::new(HashMap::new()),
            total_fees_collected: RwLock::new(HashMap::new()),
            total_distributed: RwLock::new(HashMap::new()),
            distribution_config: RwLock::new(FeeDistributionConfig::default()),
            authorized_withdrawers: RwLock::new(Vec::new()),
            withdrawal_threshold: RwLock::new(1),
            pending_approvals: RwLock::new(HashMap::new()),
            storage: Some(storage),
            bridge_sponsorship_pools: RwLock::new(HashMap::new()),
        }
    }

    /// Get a snapshot of all per-bridge sponsorship pools. Read-only.
    pub fn bridge_sponsorship_pools_snapshot(
        &self,
    ) -> HashMap<String, BridgeSponsorshipPoolRecord> {
        self.bridge_sponsorship_pools.read().clone()
    }

    /// Get the per-adapter sponsorship pool record, creating it lazily
    /// from the canonical vault-address derivation.
    pub fn get_or_create_bridge_sponsorship_pool(
        &self,
        adapter: &str,
    ) -> BridgeSponsorshipPoolRecord {
        let mut pools = self.bridge_sponsorship_pools.write();
        if let Some(rec) = pools.get(adapter) {
            return rec.clone();
        }
        // SHA-256("tenzro/bridge/sponsorship-vault" || adapter)[..20]
        let mut h = Sha256::new();
        h.update(b"tenzro/bridge/sponsorship-vault");
        h.update(adapter.as_bytes());
        let d = h.finalize();
        let mut vault = [0u8; 20];
        vault.copy_from_slice(&d[..20]);
        let rec = BridgeSponsorshipPoolRecord {
            adapter: adapter.to_string(),
            vault_address: vault,
            tnzo_balance_wei: 0,
            refill_threshold_bps: 0,
            native_outstanding_smallest_unit: 0,
        };
        pools.insert(adapter.to_string(), rec.clone());
        rec
    }

    /// Credit a user TNZO sponsorship + record the outstanding
    /// destination-native commitment. Called by the bridge-fee
    /// sponsor RPC on a successful debit.
    pub fn sponsor_bridge_fee(
        &self,
        adapter: &str,
        tnzo_paid_wei: u128,
        native_committed_smallest_unit: u128,
    ) {
        let _ = self.get_or_create_bridge_sponsorship_pool(adapter);
        let mut pools = self.bridge_sponsorship_pools.write();
        if let Some(rec) = pools.get_mut(adapter) {
            rec.tnzo_balance_wei = rec.tnzo_balance_wei.saturating_add(tnzo_paid_wei);
            rec.native_outstanding_smallest_unit = rec
                .native_outstanding_smallest_unit
                .saturating_add(native_committed_smallest_unit);
        }
    }

    /// Drain the outstanding-native commitment after the relayer
    /// submits delivery proof. Does NOT touch the TNZO balance — that
    /// is paid out to the relayer via a separate explicit drain step
    /// to keep the audit trail clean.
    pub fn drain_bridge_sponsorship_outstanding(&self, adapter: &str, native_amount: u128) {
        let mut pools = self.bridge_sponsorship_pools.write();
        if let Some(rec) = pools.get_mut(adapter) {
            rec.native_outstanding_smallest_unit = rec
                .native_outstanding_smallest_unit
                .saturating_sub(native_amount);
        }
    }

    /// Set the per-adapter refill threshold (bps). Governance-tunable
    /// via the GovernanceEngine; off-chain rebalancer triggers a
    /// treasury top-up when the relayer's destination-native balance
    /// drops below `outstanding * refill_threshold_bps / 10000`.
    pub fn set_bridge_sponsorship_refill_threshold(&self, adapter: &str, bps: u32) {
        let _ = self.get_or_create_bridge_sponsorship_pool(adapter);
        let mut pools = self.bridge_sponsorship_pools.write();
        if let Some(rec) = pools.get_mut(adapter) {
            rec.refill_threshold_bps = bps;
        }
    }

    /// Persists balance, collected, and distributed for an asset to storage
    fn persist_asset_state(
        &self,
        asset_id: &AssetId,
        balance: u128,
        collected: u128,
        distributed: u128,
    ) {
        if let Some(storage) = &self.storage {
            if let Err(e) = storage.persist_balance(asset_id, balance) {
                warn!("Failed to persist treasury balance: {}", e);
            }
            if let Err(e) = storage.persist_collected(asset_id, collected) {
                warn!("Failed to persist treasury collected: {}", e);
            }
            if let Err(e) = storage.persist_distributed(asset_id, distributed) {
                warn!("Failed to persist treasury distributed: {}", e);
            }
        }
    }

    /// Sets the multisig withdrawal threshold
    ///
    /// Requires that threshold <= number of authorized withdrawers
    pub fn set_withdrawal_threshold(&self, threshold: usize) -> Result<()> {
        let withdrawers = self.authorized_withdrawers.read();
        if threshold == 0 {
            return Err(TokenError::InvalidConfig(
                "Withdrawal threshold must be at least 1".to_string(),
            ));
        }
        if threshold > withdrawers.len() && !withdrawers.is_empty() {
            return Err(TokenError::InvalidConfig(format!(
                "Threshold {} exceeds number of authorized withdrawers {}",
                threshold,
                withdrawers.len()
            )));
        }
        drop(withdrawers);
        *self.withdrawal_threshold.write() = threshold;
        info!("Set withdrawal threshold to {}", threshold);
        Ok(())
    }

    /// Approves a pending withdrawal (multisig step).
    ///
    /// The approver must present a signature over
    /// [`withdrawal_approval_preimage`] made with the key that derives their
    /// authorized withdrawer address. The first approval pins the
    /// (asset, amount) payload; subsequent approvals must match it exactly.
    ///
    /// Returns `true` if the withdrawal has reached threshold and should be executed.
    pub fn approve_withdrawal(
        &self,
        withdrawal_id: &str,
        asset_id: &AssetId,
        amount: u128,
        approver: &Address,
        public_key: &PublicKey,
        signature: &Signature,
    ) -> Result<bool> {
        // Verify the approver is authorized
        let withdrawers = self.authorized_withdrawers.read();
        if !withdrawers.contains(approver) {
            return Err(TokenError::Unauthorized {
                reason: "Not an authorized withdrawer".to_string(),
            });
        }
        drop(withdrawers);

        // The presented public key must derive to the approver address
        if !public_key_binds_address(public_key, approver) {
            return Err(TokenError::Unauthorized {
                reason: "Public key does not derive to approver address".to_string(),
            });
        }

        // Verify the signature over the domain-separated approval preimage
        let preimage = withdrawal_approval_preimage(withdrawal_id, asset_id, amount);
        tenzro_crypto::signatures::verify(public_key, &preimage, signature).map_err(|e| {
            TokenError::Unauthorized {
                reason: format!("Invalid approval signature: {}", e),
            }
        })?;

        if amount == 0 {
            return Err(TokenError::InvalidAmount(
                "Amount must be greater than zero".to_string(),
            ));
        }

        let threshold = *self.withdrawal_threshold.read();
        let mut approvals = self.pending_approvals.write();
        let entry = approvals
            .entry(withdrawal_id.to_string())
            .or_insert_with(|| PendingWithdrawal {
                asset_id: asset_id.clone(),
                amount,
                approvers: Vec::new(),
            });

        // All approvals must bind to the same payload the first approval pinned
        if entry.asset_id != *asset_id || entry.amount != amount {
            return Err(TokenError::Unauthorized {
                reason: format!(
                    "Approval payload mismatch for withdrawal {}: pinned ({}, {}), got ({}, {})",
                    withdrawal_id,
                    entry.asset_id.as_str(),
                    entry.amount,
                    asset_id.as_str(),
                    amount
                ),
            });
        }

        // Prevent duplicate approvals from the same address
        if entry.approvers.contains(approver) {
            return Err(TokenError::InvalidAmount(format!(
                "Already approved withdrawal {}",
                withdrawal_id
            )));
        }

        entry.approvers.push(*approver);
        let approval_count = entry.approvers.len();
        info!(
            "Withdrawal {} approved by {} ({}/{})",
            withdrawal_id, approver, approval_count, threshold
        );

        Ok(approval_count >= threshold)
    }

    /// Returns the pending withdrawal record for an id, if any.
    pub fn pending_withdrawal(&self, withdrawal_id: &str) -> Option<PendingWithdrawal> {
        self.pending_approvals.read().get(withdrawal_id).cloned()
    }

    /// Clears pending approvals for a withdrawal (after execution or cancellation)
    pub fn clear_approval(&self, withdrawal_id: &str) {
        self.pending_approvals.write().remove(withdrawal_id);
    }

    /// Deposits a fee into the treasury
    ///
    /// # Arguments
    ///
    /// * `asset_id` - Asset being deposited
    /// * `amount` - Amount to deposit
    pub fn deposit_fee(&self, asset_id: &AssetId, amount: u128) -> Result<()> {
        if amount == 0 {
            return Err(TokenError::InvalidAmount(
                "Amount must be greater than zero".to_string(),
            ));
        }

        // Update balance
        let mut balances = self.balances.write();
        let current = balances.get(asset_id).copied().unwrap_or(0);
        let new_balance =
            current
                .checked_add(amount)
                .ok_or_else(|| TokenError::ArithmeticOverflow {
                    operation: "treasury deposit".to_string(),
                })?;
        balances.insert(asset_id.clone(), new_balance);
        drop(balances);

        // Update total collected
        let mut total_collected = self.total_fees_collected.write();
        let current_total = total_collected.get(asset_id).copied().unwrap_or(0);
        let new_total =
            current_total
                .checked_add(amount)
                .ok_or_else(|| TokenError::ArithmeticOverflow {
                    operation: "treasury total fees collected".to_string(),
                })?;
        total_collected.insert(asset_id.clone(), new_total);
        let distributed = self
            .total_distributed
            .read()
            .get(asset_id)
            .copied()
            .unwrap_or(0);
        drop(total_collected);

        // Persist to storage
        self.persist_asset_state(asset_id, new_balance, new_total, distributed);

        info!("Deposited {} of {} to treasury", amount, asset_id.as_str());
        Ok(())
    }

    /// Withdraws from treasury (requires authorization)
    ///
    /// When withdrawal_threshold > 1, the caller must first gather approvals
    /// via `approve_withdrawal()` before calling this method.
    ///
    /// # Arguments
    ///
    /// * `asset_id` - Asset to withdraw
    /// * `amount` - Amount to withdraw
    /// * `caller` - Address requesting withdrawal (must be authorized)
    /// * `withdrawal_id` - Unique identifier for multisig tracking (optional for threshold=1)
    pub fn withdraw(&self, asset_id: &AssetId, amount: u128, caller: &Address) -> Result<()> {
        // Check authorization
        let withdrawers = self.authorized_withdrawers.read();
        if !withdrawers.contains(caller) && caller != &self.treasury_address {
            return Err(TokenError::Unauthorized {
                reason: "Not authorized to withdraw from treasury".to_string(),
            });
        }
        drop(withdrawers);

        // Enforce multisig threshold
        let threshold = *self.withdrawal_threshold.read();
        if threshold > 1 {
            return Err(TokenError::Unauthorized {
                reason: format!(
                    "Multisig required: use approve_withdrawal() to gather {}/{} approvals first",
                    threshold, threshold
                ),
            });
        }

        if amount == 0 {
            return Err(TokenError::InvalidAmount(
                "Amount must be greater than zero".to_string(),
            ));
        }

        // Check balance
        let mut balances = self.balances.write();
        let current = balances.get(asset_id).copied().unwrap_or(0);
        if current < amount {
            return Err(TokenError::InsufficientBalance {
                required: amount,
                available: current,
            });
        }

        // Deduct balance (safe: checked above that current >= amount)
        let new_balance =
            current
                .checked_sub(amount)
                .ok_or_else(|| TokenError::ArithmeticOverflow {
                    operation: "treasury withdraw balance".to_string(),
                })?;
        balances.insert(asset_id.clone(), new_balance);
        drop(balances);

        // Update total distributed
        let mut total_distributed = self.total_distributed.write();
        let current_distributed = total_distributed.get(asset_id).copied().unwrap_or(0);
        let new_distributed = current_distributed.checked_add(amount).ok_or_else(|| {
            TokenError::ArithmeticOverflow {
                operation: "treasury total distributed".to_string(),
            }
        })?;
        total_distributed.insert(asset_id.clone(), new_distributed);
        let collected = self
            .total_fees_collected
            .read()
            .get(asset_id)
            .copied()
            .unwrap_or(0);
        drop(total_distributed);

        // Persist to storage
        self.persist_asset_state(asset_id, new_balance, collected, new_distributed);

        info!("Withdrew {} of {} from treasury", amount, asset_id.as_str());
        Ok(())
    }

    /// Executes a withdrawal after multisig approval has been reached.
    ///
    /// Verifies that the withdrawal_id has sufficient approvals, then executes.
    pub fn execute_withdrawal(
        &self,
        withdrawal_id: &str,
        asset_id: &AssetId,
        amount: u128,
    ) -> Result<()> {
        let threshold = *self.withdrawal_threshold.read();
        let approvals = self.pending_approvals.read();
        let pending = approvals.get(withdrawal_id);
        let approval_count = pending.map(|p| p.approvers.len()).unwrap_or(0);

        if approval_count < threshold {
            return Err(TokenError::Unauthorized {
                reason: format!(
                    "Insufficient approvals for withdrawal {}: have {}, need {}",
                    withdrawal_id, approval_count, threshold
                ),
            });
        }

        // Execution must match the payload the approvals were signed over
        let pending = pending.expect("approval_count >= threshold >= 1 implies record exists");
        if pending.asset_id != *asset_id || pending.amount != amount {
            return Err(TokenError::Unauthorized {
                reason: format!(
                    "Execution payload mismatch for withdrawal {}: approved ({}, {}), got ({}, {})",
                    withdrawal_id,
                    pending.asset_id.as_str(),
                    pending.amount,
                    asset_id.as_str(),
                    amount
                ),
            });
        }
        drop(approvals);

        if amount == 0 {
            return Err(TokenError::InvalidAmount(
                "Amount must be greater than zero".to_string(),
            ));
        }

        // Check balance
        let mut balances = self.balances.write();
        let current = balances.get(asset_id).copied().unwrap_or(0);
        if current < amount {
            return Err(TokenError::InsufficientBalance {
                required: amount,
                available: current,
            });
        }

        // Deduct balance (safe: checked above that current >= amount)
        let new_balance =
            current
                .checked_sub(amount)
                .ok_or_else(|| TokenError::ArithmeticOverflow {
                    operation: "treasury execute_withdrawal balance".to_string(),
                })?;
        balances.insert(asset_id.clone(), new_balance);
        drop(balances);

        // Update total distributed
        let mut total_distributed = self.total_distributed.write();
        let current_distributed = total_distributed.get(asset_id).copied().unwrap_or(0);
        let new_distributed = current_distributed.checked_add(amount).ok_or_else(|| {
            TokenError::ArithmeticOverflow {
                operation: "treasury execute_withdrawal distributed".to_string(),
            }
        })?;
        total_distributed.insert(asset_id.clone(), new_distributed);
        let collected = self
            .total_fees_collected
            .read()
            .get(asset_id)
            .copied()
            .unwrap_or(0);
        drop(total_distributed);

        // Persist to storage
        self.persist_asset_state(asset_id, new_balance, collected, new_distributed);

        // Clear pending approvals
        self.clear_approval(withdrawal_id);

        info!(
            "Executed multisig withdrawal {}: {} of {}",
            withdrawal_id,
            amount,
            asset_id.as_str()
        );
        Ok(())
    }

    /// Returns the balance for a specific asset
    pub fn balance(&self, asset_id: &AssetId) -> u128 {
        self.balances.read().get(asset_id).copied().unwrap_or(0)
    }

    /// Returns the total value across all assets (in TNZO equivalent)
    ///
    /// Note: This is a simplified version. In production, you'd use an oracle
    /// or price feed to convert different assets to TNZO equivalent.
    pub fn total_value(&self) -> u128 {
        let balances = self.balances.read();
        balances.values().sum()
    }

    /// Calculates the backing ratio (treasury_value / tnzo_supply)
    ///
    /// # Arguments
    ///
    /// * `tnzo_supply` - Current TNZO supply
    pub fn backing_ratio(&self, tnzo_supply: u128) -> f64 {
        if tnzo_supply == 0 {
            return 0.0;
        }

        let total = self.total_value();
        total as f64 / tnzo_supply as f64
    }

    /// Processes a network fee according to the distribution config
    ///
    /// # Arguments
    ///
    /// * `asset_id` - Asset of the fee
    /// * `total_fee` - Total fee amount
    ///
    /// Returns the calculated distribution
    pub fn process_network_fee(
        &self,
        asset_id: &AssetId,
        total_fee: u128,
    ) -> Result<FeeDistribution> {
        let config = self.distribution_config.read();
        config.validate()?;

        let distribution = config.calculate_distribution(total_fee);

        // Deposit treasury share
        self.deposit_fee(asset_id, distribution.treasury_amount)?;

        debug!(
            "Processed fee: treasury={}, burn={}, stakers={}",
            distribution.treasury_amount, distribution.burn_amount, distribution.staker_amount
        );

        Ok(distribution)
    }

    /// Updates the fee distribution configuration
    pub fn update_distribution_config(&self, config: FeeDistributionConfig) -> Result<()> {
        config.validate()?;
        *self.distribution_config.write() = config;
        info!("Updated fee distribution config");
        Ok(())
    }

    /// Adds an authorized withdrawer
    pub fn add_authorized_withdrawer(&self, address: Address) {
        let mut withdrawers = self.authorized_withdrawers.write();
        if !withdrawers.contains(&address) {
            withdrawers.push(address);
            info!("Added authorized withdrawer: {}", address);
        }
    }

    /// Removes an authorized withdrawer
    pub fn remove_authorized_withdrawer(&self, address: &Address) {
        let mut withdrawers = self.authorized_withdrawers.write();
        withdrawers.retain(|a| a != address);
        info!("Removed authorized withdrawer: {}", address);
    }

    /// Returns the current set of authorized withdrawers
    pub fn authorized_withdrawers(&self) -> Vec<Address> {
        self.authorized_withdrawers.read().clone()
    }

    /// Returns the current multisig withdrawal threshold
    pub fn withdrawal_threshold(&self) -> usize {
        *self.withdrawal_threshold.read()
    }

    /// Audits the treasury's supply invariant
    ///
    /// Verifies that: total_collected = current_balance + total_distributed
    ///
    /// # Arguments
    ///
    /// * `asset_id` - Asset to audit
    ///
    /// Returns `Ok(())` if the invariant holds, or an error with details if it doesn't.
    pub fn audit_supply(&self, asset_id: &AssetId) -> Result<()> {
        let balances = self.balances.read();
        let total_collected = self.total_fees_collected.read();
        let total_distributed = self.total_distributed.read();

        let current_balance = balances.get(asset_id).copied().unwrap_or(0);
        let collected = total_collected.get(asset_id).copied().unwrap_or(0);
        let distributed = total_distributed.get(asset_id).copied().unwrap_or(0);

        // Invariant: collected = balance + distributed
        // Or equivalently: balance = collected - distributed
        let expected_balance =
            collected
                .checked_sub(distributed)
                .ok_or_else(|| TokenError::TreasuryError {
                    message: format!(
                        "Distributed ({}) exceeds collected ({}) for {}",
                        distributed,
                        collected,
                        asset_id.as_str()
                    ),
                })?;

        if current_balance != expected_balance {
            return Err(TokenError::TreasuryError {
                message: format!(
                    "Supply invariant violation for {}: balance {} != expected {} (collected {} - distributed {})",
                    asset_id.as_str(),
                    current_balance,
                    expected_balance,
                    collected,
                    distributed
                ),
            });
        }

        debug!(
            "Supply audit passed for {}: balance={}, collected={}, distributed={}",
            asset_id.as_str(),
            current_balance,
            collected,
            distributed
        );
        Ok(())
    }

    /// Audits all assets in the treasury
    ///
    /// Returns the list of assets that failed the audit.
    pub fn audit_all_supplies(&self) -> Vec<(AssetId, String)> {
        let balances = self.balances.read();
        let mut failures = Vec::new();

        for asset_id in balances.keys() {
            if let Err(e) = self.audit_supply(asset_id) {
                failures.push((asset_id.clone(), e.to_string()));
            }
        }

        if !failures.is_empty() {
            warn!("Treasury audit found {} violations", failures.len());
        } else {
            info!("Treasury audit passed for all {} assets", balances.len());
        }

        failures
    }

    /// Returns treasury statistics
    pub fn stats(&self) -> TreasuryStats {
        let balances = self.balances.read().clone();
        let total_collected = self.total_fees_collected.read().clone();
        let total_distributed = self.total_distributed.read().clone();

        TreasuryStats {
            treasury_address: self.treasury_address,
            balances,
            total_fees_collected: total_collected,
            total_distributed,
            total_value: self.total_value(),
        }
    }
}

/// Treasury statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TreasuryStats {
    /// Treasury address
    pub treasury_address: Address,
    /// Current balances per asset
    pub balances: HashMap<AssetId, u128>,
    /// Total fees collected per asset
    pub total_fees_collected: HashMap<AssetId, u128>,
    /// Total distributed per asset
    pub total_distributed: HashMap<AssetId, u128>,
    /// Total value (in TNZO equivalent)
    pub total_value: u128,
}

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

    #[test]
    fn test_fee_distribution_config() {
        let config = FeeDistributionConfig::default();
        assert!(config.validate().is_ok());

        let distribution = config.calculate_distribution(10000);
        assert_eq!(distribution.treasury_amount, 4000);
        assert_eq!(distribution.burn_amount, 3000);
        assert_eq!(distribution.staker_amount, 3000);
    }

    #[test]
    fn test_treasury_deposit() {
        let treasury = NetworkTreasury::new(Address::new([1u8; 32]));
        let asset_id = AssetId::tnzo();

        treasury.deposit_fee(&asset_id, 1000).unwrap();
        assert_eq!(treasury.balance(&asset_id), 1000);
    }

    #[test]
    fn test_treasury_withdraw() {
        let treasury_addr = Address::new([1u8; 32]);
        let treasury = NetworkTreasury::new(treasury_addr);
        let asset_id = AssetId::tnzo();

        treasury.deposit_fee(&asset_id, 1000).unwrap();
        treasury.withdraw(&asset_id, 500, &treasury_addr).unwrap();
        assert_eq!(treasury.balance(&asset_id), 500);
    }

    /// Generates an Ed25519 keypair and the canonical 32-byte withdrawer
    /// address it derives to (20-byte crypto address, zero-padded tail).
    fn test_signer() -> (tenzro_crypto::KeyPair, Address) {
        let keypair = tenzro_crypto::KeyPair::generate(tenzro_crypto::KeyType::Ed25519).unwrap();
        let derived = keypair.public_key().to_address();
        let mut addr_bytes = [0u8; 32];
        addr_bytes[..20].copy_from_slice(derived.as_bytes());
        (keypair, Address::new(addr_bytes))
    }

    fn sign_approval(
        keypair: &tenzro_crypto::KeyPair,
        withdrawal_id: &str,
        asset_id: &AssetId,
        amount: u128,
    ) -> Signature {
        use tenzro_crypto::Signer as _;
        let kp = tenzro_crypto::KeyPair::from_secret_key(keypair.secret_key().clone()).unwrap();
        let signer = tenzro_crypto::signatures::Ed25519SignerImpl::new(kp).unwrap();
        signer
            .sign(&withdrawal_approval_preimage(
                withdrawal_id,
                asset_id,
                amount,
            ))
            .unwrap()
    }

    #[test]
    fn test_multisig_withdrawal() {
        let treasury_addr = Address::new([0u8; 32]);
        let treasury = NetworkTreasury::new(treasury_addr);
        let asset_id = AssetId::tnzo();

        let (kp1, signer1) = test_signer();
        let (kp2, signer2) = test_signer();
        let (_kp3, signer3) = test_signer();

        // Add authorized withdrawers
        treasury.add_authorized_withdrawer(signer1);
        treasury.add_authorized_withdrawer(signer2);
        treasury.add_authorized_withdrawer(signer3);

        // Set 2-of-3 multisig
        treasury.set_withdrawal_threshold(2).unwrap();

        // Deposit funds
        treasury.deposit_fee(&asset_id, 1000).unwrap();

        // Single-sig withdraw should be blocked
        assert!(treasury.withdraw(&asset_id, 500, &signer1).is_err());

        // First approval — not yet reached threshold
        let sig1 = sign_approval(&kp1, "w1", &asset_id, 500);
        let reached = treasury
            .approve_withdrawal("w1", &asset_id, 500, &signer1, kp1.public_key(), &sig1)
            .unwrap();
        assert!(!reached);

        // Execute before threshold — should fail
        assert!(treasury.execute_withdrawal("w1", &asset_id, 500).is_err());

        // Second approval — reached threshold
        let sig2 = sign_approval(&kp2, "w1", &asset_id, 500);
        let reached = treasury
            .approve_withdrawal("w1", &asset_id, 500, &signer2, kp2.public_key(), &sig2)
            .unwrap();
        assert!(reached);

        // Execute with a different amount than approved — must fail
        assert!(treasury.execute_withdrawal("w1", &asset_id, 999).is_err());

        // Execute now — should succeed
        treasury.execute_withdrawal("w1", &asset_id, 500).unwrap();
        assert_eq!(treasury.balance(&asset_id), 500);
    }

    #[test]
    fn test_no_duplicate_approvals() {
        let treasury_addr = Address::new([0u8; 32]);
        let treasury = NetworkTreasury::new(treasury_addr);
        let asset_id = AssetId::tnzo();

        let (kp, signer) = test_signer();
        treasury.add_authorized_withdrawer(signer);
        treasury.set_withdrawal_threshold(1).unwrap();

        let sig = sign_approval(&kp, "w1", &asset_id, 100);
        treasury
            .approve_withdrawal("w1", &asset_id, 100, &signer, kp.public_key(), &sig)
            .unwrap();
        // Duplicate approval should fail
        assert!(
            treasury
                .approve_withdrawal("w1", &asset_id, 100, &signer, kp.public_key(), &sig)
                .is_err()
        );
    }

    #[test]
    fn test_approval_rejects_bad_signature_and_wrong_key() {
        let treasury = NetworkTreasury::new(Address::new([0u8; 32]));
        let asset_id = AssetId::tnzo();

        let (kp, signer) = test_signer();
        let (other_kp, _other_addr) = test_signer();
        treasury.add_authorized_withdrawer(signer);
        treasury.set_withdrawal_threshold(1).unwrap();

        // Signature over a different payload must be rejected
        let sig_wrong_payload = sign_approval(&kp, "w1", &asset_id, 999);
        assert!(
            treasury
                .approve_withdrawal(
                    "w1",
                    &asset_id,
                    100,
                    &signer,
                    kp.public_key(),
                    &sig_wrong_payload
                )
                .is_err()
        );

        // A key that does not derive to the approver address must be rejected,
        // even with a valid signature over the right payload
        let sig_other = sign_approval(&other_kp, "w1", &asset_id, 100);
        assert!(
            treasury
                .approve_withdrawal(
                    "w1",
                    &asset_id,
                    100,
                    &signer,
                    other_kp.public_key(),
                    &sig_other
                )
                .is_err()
        );

        // Payload pinning: valid first approval, then a conflicting payload
        let sig = sign_approval(&kp, "w1", &asset_id, 100);
        treasury
            .approve_withdrawal("w1", &asset_id, 100, &signer, kp.public_key(), &sig)
            .unwrap();
        let pinned = treasury.pending_withdrawal("w1").unwrap();
        assert_eq!(pinned.amount, 100);
        assert_eq!(pinned.approvers.len(), 1);
    }

    #[test]
    fn test_supply_audit_passes() {
        let treasury = NetworkTreasury::new(Address::new([1u8; 32]));
        let asset_id = AssetId::tnzo();

        // Deposit 1000
        treasury.deposit_fee(&asset_id, 1000).unwrap();

        // Audit should pass
        assert!(treasury.audit_supply(&asset_id).is_ok());

        // Add authorized withdrawer and set threshold to 1
        treasury.add_authorized_withdrawer(treasury.treasury_address);

        // Withdraw 400 (single-sig)
        treasury
            .withdraw(&asset_id, 400, &treasury.treasury_address)
            .unwrap();

        // Audit should still pass: balance=600, collected=1000, distributed=400
        assert!(treasury.audit_supply(&asset_id).is_ok());
    }

    #[test]
    fn test_supply_audit_detects_violation() {
        let treasury = NetworkTreasury::new(Address::new([1u8; 32]));
        let asset_id = AssetId::tnzo();

        // Deposit 1000
        treasury.deposit_fee(&asset_id, 1000).unwrap();

        // Manually corrupt the distributed counter (simulate a bug)
        {
            let mut distributed = treasury.total_distributed.write();
            distributed.insert(asset_id.clone(), 1500); // Exceeds collected
        }

        // Audit should fail
        assert!(treasury.audit_supply(&asset_id).is_err());
    }

    #[test]
    fn test_audit_all_supplies() {
        let treasury = NetworkTreasury::new(Address::new([1u8; 32]));
        let asset1 = AssetId::tnzo();
        let asset2 = AssetId::usd_stablecoin();

        // Deposit to both assets
        treasury.deposit_fee(&asset1, 1000).unwrap();
        treasury.deposit_fee(&asset2, 2000).unwrap();

        // All audits should pass
        let failures = treasury.audit_all_supplies();
        assert_eq!(failures.len(), 0);

        // Corrupt one asset
        {
            let mut distributed = treasury.total_distributed.write();
            distributed.insert(asset1.clone(), 1500);
        }

        // Should detect exactly one failure
        let failures = treasury.audit_all_supplies();
        assert_eq!(failures.len(), 1);
        assert_eq!(failures[0].0, asset1);
    }
}