contract.rs

use cosmwasm_std::{
    entry_point, to_binary, Addr, BankMsg, Binary, Coin,
    Deps, DepsMut, Env, MessageInfo, Response, StdResult, Uint128, StdError, WasmMsg,
    Order,
};
use cw2::set_contract_version;
use cw20::{Cw20ReceiveMsg, TokenInfoResponse};
use cw_storage_plus::{Bound, Item};

use crate::error::ContractError;
use crate::msg::{
    ExecuteMsg, InstantiateMsg, MigrateMsg, QueryMsg, SwapTokenForBwick, PoolResponse,
    AllPoolsResponse, SimulateSwapResponse, ConfigResponse, AugmentedFeeStatusResponse,
    LpBalanceResponse, SimulateAddLiquidityResponse,
};
use crate::state::{Config, Pool, AugmentedFeeConfig, CONFIG, POOLS, OPEN_LP};

const CONTRACT_NAME: &str = "crates.io:bwick-amm";
const CONTRACT_VERSION: &str = env!("CARGO_PKG_VERSION");
/// Maximum total fee (base + augmented) in basis points: 1000 bps = 10%
const MAX_TOTAL_FEE_BPS: u16 = 1000;

#[cfg_attr(not(feature = "library"), entry_point)]
pub fn instantiate(
    deps: DepsMut,
    _env: Env,
    _info: MessageInfo,
    msg: InstantiateMsg,
) -> Result<Response, ContractError> {
    // Validate authorized creators addresses
    let mut authorized_creators = Vec::new();
    for creator_str in msg.authorized_creators.iter() {
        let validated = deps.api.addr_validate(creator_str)?;
        authorized_creators.push(validated);
    }

    // Validate swap_fee_bps against fee cap
    // Base fee alone cannot exceed the total fee cap since augmented fees add on top
    if msg.swap_fee_bps > MAX_TOTAL_FEE_BPS {
        return Err(ContractError::TotalFeeCapExceeded {
            swap_fee_bps: msg.swap_fee_bps,
            augmented_fee_bps: 0,
            total: msg.swap_fee_bps,
        });
    }

    // Save config
    let config = Config {
        authorized_creators: authorized_creators.clone(),
        swap_fee_bps: msg.swap_fee_bps,
        max_wallet_bps: msg.max_wallet_bps,
    };
    CONFIG.save(deps.storage, &config)?;

    // Set contract version
    set_contract_version(deps.storage, CONTRACT_NAME, CONTRACT_VERSION)?;

    // Create comma-separated list for attribute
    let creators_list = authorized_creators
        .iter()
        .map(|a| a.to_string())
        .collect::<Vec<_>>()
        .join(",");

    Ok(Response::new()
        .add_attribute("action", "instantiate")
        .add_attribute("authorized_creators", creators_list)
        .add_attribute("swap_fee_bps", msg.swap_fee_bps.to_string()))
}

// Legacy config format for migration compatibility
#[cosmwasm_schema::cw_serde]
struct LegacyConfig {
    pub bonding_curve_contract: Addr,
    pub swap_fee_bps: u16,
}

#[cfg_attr(not(feature = "library"), entry_point)]
pub fn migrate(
    deps: DepsMut,
    _env: Env,
    msg: MigrateMsg,
) -> Result<Response, ContractError> {
    // Try loading as new format first, fall back to legacy
    let mut config = match CONFIG.load(deps.storage) {
        Ok(c) => c,
        Err(_) => {
            // Try legacy format
            let legacy_item: Item<LegacyConfig> = Item::new("config");
            let legacy = legacy_item.load(deps.storage)?;
            let new_config = Config {
                authorized_creators: vec![legacy.bonding_curve_contract],
                swap_fee_bps: legacy.swap_fee_bps,
                max_wallet_bps: 300,
            };
            CONFIG.save(deps.storage, &new_config)?;
            new_config
        }
    };

    // Apply add/remove operations from MigrateMsg
    if let Some(add_addr) = msg.add_authorized_creator {
        let validated = deps.api.addr_validate(&add_addr)?;
        if !config.authorized_creators.contains(&validated) {
            config.authorized_creators.push(validated);
            CONFIG.save(deps.storage, &config)?;
        }
    }

    if let Some(remove_addr) = msg.remove_authorized_creator {
        let validated = deps.api.addr_validate(&remove_addr)?;
        config.authorized_creators.retain(|a| a != &validated);
        CONFIG.save(deps.storage, &config)?;
    }

    if let Some(max_wallet) = msg.max_wallet_bps {
        config.max_wallet_bps = max_wallet;
        CONFIG.save(deps.storage, &config)?;
    }

    // v2: backfill locked_lp_supply for existing pools that were created before this field
    // existed. Treat all pre-existing LP as locked (since open LPs weren't a thing yet).
    let pool_keys: Vec<Addr> = POOLS
        .keys(deps.storage, None, None, Order::Ascending)
        .filter_map(|r| r.ok())
        .collect();
    for k in pool_keys {
        let mut pool = POOLS.load(deps.storage, &k)?;
        if pool.locked_lp_supply == 0 && pool.lp_total_supply > 0 {
            pool.locked_lp_supply = pool.lp_total_supply;
            POOLS.save(deps.storage, &k, &pool)?;
        }
    }

    set_contract_version(deps.storage, CONTRACT_NAME, CONTRACT_VERSION)?;
    Ok(Response::new().add_attribute("action", "migrate"))
}

#[cfg_attr(not(feature = "library"), entry_point)]
pub fn execute(
    deps: DepsMut,
    env: Env,
    info: MessageInfo,
    msg: ExecuteMsg,
) -> Result<Response, ContractError> {
    match msg {
        ExecuteMsg::CreatePool {
            token_address,
            bwick_amount,
            token_amount,
            augmented_fee_bps,
            lp_target_ubwick,
            locked_lp_unlock_at,
            locked_lp_recipient,
        } => execute_create_pool(
            deps,
            env,
            info,
            token_address,
            bwick_amount,
            token_amount,
            augmented_fee_bps,
            lp_target_ubwick,
            locked_lp_unlock_at,
            locked_lp_recipient,
        ),
        ExecuteMsg::Swap {
            token_address,
            offer_bwick,
            min_output,
        } => execute_swap(deps, env, info, token_address, offer_bwick, min_output),
        ExecuteMsg::AddLiquidity {
            token_address,
            max_token_amount,
        } => execute_add_liquidity(deps, env, info, token_address, max_token_amount),
        ExecuteMsg::RemoveLiquidity {
            token_address,
            lp_amount,
        } => execute_remove_liquidity(deps, env, info, token_address, lp_amount),
        ExecuteMsg::WithdrawLockedLp {
            token_address,
            recipient,
        } => execute_withdraw_locked_lp(deps, env, info, token_address, recipient),
        ExecuteMsg::Receive(cw20_msg) => execute_receive(deps, env, info, cw20_msg),
    }
}

fn execute_create_pool(
    deps: DepsMut,
    env: Env,
    info: MessageInfo,
    token_address: String,
    bwick_amount: Uint128,
    token_amount: Uint128,
    augmented_fee_bps: Option<u16>,
    lp_target_ubwick: Option<String>,
    locked_lp_unlock_at: Option<u64>,
    locked_lp_recipient: Option<String>,
) -> Result<Response, ContractError> {
    // Authorization check: only authorized creators can create pools
    let config = CONFIG.load(deps.storage)?;
    if !config.authorized_creators.contains(&info.sender) {
        return Err(ContractError::UnauthorizedPoolCreation {});
    }

    // Validate amounts
    if bwick_amount.is_zero() || token_amount.is_zero() {
        return Err(ContractError::ZeroAmount {});
    }

    // Validate token address
    let token_addr = deps.api.addr_validate(&token_address)?;

    // Check if token is valid CW20
    validate_cw20_token(deps.as_ref(), &token_addr)?;

    // Check pool doesn't already exist
    if POOLS.may_load(deps.storage, &token_addr)?.is_some() {
        return Err(ContractError::PoolAlreadyExists {
            token: token_address,
        });
    }

    // Verify BWICK funds sent with message
    let bwick_sent = extract_bwick_from_funds(&info.funds)?;
    if bwick_sent != bwick_amount {
        return Err(ContractError::Std(StdError::generic_err(format!(
            "BWICK amount mismatch: expected {}, got {}",
            bwick_amount, bwick_sent
        ))));
    }

    // Calculate initial LP supply: sqrt(bwick_amount * token_amount)
    let lp_total_supply = integer_sqrt(bwick_amount.u128() * token_amount.u128());

    // Process augmented fee configuration
    let augmented_fee = match (augmented_fee_bps, lp_target_ubwick) {
        (Some(fee_bps), Some(target_str)) => {
            // Both provided - validate and create config
            if fee_bps == 0 || fee_bps > 500 {
                return Err(ContractError::AugmentedFeeInvalid {});
            }
            // Validate total fee cap: base + augmented must not exceed 1000 bps (10%)
            let total_fee = config.swap_fee_bps + fee_bps;
            if total_fee > MAX_TOTAL_FEE_BPS {
                return Err(ContractError::TotalFeeCapExceeded {
                    swap_fee_bps: config.swap_fee_bps,
                    augmented_fee_bps: fee_bps,
                    total: total_fee,
                });
            }
            let target = target_str.parse::<u128>().map_err(|_| {
                ContractError::Std(StdError::generic_err("Invalid lp_target_ubwick format"))
            })?;
            if target == 0 {
                return Err(ContractError::AugmentedFeeInvalid {});
            }
            Some(AugmentedFeeConfig {
                augmented_fee_bps: fee_bps,
                lp_target_ubwick: target,
                active: true,
            })
        }
        (None, None) => None, // Neither provided - normal pool
        _ => return Err(ContractError::AugmentedFeeIncomplete {}), // Only one provided
    };

    // v3: validate lock config. If unlock_at is set, recipient is required.
    // If recipient is set without unlock_at, we silently ignore it (treated
    // as permanent lock) — that keeps the existing callers backwards-compat.
    let (locked_lp_unlock_at_v, locked_lp_recipient_v) = match locked_lp_unlock_at {
        Some(t) => {
            let recipient_str = locked_lp_recipient
                .ok_or(ContractError::LockedLpRecipientRequired {})?;
            let recipient_addr = deps.api.addr_validate(&recipient_str)?;
            (Some(t), Some(recipient_addr))
        }
        None => (None, None),
    };

    // Create pool. The initial seed is fully locked: locked_lp_supply == lp_total_supply.
    // Open LPs added later (via AddLiquidity) are tracked separately in OPEN_LP.
    let pool = Pool {
        token_address: token_addr.clone(),
        bwick_reserve: bwick_amount.u128(),
        token_reserve: token_amount.u128(),
        lp_token_address: env.contract.address.clone(),
        lp_total_supply,
        locked_lp_supply: lp_total_supply,
        locked_lp_unlock_at: locked_lp_unlock_at_v,
        locked_lp_recipient: locked_lp_recipient_v,
        augmented_fee,
    };

    // Save pool
    POOLS.save(deps.storage, &token_addr, &pool)?;

    let mut response = Response::new()
        .add_attribute("action", "create_pool")
        .add_attribute("token_address", token_address)
        .add_attribute("bwick_reserve", bwick_amount)
        .add_attribute("token_reserve", token_amount)
        .add_attribute("lp_total_supply", lp_total_supply.to_string());

    // Add augmented fee attributes if configured
    if let Some(ref aug_fee) = pool.augmented_fee {
        response = response
            .add_attribute("augmented_fee_bps", aug_fee.augmented_fee_bps.to_string())
            .add_attribute("lp_target_ubwick", aug_fee.lp_target_ubwick.to_string());
    }

    // v3: surface lock metadata so indexers don't have to query the pool back.
    if let (Some(t), Some(r)) = (pool.locked_lp_unlock_at, pool.locked_lp_recipient.as_ref()) {
        response = response
            .add_attribute("locked_lp_unlock_at", t.to_string())
            .add_attribute("locked_lp_recipient", r.to_string());
    } else {
        response = response.add_attribute("locked_lp_unlock_at", "forever");
    }

    Ok(response)
}

/// Check that buyer's holding after purchase won't exceed max_wallet_bps of token's total supply.
/// Queries the CW20 contract for total_supply and the buyer's current balance.
/// Skips the check if max_wallet_bps is 0 (disabled).
fn check_max_wallet_amm(
    deps: &DepsMut,
    token_address: &str,
    buyer_address: &str,
    tokens_out: u128,
    max_wallet_bps: u16,
) -> Result<(), ContractError> {
    if max_wallet_bps == 0 {
        return Ok(());
    }

    // Query CW20 total supply
    let token_info: TokenInfoResponse = deps.querier.query_wasm_smart(
        token_address,
        &cw20::Cw20QueryMsg::TokenInfo {},
    )?;
    let total_supply = token_info.total_supply.u128();
    let limit = total_supply * (max_wallet_bps as u128) / 10000;

    // Query buyer's current CW20 balance
    let balance_resp: cw20::BalanceResponse = deps.querier.query_wasm_smart(
        token_address,
        &cw20::Cw20QueryMsg::Balance {
            address: buyer_address.to_string(),
        },
    )?;
    let current = balance_resp.balance.u128();
    let would_become = current + tokens_out;

    if would_become > limit {
        return Err(ContractError::MaxWalletExceeded {
            current,
            buying: tokens_out,
            would_become,
            limit,
        });
    }

    Ok(())
}

// ── v2: open LP add/remove ────────────────────────────────────────────────────

fn execute_add_liquidity(
    deps: DepsMut,
    _env: Env,
    info: MessageInfo,
    token_address: String,
    max_token_amount: Uint128,
) -> Result<Response, ContractError> {
    let token_addr = deps.api.addr_validate(&token_address)?;
    let mut pool = POOLS.load(deps.storage, &token_addr).map_err(|_| {
        ContractError::PoolNotFound { token: token_address.clone() }
    })?;

    if pool.lp_total_supply == 0 || pool.bwick_reserve == 0 || pool.token_reserve == 0 {
        return Err(ContractError::Std(StdError::generic_err(
            "pool not initialized; cannot add liquidity",
        )));
    }

    // BWICK side: from attached funds (convert to u128 for arithmetic).
    let bwick_added: u128 = extract_bwick_from_funds(&info.funds)?.u128();
    if bwick_added == 0 {
        return Err(ContractError::ZeroAmount {});
    }

    // Required token amount to maintain pool ratio (round up, conservative for the pool).
    // token_required = ceil(bwick_added * token_reserve / bwick_reserve)
    let token_required_num = bwick_added
        .checked_mul(pool.token_reserve)
        .ok_or_else(|| ContractError::Std(StdError::generic_err("overflow in token requirement")))?;
    let token_required = (token_required_num + pool.bwick_reserve - 1) / pool.bwick_reserve;
    if token_required > max_token_amount.u128() {
        return Err(ContractError::Std(StdError::generic_err(format!(
            "slippage: token_required={} exceeds max_token_amount={}",
            token_required, max_token_amount.u128()
        ))));
    }

    // LP shares minted: proportional to BWICK side (round down so existing LPs aren't diluted by rounding).
    let shares_minted = pool.lp_total_supply
        .checked_mul(bwick_added)
        .ok_or_else(|| ContractError::Std(StdError::generic_err("overflow in lp shares")))?
        / pool.bwick_reserve;
    if shares_minted == 0 {
        return Err(ContractError::Std(StdError::generic_err("deposit too small to mint shares")));
    }

    // Pull tokens from caller via TransferFrom (caller must have set allowance).
    let pull_tokens = WasmMsg::Execute {
        contract_addr: token_address.clone(),
        msg: to_binary(&cw20::Cw20ExecuteMsg::TransferFrom {
            owner: info.sender.to_string(),
            recipient: _env.contract.address.to_string(),
            amount: Uint128::from(token_required),
        })?,
        funds: vec![],
    };

    // Update reserves and supply.
    pool.bwick_reserve = pool.bwick_reserve.checked_add(bwick_added)
        .ok_or_else(|| ContractError::Std(StdError::generic_err("bwick_reserve overflow")))?;
    pool.token_reserve = pool.token_reserve.checked_add(token_required)
        .ok_or_else(|| ContractError::Std(StdError::generic_err("token_reserve overflow")))?;
    pool.lp_total_supply = pool.lp_total_supply.checked_add(shares_minted)
        .ok_or_else(|| ContractError::Std(StdError::generic_err("lp_total_supply overflow")))?;
    POOLS.save(deps.storage, &token_addr, &pool)?;

    // Credit shares to caller's open LP balance.
    let prev = OPEN_LP.may_load(deps.storage, (&token_addr, &info.sender))?.unwrap_or(0);
    OPEN_LP.save(deps.storage, (&token_addr, &info.sender), &(prev + shares_minted))?;

    Ok(Response::new()
        .add_message(pull_tokens)
        .add_attribute("action", "add_liquidity")
        .add_attribute("token_address", token_address)
        .add_attribute("provider", info.sender)
        .add_attribute("bwick_added", bwick_added.to_string())
        .add_attribute("token_added", token_required.to_string())
        .add_attribute("lp_shares_minted", shares_minted.to_string())
        .add_attribute("new_lp_total_supply", pool.lp_total_supply.to_string()))
}

fn execute_remove_liquidity(
    deps: DepsMut,
    _env: Env,
    info: MessageInfo,
    token_address: String,
    lp_amount: Uint128,
) -> Result<Response, ContractError> {
    let token_addr = deps.api.addr_validate(&token_address)?;
    let mut pool = POOLS.load(deps.storage, &token_addr).map_err(|_| {
        ContractError::PoolNotFound { token: token_address.clone() }
    })?;

    let lp = lp_amount.u128();
    if lp == 0 {
        return Err(ContractError::ZeroAmount {});
    }

    // Caller must own at least `lp` open LP shares.
    let holder_shares = OPEN_LP
        .may_load(deps.storage, (&token_addr, &info.sender))?
        .unwrap_or(0);
    if lp > holder_shares {
        return Err(ContractError::Std(StdError::generic_err(format!(
            "insufficient lp balance: have {}, need {}",
            holder_shares, lp
        ))));
    }

    // Compute proportional withdrawal (round down so the pool keeps the dust).
    if pool.lp_total_supply == 0 {
        return Err(ContractError::Std(StdError::generic_err("pool empty")));
    }
    let bwick_out = pool.bwick_reserve
        .checked_mul(lp)
        .ok_or_else(|| ContractError::Std(StdError::generic_err("overflow in bwick_out")))?
        / pool.lp_total_supply;
    let token_out = pool.token_reserve
        .checked_mul(lp)
        .ok_or_else(|| ContractError::Std(StdError::generic_err("overflow in token_out")))?
        / pool.lp_total_supply;
    if bwick_out == 0 || token_out == 0 {
        return Err(ContractError::Std(StdError::generic_err(
            "lp_amount too small to withdraw any reserves",
        )));
    }

    // Update reserves and supply.
    pool.bwick_reserve = pool.bwick_reserve.checked_sub(bwick_out)
        .ok_or_else(|| ContractError::Std(StdError::generic_err("bwick_reserve underflow")))?;
    pool.token_reserve = pool.token_reserve.checked_sub(token_out)
        .ok_or_else(|| ContractError::Std(StdError::generic_err("token_reserve underflow")))?;
    pool.lp_total_supply = pool.lp_total_supply.checked_sub(lp)
        .ok_or_else(|| ContractError::Std(StdError::generic_err("lp_total_supply underflow")))?;
    POOLS.save(deps.storage, &token_addr, &pool)?;

    // Decrement caller's open LP balance.
    let new_balance = holder_shares - lp;
    if new_balance == 0 {
        OPEN_LP.remove(deps.storage, (&token_addr, &info.sender));
    } else {
        OPEN_LP.save(deps.storage, (&token_addr, &info.sender), &new_balance)?;
    }

    // Send BWICK + tokens back to caller.
    let send_bwick = BankMsg::Send {
        to_address: info.sender.to_string(),
        amount: vec![Coin { denom: "ubwick".to_string(), amount: Uint128::from(bwick_out) }],
    };
    let send_tokens = WasmMsg::Execute {
        contract_addr: token_address.clone(),
        msg: to_binary(&cw20::Cw20ExecuteMsg::Transfer {
            recipient: info.sender.to_string(),
            amount: Uint128::from(token_out),
        })?,
        funds: vec![],
    };

    Ok(Response::new()
        .add_message(send_bwick)
        .add_message(send_tokens)
        .add_attribute("action", "remove_liquidity")
        .add_attribute("token_address", token_address)
        .add_attribute("provider", info.sender)
        .add_attribute("lp_burned", lp.to_string())
        .add_attribute("bwick_returned", bwick_out.to_string())
        .add_attribute("tokens_returned", token_out.to_string())
        .add_attribute("new_lp_total_supply", pool.lp_total_supply.to_string()))
}

// v3: withdraw the locked LP seed once the configured timestamp has passed.
// This is the only path that can decrement `locked_lp_supply`. Open LP
// holders keep their proportional claim because we decrement
// `lp_total_supply` by the same amount we burn from `locked_lp_supply`.
fn execute_withdraw_locked_lp(
    deps: DepsMut,
    env: Env,
    info: MessageInfo,
    token_address: String,
    recipient: Option<String>,
) -> Result<Response, ContractError> {
    let token_addr = deps.api.addr_validate(&token_address)?;
    let mut pool = POOLS.load(deps.storage, &token_addr).map_err(|_| {
        ContractError::PoolNotFound { token: token_address.clone() }
    })?;

    let unlock_at = pool.locked_lp_unlock_at.ok_or(ContractError::LockedLpPermanent {})?;
    let configured_recipient = pool
        .locked_lp_recipient
        .clone()
        .ok_or(ContractError::LockedLpPermanent {})?;

    let now = env.block.time.seconds();
    if now < unlock_at {
        return Err(ContractError::LockedLpStillLocked { unlock_at, now });
    }
    if info.sender != configured_recipient {
        return Err(ContractError::LockedLpUnauthorized {});
    }
    if pool.locked_lp_supply == 0 {
        return Err(ContractError::LockedLpAlreadyWithdrawn {});
    }
    if pool.lp_total_supply == 0 {
        return Err(ContractError::Std(StdError::generic_err("pool empty")));
    }

    let lp = pool.locked_lp_supply;
    let bwick_out = pool
        .bwick_reserve
        .checked_mul(lp)
        .ok_or_else(|| ContractError::Std(StdError::generic_err("overflow in bwick_out")))?
        / pool.lp_total_supply;
    let token_out = pool
        .token_reserve
        .checked_mul(lp)
        .ok_or_else(|| ContractError::Std(StdError::generic_err("overflow in token_out")))?
        / pool.lp_total_supply;

    pool.bwick_reserve = pool
        .bwick_reserve
        .checked_sub(bwick_out)
        .ok_or_else(|| ContractError::Std(StdError::generic_err("bwick_reserve underflow")))?;
    pool.token_reserve = pool
        .token_reserve
        .checked_sub(token_out)
        .ok_or_else(|| ContractError::Std(StdError::generic_err("token_reserve underflow")))?;
    pool.lp_total_supply = pool
        .lp_total_supply
        .checked_sub(lp)
        .ok_or_else(|| ContractError::Std(StdError::generic_err("lp_total_supply underflow")))?;
    pool.locked_lp_supply = 0;
    POOLS.save(deps.storage, &token_addr, &pool)?;

    let to_addr = match recipient {
        Some(r) => deps.api.addr_validate(&r)?,
        None => configured_recipient,
    };

    let send_bwick = BankMsg::Send {
        to_address: to_addr.to_string(),
        amount: vec![Coin { denom: "ubwick".to_string(), amount: Uint128::from(bwick_out) }],
    };
    let send_tokens = WasmMsg::Execute {
        contract_addr: token_address.clone(),
        msg: to_binary(&cw20::Cw20ExecuteMsg::Transfer {
            recipient: to_addr.to_string(),
            amount: Uint128::from(token_out),
        })?,
        funds: vec![],
    };

    Ok(Response::new()
        .add_message(send_bwick)
        .add_message(send_tokens)
        .add_attribute("action", "withdraw_locked_lp")
        .add_attribute("token_address", token_address)
        .add_attribute("recipient", to_addr)
        .add_attribute("lp_burned", lp.to_string())
        .add_attribute("bwick_returned", bwick_out.to_string())
        .add_attribute("tokens_returned", token_out.to_string())
        .add_attribute("new_lp_total_supply", pool.lp_total_supply.to_string()))
}

fn execute_swap(
    deps: DepsMut,
    _env: Env,
    info: MessageInfo,
    token_address: String,
    offer_bwick: bool,
    min_output: Uint128,
) -> Result<Response, ContractError> {
    let config = CONFIG.load(deps.storage)?;
    let token_addr = deps.api.addr_validate(&token_address)?;

    if !offer_bwick {
        return Err(ContractError::Std(StdError::generic_err(
            "For Token->BWICK swaps, use CW20 Send to this contract with SwapTokenForBwick message",
        )));
    }

    // BWICK -> Token swap
    let mut pool = POOLS.load(deps.storage, &token_addr)?;

    // Extract BWICK input amount
    let input_amount = extract_bwick_from_funds(&info.funds)?;
    if input_amount.is_zero() {
        return Err(ContractError::ZeroAmount {});
    }

    // Calculate total fee (base + augmented if active)
    let base_fee_bps = config.swap_fee_bps;
    let total_fee_bps = if let Some(ref aug_fee) = pool.augmented_fee {
        if aug_fee.active {
            base_fee_bps + aug_fee.augmented_fee_bps
        } else {
            base_fee_bps
        }
    } else {
        base_fee_bps
    };

    // Calculate output using constant product
    let (output_amount, fee_amount) = calculate_swap_output(
        pool.bwick_reserve,
        pool.token_reserve,
        input_amount.u128(),
        total_fee_bps,
    )?;

    // Check slippage protection
    if Uint128::new(output_amount) < min_output {
        return Err(ContractError::SlippageExceeded {
            expected: min_output.u128(),
            actual: output_amount,
        });
    }

    // Check max wallet holding limit (BWICK->Token direction only = buying tokens)
    check_max_wallet_amm(
        &deps,
        &token_address,
        info.sender.as_str(),
        output_amount,
        config.max_wallet_bps,
    )?;

    // Update pool reserves
    // Fee stays in pool (auto-compound) - input includes fee
    pool.bwick_reserve += input_amount.u128();
    pool.token_reserve -= output_amount;

    // Check if augmented fee should auto-disable
    let mut response = Response::new();
    if let Some(ref mut aug_fee) = pool.augmented_fee {
        if aug_fee.active {
            let pool_value = pool.bwick_reserve * 2;
            if pool_value >= aug_fee.lp_target_ubwick {
                aug_fee.active = false;
                response = response
                    .add_attribute("augmented_fee_disabled", "true")
                    .add_attribute("pool_value", pool_value.to_string())
                    .add_attribute("lp_target", aug_fee.lp_target_ubwick.to_string());
            }
        }
    }

    // Save updated pool
    POOLS.save(deps.storage, &token_addr, &pool)?;

    // Transfer CW20 tokens to sender
    let transfer_msg = WasmMsg::Execute {
        contract_addr: token_address.clone(),
        msg: to_binary(&cw20::Cw20ExecuteMsg::Transfer {
            recipient: info.sender.to_string(),
            amount: Uint128::new(output_amount),
        })?,
        funds: vec![],
    };

    // Calculate augmented fee portion for attribute
    let augmented_fee_amount = if total_fee_bps > base_fee_bps {
        (input_amount.u128() * (total_fee_bps - base_fee_bps) as u128) / 10000
    } else {
        0
    };

    Ok(response
        .add_message(transfer_msg)
        .add_attribute("action", "swap")
        .add_attribute("direction", "bwick_to_token")
        .add_attribute("token_address", token_address)
        .add_attribute("input_amount", input_amount)
        .add_attribute("output_amount", output_amount.to_string())
        .add_attribute("fee_amount", fee_amount.to_string())
        .add_attribute("augmented_fee_amount", augmented_fee_amount.to_string()))
}

fn execute_receive(
    deps: DepsMut,
    _env: Env,
    info: MessageInfo,
    cw20_msg: Cw20ReceiveMsg,
) -> Result<Response, ContractError> {
    let config = CONFIG.load(deps.storage)?;

    // info.sender is the CW20 contract address
    let token_addr = info.sender;

    // cw20_msg.sender is the original sender (user who initiated CW20 Send)
    let user_addr = deps.api.addr_validate(&cw20_msg.sender)?;

    // Decode the message
    let swap_msg: SwapTokenForBwick = cosmwasm_std::from_slice(&cw20_msg.msg)?;

    // Token -> BWICK swap
    let mut pool = POOLS.load(deps.storage, &token_addr)?;

    let input_amount = cw20_msg.amount.u128();
    if input_amount == 0 {
        return Err(ContractError::ZeroAmount {});
    }

    // Calculate total fee (base + augmented if active)
    let base_fee_bps = config.swap_fee_bps;
    let total_fee_bps = if let Some(ref aug_fee) = pool.augmented_fee {
        if aug_fee.active {
            base_fee_bps + aug_fee.augmented_fee_bps
        } else {
            base_fee_bps
        }
    } else {
        base_fee_bps
    };

    // Calculate output using constant product
    let (output_amount, fee_amount) = calculate_swap_output(
        pool.token_reserve,
        pool.bwick_reserve,
        input_amount,
        total_fee_bps,
    )?;

    // Check slippage protection
    if Uint128::new(output_amount) < swap_msg.min_output {
        return Err(ContractError::SlippageExceeded {
            expected: swap_msg.min_output.u128(),
            actual: output_amount,
        });
    }

    // Update pool reserves
    // Fee stays in pool (auto-compound) - input includes fee
    pool.token_reserve += input_amount;
    pool.bwick_reserve -= output_amount;

    // Check if augmented fee should auto-disable
    let mut response = Response::new();
    if let Some(ref mut aug_fee) = pool.augmented_fee {
        if aug_fee.active {
            let pool_value = pool.bwick_reserve * 2;
            if pool_value >= aug_fee.lp_target_ubwick {
                aug_fee.active = false;
                response = response
                    .add_attribute("augmented_fee_disabled", "true")
                    .add_attribute("pool_value", pool_value.to_string())
                    .add_attribute("lp_target", aug_fee.lp_target_ubwick.to_string());
            }
        }
    }

    // Save updated pool
    POOLS.save(deps.storage, &token_addr, &pool)?;

    // Transfer native BWICK to the original sender
    let send_msg = BankMsg::Send {
        to_address: user_addr.to_string(),
        amount: vec![Coin {
            denom: "ubwick".to_string(),
            amount: Uint128::new(output_amount),
        }],
    };

    // Calculate augmented fee portion for attribute
    let augmented_fee_amount = if total_fee_bps > base_fee_bps {
        (input_amount * (total_fee_bps - base_fee_bps) as u128) / 10000
    } else {
        0
    };

    Ok(response
        .add_message(send_msg)
        .add_attribute("action", "swap")
        .add_attribute("direction", "token_to_bwick")
        .add_attribute("token_address", token_addr.to_string())
        .add_attribute("input_amount", input_amount.to_string())
        .add_attribute("output_amount", output_amount.to_string())
        .add_attribute("fee_amount", fee_amount.to_string())
        .add_attribute("augmented_fee_amount", augmented_fee_amount.to_string())
        .add_attribute("recipient", user_addr.to_string()))
}

// Helper functions

/// Calculate output amount using constant product formula
/// Returns (output_amount, fee_amount)
fn calculate_swap_output(
    input_reserve: u128,
    output_reserve: u128,
    input_amount: u128,
    fee_bps: u16,
) -> Result<(u128, u128), ContractError> {
    // Deduct fee from input first
    let fee_amount = input_amount * (fee_bps as u128) / 10000;
    let input_after_fee = input_amount - fee_amount;

    // Constant product: x * y = k
    // new_input_reserve * new_output_reserve = k
    // (input_reserve + input_after_fee) * (output_reserve - output_amount) = input_reserve * output_reserve
    // output_amount = output_reserve - (input_reserve * output_reserve) / (input_reserve + input_after_fee)
    // output_amount = output_reserve * input_after_fee / (input_reserve + input_after_fee)

    let numerator = output_reserve * input_after_fee;
    let denominator = input_reserve + input_after_fee;

    if denominator == 0 {
        return Err(ContractError::InsufficientLiquidity {});
    }

    let output_amount = numerator / denominator;

    if output_amount == 0 {
        return Err(ContractError::InsufficientLiquidity {});
    }

    if output_amount >= output_reserve {
        return Err(ContractError::InsufficientLiquidity {});
    }

    Ok((output_amount, fee_amount))
}

/// Calculate integer square root using Newton's method
fn integer_sqrt(n: u128) -> u128 {
    if n == 0 {
        return 0;
    }

    let mut x = n;
    let mut y = x.div_ceil(2);

    while y < x {
        x = y;
        y = (x + n / x).div_ceil(2);
    }

    x
}

/// Validate that the token is a valid CW20 contract
fn validate_cw20_token(deps: Deps, token_addr: &Addr) -> Result<(), ContractError> {
    // Query token_info to verify it's a valid CW20
    let query_msg = cw20::Cw20QueryMsg::TokenInfo {};
    let _: TokenInfoResponse = deps
        .querier
        .query_wasm_smart(token_addr.to_string(), &query_msg)
        .map_err(|_| {
            ContractError::Std(StdError::generic_err(format!(
                "Invalid CW20 token address: {}",
                token_addr
            )))
        })?;
    Ok(())
}

/// Extract BWICK amount from MessageInfo funds
fn extract_bwick_from_funds(funds: &[Coin]) -> Result<Uint128, ContractError> {
    // Must have exactly one coin with denom "ubwick"
    if funds.len() != 1 {
        return Err(ContractError::InvalidToken {});
    }

    let coin = &funds[0];
    if coin.denom != "ubwick" {
        return Err(ContractError::InvalidToken {});
    }

    Ok(coin.amount)
}

#[cfg_attr(not(feature = "library"), entry_point)]
pub fn query(deps: Deps, env: Env, msg: QueryMsg) -> StdResult<Binary> {
    match msg {
        QueryMsg::Pool { token_address } => to_binary(&query_pool(deps, token_address)?),
        QueryMsg::AllPools { start_after, limit } => {
            to_binary(&query_all_pools(deps, start_after, limit)?)
        }
        QueryMsg::SimulateSwap { token_address, offer_bwick, offer_amount } => {
            to_binary(&query_simulate_swap(deps, token_address, offer_bwick, offer_amount)?)
        }
        QueryMsg::Config {} => to_binary(&query_config(deps)?),
        QueryMsg::AugmentedFeeStatus { token_address } => {
            to_binary(&query_augmented_fee_status(deps, token_address)?)
        }
        QueryMsg::LpBalance { token_address, holder } => {
            to_binary(&query_lp_balance(deps, token_address, holder)?)
        }
        QueryMsg::SimulateAddLiquidity { token_address, bwick_amount } => {
            to_binary(&query_simulate_add_liquidity(deps, token_address, bwick_amount)?)
        }
        QueryMsg::LockedLp { token_address } => {
            to_binary(&query_locked_lp(deps, env, token_address)?)
        }
    }
}

fn query_locked_lp(
    deps: Deps,
    env: Env,
    token_address: String,
) -> StdResult<crate::msg::LockedLpResponse> {
    let token_addr = deps.api.addr_validate(&token_address)?;
    let pool = POOLS
        .load(deps.storage, &token_addr)
        .map_err(|_| StdError::not_found(format!("Pool for token {}", token_address)))?;
    let now = env.block.time.seconds();
    let (bwick_share, token_share) = if pool.lp_total_supply > 0 && pool.locked_lp_supply > 0 {
        (
            pool.bwick_reserve.saturating_mul(pool.locked_lp_supply) / pool.lp_total_supply,
            pool.token_reserve.saturating_mul(pool.locked_lp_supply) / pool.lp_total_supply,
        )
    } else {
        (0u128, 0u128)
    };
    let unlocked = matches!(pool.locked_lp_unlock_at, Some(t) if now >= t);
    Ok(crate::msg::LockedLpResponse {
        token_address: token_addr,
        locked_lp_supply: Uint128::from(pool.locked_lp_supply),
        unlock_at: pool.locked_lp_unlock_at,
        recipient: pool.locked_lp_recipient,
        unlocked,
        bwick_share: Uint128::from(bwick_share),
        token_share: Uint128::from(token_share),
    })
}

fn query_lp_balance(
    deps: Deps,
    token_address: String,
    holder: String,
) -> StdResult<LpBalanceResponse> {
    let token_addr = deps.api.addr_validate(&token_address)?;
    let holder_addr = deps.api.addr_validate(&holder)?;
    let pool = POOLS.load(deps.storage, &token_addr)
        .map_err(|_| StdError::not_found(format!("Pool for token {}", token_address)))?;
    let shares = OPEN_LP.may_load(deps.storage, (&token_addr, &holder_addr))?.unwrap_or(0);

    let (bwick_share, token_share) = if pool.lp_total_supply > 0 && shares > 0 {
        (
            pool.bwick_reserve.saturating_mul(shares) / pool.lp_total_supply,
            pool.token_reserve.saturating_mul(shares) / pool.lp_total_supply,
        )
    } else {
        (0u128, 0u128)
    };

    Ok(LpBalanceResponse {
        holder: holder_addr,
        token_address: token_addr,
        shares: Uint128::from(shares),
        bwick_share: Uint128::from(bwick_share),
        token_share: Uint128::from(token_share),
        lp_total_supply: Uint128::from(pool.lp_total_supply),
        locked_lp_supply: Uint128::from(pool.locked_lp_supply),
    })
}

fn query_simulate_add_liquidity(
    deps: Deps,
    token_address: String,
    bwick_amount: Uint128,
) -> StdResult<SimulateAddLiquidityResponse> {
    let token_addr = deps.api.addr_validate(&token_address)?;
    let pool = POOLS.load(deps.storage, &token_addr)
        .map_err(|_| StdError::not_found(format!("Pool for token {}", token_address)))?;
    if pool.lp_total_supply == 0 || pool.bwick_reserve == 0 {
        return Err(StdError::generic_err("pool not initialized"));
    }
    let bwick_added = bwick_amount.u128();
    if bwick_added == 0 {
        return Err(StdError::generic_err("bwick_amount must be positive"));
    }
    let token_required_num = bwick_added.saturating_mul(pool.token_reserve);
    let token_required = (token_required_num + pool.bwick_reserve - 1) / pool.bwick_reserve;
    let lp_shares = pool.lp_total_supply.saturating_mul(bwick_added) / pool.bwick_reserve;
    Ok(SimulateAddLiquidityResponse {
        token_amount_required: Uint128::from(token_required),
        lp_shares_minted: Uint128::from(lp_shares),
        bwick_reserve_before: Uint128::from(pool.bwick_reserve),
        token_reserve_before: Uint128::from(pool.token_reserve),
    })
}

fn query_pool(deps: Deps, token_address: String) -> StdResult<PoolResponse> {
    let token_addr = deps.api.addr_validate(&token_address)?;
    let pool = POOLS.load(deps.storage, &token_addr)
        .map_err(|_| StdError::not_found(format!("Pool for token {}", token_address)))?;

    // Calculate price: BWICK per token
    let price = if pool.token_reserve > 0 {
        format!("{:.6}", pool.bwick_reserve as f64 / pool.token_reserve as f64)
    } else {
        "0".to_string()
    };

    Ok(PoolResponse {
        token_address: pool.token_address,
        bwick_reserve: Uint128::from(pool.bwick_reserve),
        token_reserve: Uint128::from(pool.token_reserve),
        lp_token_address: pool.lp_token_address,
        lp_total_supply: Uint128::from(pool.lp_total_supply),
        price,
        locked_lp_unlock_at: pool.locked_lp_unlock_at,
    })
}

fn query_all_pools(
    deps: Deps,
    start_after: Option<String>,
    limit: Option<u32>,
) -> StdResult<AllPoolsResponse> {
    let limit = limit.unwrap_or(10).min(30) as usize;
    let start = start_after
        .map(|s| deps.api.addr_validate(&s))
        .transpose()?;

    let pools: Vec<PoolResponse> = POOLS
        .range(
            deps.storage,
            start.as_ref().map(Bound::exclusive),
            None,
            Order::Ascending,
        )
        .take(limit)
        .map(|item| {
            let (_, pool) = item?;
            let price = if pool.token_reserve > 0 {
                format!("{:.6}", pool.bwick_reserve as f64 / pool.token_reserve as f64)
            } else {
                "0".to_string()
            };
            Ok(PoolResponse {
                token_address: pool.token_address,
                bwick_reserve: Uint128::from(pool.bwick_reserve),
                token_reserve: Uint128::from(pool.token_reserve),
                lp_token_address: pool.lp_token_address,
                lp_total_supply: Uint128::from(pool.lp_total_supply),
                price,
                locked_lp_unlock_at: pool.locked_lp_unlock_at,
            })
        })
        .collect::<StdResult<Vec<_>>>()?;

    Ok(AllPoolsResponse { pools })
}

fn query_simulate_swap(
    deps: Deps,
    token_address: String,
    offer_bwick: bool,
    offer_amount: Uint128,
) -> StdResult<SimulateSwapResponse> {
    let token_addr = deps.api.addr_validate(&token_address)?;
    let pool = POOLS.load(deps.storage, &token_addr)
        .map_err(|_| StdError::not_found(format!("Pool for token {}", token_address)))?;
    let config = CONFIG.load(deps.storage)?;

    let (input_reserve, output_reserve) = if offer_bwick {
        (pool.bwick_reserve, pool.token_reserve)
    } else {
        (pool.token_reserve, pool.bwick_reserve)
    };

    // Calculate total fee (base + augmented if active)
    let base_fee_bps = config.swap_fee_bps;
    let total_fee_bps = if let Some(ref aug_fee) = pool.augmented_fee {
        if aug_fee.active {
            base_fee_bps + aug_fee.augmented_fee_bps
        } else {
            base_fee_bps
        }
    } else {
        base_fee_bps
    };

    let offer = offer_amount.u128();
    let (output_amount, fee_amount) = calculate_swap_output(
        input_reserve,
        output_reserve,
        offer,
        total_fee_bps,
    ).map_err(|e| StdError::generic_err(format!("Swap calculation failed: {:?}", e)))?;

    // Calculate price impact
    // Price impact = (new_price - old_price) / old_price
    // old_price = output_reserve / input_reserve
    // new_price = (output_reserve - output) / (input_reserve + offer - fee)
    let old_price = output_reserve as f64 / input_reserve as f64;
    let new_input = input_reserve + offer - fee_amount;
    let new_output = output_reserve - output_amount;
    let new_price = new_output as f64 / new_input as f64;
    let price_impact = ((old_price - new_price) / old_price * 100.0).abs();

    // Calculate augmented fee portion
    let augmented_fee_amount = if total_fee_bps > base_fee_bps {
        (offer * (total_fee_bps - base_fee_bps) as u128) / 10000
    } else {
        0
    };

    Ok(SimulateSwapResponse {
        output_amount: Uint128::from(output_amount),
        fee_amount: Uint128::from(fee_amount),
        price_impact: format!("{:.2}%", price_impact),
        augmented_fee_amount: Uint128::from(augmented_fee_amount),
    })
}

fn query_config(deps: Deps) -> StdResult<ConfigResponse> {
    let config = CONFIG.load(deps.storage)?;
    Ok(ConfigResponse {
        authorized_creators: config.authorized_creators,
        swap_fee_bps: config.swap_fee_bps,
        max_wallet_bps: config.max_wallet_bps,
    })
}

fn query_augmented_fee_status(
    deps: Deps,
    token_address: String,
) -> StdResult<AugmentedFeeStatusResponse> {
    let token_addr = deps.api.addr_validate(&token_address)?;
    let pool = POOLS.load(deps.storage, &token_addr)
        .map_err(|_| StdError::not_found(format!("Pool for token {}", token_address)))?;

    let current_pool_value = pool.bwick_reserve * 2;

    match pool.augmented_fee {
        Some(aug_fee) => {
            let progress_percent = if aug_fee.lp_target_ubwick > 0 {
                format!(
                    "{:.2}",
                    (current_pool_value as f64 / aug_fee.lp_target_ubwick as f64) * 100.0
                )
            } else {
                "0.00".to_string()
            };

            Ok(AugmentedFeeStatusResponse {
                active: aug_fee.active,
                augmented_fee_bps: aug_fee.augmented_fee_bps,
                lp_target_ubwick: Uint128::from(aug_fee.lp_target_ubwick),
                current_pool_value_ubwick: Uint128::from(current_pool_value),
                progress_percent,
            })
        }
        None => Ok(AugmentedFeeStatusResponse {
            active: false,
            augmented_fee_bps: 0,
            lp_target_ubwick: Uint128::zero(),
            current_pool_value_ubwick: Uint128::from(current_pool_value),
            progress_percent: "N/A".to_string(),
        }),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use cosmwasm_std::testing::mock_dependencies;
    use cosmwasm_std::Addr;

    const BONDING_CURVE: &str = "bonding_curve_contract";
    const TOKEN: &str = "token_contract";

    /// Setup contract state directly without going through instantiate
    /// This avoids bech32 address validation in mock_dependencies
    fn setup_contract_state(deps: DepsMut) {
        let config = Config {
            authorized_creators: vec![Addr::unchecked(BONDING_CURVE)],
            swap_fee_bps: 100, // 1% fee
            max_wallet_bps: 300, // 3% max wallet
        };
        CONFIG.save(deps.storage, &config).unwrap();
    }

    fn setup_pool(deps: DepsMut, bwick_reserve: u128, token_reserve: u128) {
        let token_addr = Addr::unchecked(TOKEN);
        let pool = Pool {
            token_address: token_addr.clone(),
            bwick_reserve,
            token_reserve,
            lp_token_address: Addr::unchecked("contract"),
            lp_total_supply: integer_sqrt(bwick_reserve * token_reserve),
            locked_lp_supply: 0,
            locked_lp_unlock_at: None,
            locked_lp_recipient: None,
            augmented_fee: None,
        };
        POOLS.save(deps.storage, &token_addr, &pool).unwrap();
    }

    // ==================== v3: LP lock tests ====================

    fn setup_locked_pool(
        deps: DepsMut,
        bwick_reserve: u128,
        token_reserve: u128,
        locked_supply: u128,
        unlock_at: Option<u64>,
        recipient: Option<&str>,
    ) {
        let token_addr = Addr::unchecked(TOKEN);
        let pool = Pool {
            token_address: token_addr.clone(),
            bwick_reserve,
            token_reserve,
            lp_token_address: Addr::unchecked("contract"),
            lp_total_supply: locked_supply,
            locked_lp_supply: locked_supply,
            locked_lp_unlock_at: unlock_at,
            locked_lp_recipient: recipient.map(Addr::unchecked),
            augmented_fee: None,
        };
        POOLS.save(deps.storage, &token_addr, &pool).unwrap();
    }

    #[test]
    fn withdraw_locked_lp_rejects_permanent_lock() {
        let mut deps = mock_dependencies();
        setup_contract_state(deps.as_mut());
        setup_locked_pool(deps.as_mut(), 1000, 1000, 1000, None, None);
        let info = cosmwasm_std::testing::mock_info("anyone", &[]);
        let env = cosmwasm_std::testing::mock_env();
        let err = execute_withdraw_locked_lp(
            deps.as_mut(),
            env,
            info,
            TOKEN.to_string(),
            None,
        )
        .unwrap_err();
        assert!(matches!(err, ContractError::LockedLpPermanent {}));
    }

    #[test]
    fn withdraw_locked_lp_rejects_before_unlock() {
        let mut deps = mock_dependencies();
        setup_contract_state(deps.as_mut());
        let mut env = cosmwasm_std::testing::mock_env();
        let unlock_at = env.block.time.seconds() + 10_000;
        setup_locked_pool(
            deps.as_mut(),
            1000,
            1000,
            1000,
            Some(unlock_at),
            Some("creator"),
        );
        let info = cosmwasm_std::testing::mock_info("creator", &[]);
        env.block.time = cosmwasm_std::Timestamp::from_seconds(unlock_at - 1);
        let err = execute_withdraw_locked_lp(
            deps.as_mut(),
            env,
            info,
            TOKEN.to_string(),
            None,
        )
        .unwrap_err();
        assert!(matches!(err, ContractError::LockedLpStillLocked { .. }));
    }

    #[test]
    fn withdraw_locked_lp_rejects_wrong_caller() {
        let mut deps = mock_dependencies();
        setup_contract_state(deps.as_mut());
        let mut env = cosmwasm_std::testing::mock_env();
        let unlock_at = env.block.time.seconds();
        setup_locked_pool(
            deps.as_mut(),
            1000,
            1000,
            1000,
            Some(unlock_at),
            Some("creator"),
        );
        let info = cosmwasm_std::testing::mock_info("imposter", &[]);
        env.block.time = cosmwasm_std::Timestamp::from_seconds(unlock_at + 1);
        let err = execute_withdraw_locked_lp(
            deps.as_mut(),
            env,
            info,
            TOKEN.to_string(),
            None,
        )
        .unwrap_err();
        assert!(matches!(err, ContractError::LockedLpUnauthorized {}));
    }

    #[test]
    fn withdraw_locked_lp_succeeds_after_unlock_and_drains_seed() {
        let mut deps = mock_dependencies();
        setup_contract_state(deps.as_mut());
        let mut env = cosmwasm_std::testing::mock_env();
        let unlock_at = env.block.time.seconds();
        // pool: 1000 BWICK, 5000 tokens, 1000 LP, all locked.
        setup_locked_pool(
            deps.as_mut(),
            1000,
            5000,
            1000,
            Some(unlock_at),
            Some("creator"),
        );
        env.block.time = cosmwasm_std::Timestamp::from_seconds(unlock_at);
        let info = cosmwasm_std::testing::mock_info("creator", &[]);
        let res = execute_withdraw_locked_lp(
            deps.as_mut(),
            env,
            info,
            TOKEN.to_string(),
            None,
        )
        .unwrap();
        // 2 messages: bank send BWICK, wasm transfer tokens.
        assert_eq!(res.messages.len(), 2);
        let pool = POOLS
            .load(deps.as_ref().storage, &Addr::unchecked(TOKEN))
            .unwrap();
        assert_eq!(pool.locked_lp_supply, 0);
        assert_eq!(pool.lp_total_supply, 0);
        assert_eq!(pool.bwick_reserve, 0);
        assert_eq!(pool.token_reserve, 0);

        // A second withdraw attempt rejects with AlreadyWithdrawn (since
        // locked_lp_supply is 0). Note: lp_total_supply is also 0 now so
        // either check could trip; we want the dedicated error.
        let info2 = cosmwasm_std::testing::mock_info("creator", &[]);
        let env2 = cosmwasm_std::testing::mock_env();
        let err = execute_withdraw_locked_lp(
            deps.as_mut(),
            env2,
            info2,
            TOKEN.to_string(),
            None,
        )
        .unwrap_err();
        assert!(matches!(err, ContractError::LockedLpAlreadyWithdrawn {}));
    }

    // ==================== Arithmetic Edge Case Tests ====================

    #[test]
    fn test_calculate_swap_output_normal() {
        // Normal swap: 1000 BWICK in, pool has 10000 BWICK and 10000 tokens
        // Fee: 1000 * 100 / 10000 = 10, input after fee = 990
        // Output = 10000 * 990 / (10000 + 990) = 9900000 / 10990 = 900 (integer division)
        let result = calculate_swap_output(10000, 10000, 1000, 100).unwrap();
        assert_eq!(result.0, 900); // output amount (integer division rounds down)
        assert_eq!(result.1, 10);  // fee amount
    }

    #[test]
    fn test_calculate_swap_output_large_values() {
        // Test with large but safe values (10^18 scale typical for tokens)
        let input_reserve: u128 = 1_000_000_000_000_000_000; // 1e18
        let output_reserve: u128 = 1_000_000_000_000_000_000; // 1e18
        let input_amount: u128 = 1_000_000_000_000_000; // 1e15 (0.1% of pool)

        let result = calculate_swap_output(input_reserve, output_reserve, input_amount, 100);
        assert!(result.is_ok());
        let (output, fee) = result.unwrap();
        assert!(output > 0);
        assert!(output < output_reserve);
        assert_eq!(fee, input_amount / 100); // 1% fee
    }

    #[test]
    fn test_calculate_swap_output_max_fee() {
        // Test with maximum fee (100% = 10000 bps)
        // All input becomes fee, so input_after_fee = 0
        // This should still work (output will be 0, triggering InsufficientLiquidity)
        let result = calculate_swap_output(10000, 10000, 1000, 10000);
        assert!(result.is_err());
        match result {
            Err(ContractError::InsufficientLiquidity {}) => {}
            _ => panic!("Expected InsufficientLiquidity error"),
        }
    }

    #[test]
    fn test_calculate_swap_output_zero_fee() {
        // Test with 0% fee
        // Output = 10000 * 1000 / (10000 + 1000) = 10000000 / 11000 = 909
        let result = calculate_swap_output(10000, 10000, 1000, 0).unwrap();
        assert_eq!(result.0, 909); // output amount
        assert_eq!(result.1, 0);   // fee amount
    }

    #[test]
    fn test_calculate_swap_output_small_input() {
        // Very small input relative to pool
        // Input = 1, fee = 0 (rounds down), output should be minimal but > 0
        let result = calculate_swap_output(1_000_000, 1_000_000, 1, 100);
        // Output = 1_000_000 * 1 / (1_000_000 + 1) = 0 (rounds down)
        assert!(result.is_err()); // InsufficientLiquidity because output = 0
    }

    #[test]
    fn test_fee_calculation_boundaries() {
        // Test fee calculation at various boundaries
        // 1 bps = 0.01% fee
        let (_, fee1) = calculate_swap_output(10000, 10000, 10000, 1).unwrap();
        assert_eq!(fee1, 1); // 10000 * 1 / 10000 = 1

        // 9999 bps = 99.99% fee
        let result = calculate_swap_output(10000, 10000, 10000, 9999);
        // Fee = 10000 * 9999 / 10000 = 9999, input after fee = 1
        // Output = 10000 * 1 / (10000 + 1) = 0 (rounds down)
        assert!(result.is_err()); // InsufficientLiquidity
    }

    // ==================== Zero/Empty Handling Tests ====================

    #[test]
    fn test_calculate_swap_zero_input() {
        // Zero input amount
        let result = calculate_swap_output(10000, 10000, 0, 100);
        assert!(result.is_err());
        match result {
            Err(ContractError::InsufficientLiquidity {}) => {}
            _ => panic!("Expected InsufficientLiquidity error for zero input"),
        }
    }

    #[test]
    fn test_calculate_swap_zero_input_reserve() {
        // Zero input reserve (empty pool side)
        let result = calculate_swap_output(0, 10000, 1000, 100);
        // denominator = 0 + (1000 - 10) = 990, not zero
        // numerator = 10000 * 990 = 9900000
        // output = 9900000 / 990 = 10000 >= output_reserve
        assert!(result.is_err());
        match result {
            Err(ContractError::InsufficientLiquidity {}) => {}
            _ => panic!("Expected InsufficientLiquidity error"),
        }
    }

    #[test]
    fn test_calculate_swap_zero_output_reserve() {
        // Zero output reserve (would drain entire side)
        let result = calculate_swap_output(10000, 0, 1000, 100);
        // output = 0 * 990 / 10990 = 0
        assert!(result.is_err());
        match result {
            Err(ContractError::InsufficientLiquidity {}) => {}
            _ => panic!("Expected InsufficientLiquidity error"),
        }
    }

    #[test]
    fn test_query_nonexistent_pool() {
        let mut deps = mock_dependencies();
        setup_contract_state(deps.as_mut());

        // Query pool that doesn't exist - use a valid-looking address
        // Since mock_dependencies doesn't validate bech32, just check the pool doesn't exist
        let token_addr = Addr::unchecked("nonexistent_token");
        let result = POOLS.may_load(deps.as_ref().storage, &token_addr).unwrap();
        assert!(result.is_none());
    }

    #[test]
    fn test_simulate_swap_nonexistent_pool() {
        let mut deps = mock_dependencies();
        setup_contract_state(deps.as_mut());

        // Directly test the logic - pool doesn't exist
        let token_addr = Addr::unchecked("nonexistent_token");
        let result = POOLS.may_load(deps.as_ref().storage, &token_addr).unwrap();
        assert!(result.is_none());
    }

    // ==================== Slippage Edge Case Tests ====================

    #[test]
    fn test_slippage_exact_match() {
        // Test when output exactly equals min_receive
        // Calculate expected output first
        let (expected_output, _) = calculate_swap_output(10000, 10000, 1000, 100).unwrap();
        assert_eq!(expected_output, 900); // 9900000 / 10990 = 900 (integer division)

        // min_output exactly equals expected output - should pass
        let min_output = Uint128::new(expected_output);
        let actual = Uint128::new(expected_output);

        // This is the check from execute_swap
        assert!(actual >= min_output); // Should pass (equal)
    }

    #[test]
    fn test_slippage_off_by_one_fail() {
        // Test when output is 1 less than min_receive - should fail
        let (expected_output, _) = calculate_swap_output(10000, 10000, 1000, 100).unwrap();
        assert_eq!(expected_output, 900); // 9900000 / 10990 = 900

        // min_output is 1 more than actual output
        let min_output = Uint128::new(expected_output + 1);
        let actual = Uint128::new(expected_output);

        // This simulates the slippage check
        assert!(actual < min_output); // Would fail slippage check
    }

    #[test]
    fn test_slippage_with_generous_tolerance() {
        // Test with min_output well below actual
        let (expected_output, _) = calculate_swap_output(10000, 10000, 1000, 100).unwrap();

        let min_output = Uint128::new(expected_output / 2); // 50% tolerance
        let actual = Uint128::new(expected_output);

        assert!(actual >= min_output); // Should pass easily
    }

    // ==================== Integer Square Root Tests ====================

    #[test]
    fn test_integer_sqrt_zero() {
        assert_eq!(integer_sqrt(0), 0);
    }

    #[test]
    fn test_integer_sqrt_one() {
        assert_eq!(integer_sqrt(1), 1);
    }

    #[test]
    fn test_integer_sqrt_perfect_squares() {
        assert_eq!(integer_sqrt(4), 2);
        assert_eq!(integer_sqrt(9), 3);
        assert_eq!(integer_sqrt(16), 4);
        assert_eq!(integer_sqrt(100), 10);
        assert_eq!(integer_sqrt(10000), 100);
        assert_eq!(integer_sqrt(1_000_000), 1000);
    }

    #[test]
    fn test_integer_sqrt_non_perfect_squares() {
        // Newton's method with div_ceil - actual behavior verified
        // Note: This implementation uses ceiling division which affects convergence
        assert_eq!(integer_sqrt(2), 1);   // sqrt(2) = 1.41 -> 1
        assert_eq!(integer_sqrt(3), 2);   // sqrt(3) = 1.73 -> 2
        assert_eq!(integer_sqrt(5), 2);   // sqrt(5) = 2.23 -> 2
        assert_eq!(integer_sqrt(8), 3);   // sqrt(8) = 2.83 -> 3
        assert_eq!(integer_sqrt(99), 10); // sqrt(99) = 9.95 -> 10
        assert_eq!(integer_sqrt(101), 10); // sqrt(101) = 10.05 -> 10
    }

    #[test]
    fn test_integer_sqrt_large_value() {
        // Test with large value (10^36 which is sqrt(10^18) * sqrt(10^18))
        let large: u128 = 1_000_000_000_000_000_000_000_000_000_000_000_000; // 10^36
        let sqrt = integer_sqrt(large);
        assert_eq!(sqrt, 1_000_000_000_000_000_000); // 10^18
    }

    // ==================== Instantiation Tests ====================

    #[test]
    fn test_instantiate_stores_config() {
        // Test that config is stored correctly (bypass addr_validate by storing directly)
        let mut deps = mock_dependencies();
        let config = Config {
            authorized_creators: vec![Addr::unchecked(BONDING_CURVE)],
            swap_fee_bps: 100,
            max_wallet_bps: 300,
        };
        CONFIG.save(deps.as_mut().storage, &config).unwrap();

        let loaded = CONFIG.load(deps.as_ref().storage).unwrap();
        assert_eq!(loaded.authorized_creators, vec![Addr::unchecked(BONDING_CURVE)]);
        assert_eq!(loaded.swap_fee_bps, 100);
    }

    #[test]
    fn test_instantiate_fee_validation_logic() {
        // Test the fee validation logic directly
        let valid_fee: u16 = 10000;
        let invalid_fee: u16 = 10001;

        assert!(valid_fee <= 10000);
        assert!(invalid_fee > 10000);
    }

    #[test]
    fn test_instantiate_fee_boundary() {
        // Test fee boundary values
        assert!(100 <= 10000);   // 1% fee is valid
        assert!(0 <= 10000);     // 0% fee is valid
        assert!(10000 <= 10000); // 100% fee is valid (edge case)
        assert!(10001 > 10000);  // 100.01% is invalid
    }

    // ==================== Access Control Tests ====================

    #[test]
    fn test_extract_bwick_from_funds_success() {
        let funds = vec![Coin {
            denom: "ubwick".to_string(),
            amount: Uint128::new(1000),
        }];
        let result = extract_bwick_from_funds(&funds);
        assert!(result.is_ok());
        assert_eq!(result.unwrap(), Uint128::new(1000));
    }

    #[test]
    fn test_extract_bwick_from_funds_wrong_denom() {
        let funds = vec![Coin {
            denom: "uatom".to_string(),
            amount: Uint128::new(1000),
        }];
        let result = extract_bwick_from_funds(&funds);
        assert!(result.is_err());
        match result {
            Err(ContractError::InvalidToken {}) => {}
            _ => panic!("Expected InvalidToken error"),
        }
    }

    #[test]
    fn test_extract_bwick_from_funds_multiple_coins() {
        let funds = vec![
            Coin {
                denom: "ubwick".to_string(),
                amount: Uint128::new(1000),
            },
            Coin {
                denom: "uatom".to_string(),
                amount: Uint128::new(500),
            },
        ];
        let result = extract_bwick_from_funds(&funds);
        assert!(result.is_err());
        match result {
            Err(ContractError::InvalidToken {}) => {}
            _ => panic!("Expected InvalidToken error"),
        }
    }

    #[test]
    fn test_extract_bwick_from_funds_empty() {
        let funds: Vec<Coin> = vec![];
        let result = extract_bwick_from_funds(&funds);
        assert!(result.is_err());
        match result {
            Err(ContractError::InvalidToken {}) => {}
            _ => panic!("Expected InvalidToken error"),
        }
    }

    // ==================== Query Tests ====================

    #[test]
    fn test_query_config() {
        let mut deps = mock_dependencies();
        setup_contract_state(deps.as_mut());

        let result = query_config(deps.as_ref()).unwrap();
        assert_eq!(result.authorized_creators, vec![Addr::unchecked(BONDING_CURVE)]);
        assert_eq!(result.swap_fee_bps, 100);
    }

    #[test]
    fn test_query_pool_direct() {
        // Test pool loading directly without address validation
        let mut deps = mock_dependencies();
        setup_contract_state(deps.as_mut());
        setup_pool(deps.as_mut(), 10000, 5000);

        // Load pool directly
        let token_addr = Addr::unchecked(TOKEN);
        let pool = POOLS.load(deps.as_ref().storage, &token_addr).unwrap();
        assert_eq!(pool.bwick_reserve, 10000);
        assert_eq!(pool.token_reserve, 5000);

        // Verify price calculation logic (10000/5000 = 2.0)
        let price = if pool.token_reserve > 0 {
            format!("{:.6}", pool.bwick_reserve as f64 / pool.token_reserve as f64)
        } else {
            "0".to_string()
        };
        assert_eq!(price, "2.000000");
    }

    #[test]
    fn test_query_all_pools_pagination_direct() {
        let mut deps = mock_dependencies();
        setup_contract_state(deps.as_mut());

        // Create multiple pools
        for i in 0..5 {
            let token_addr = Addr::unchecked(format!("token{}", i));
            let pool = Pool {
                token_address: token_addr.clone(),
                bwick_reserve: 10000,
                token_reserve: 10000,
                lp_token_address: Addr::unchecked("contract"),
                lp_total_supply: 10000,
                locked_lp_supply: 0,
                locked_lp_unlock_at: None,
                locked_lp_recipient: None,
                augmented_fee: None,
            };
            POOLS.save(deps.as_mut().storage, &token_addr, &pool).unwrap();
        }

        // Query with limit - test range iteration directly
        let pools: Vec<_> = POOLS
            .range(deps.as_ref().storage, None, None, Order::Ascending)
            .take(2)
            .collect::<StdResult<Vec<_>>>()
            .unwrap();
        assert_eq!(pools.len(), 2);

        // Query all
        let pools: Vec<_> = POOLS
            .range(deps.as_ref().storage, None, None, Order::Ascending)
            .take(10)
            .collect::<StdResult<Vec<_>>>()
            .unwrap();
        assert_eq!(pools.len(), 5);
    }

    #[test]
    fn test_query_all_pools_max_limit_direct() {
        let mut deps = mock_dependencies();
        setup_contract_state(deps.as_mut());

        // Create 35 pools
        for i in 0..35 {
            let token_addr = Addr::unchecked(format!("token{:02}", i));
            let pool = Pool {
                token_address: token_addr.clone(),
                bwick_reserve: 10000,
                token_reserve: 10000,
                lp_token_address: Addr::unchecked("contract"),
                lp_total_supply: 10000,
                locked_lp_supply: 0,
                locked_lp_unlock_at: None,
                locked_lp_recipient: None,
                augmented_fee: None,
            };
            POOLS.save(deps.as_mut().storage, &token_addr, &pool).unwrap();
        }

        // Request limit > 30 should cap at 30 - test the limit calculation
        let requested_limit: u32 = 50;
        let actual_limit = requested_limit.min(30) as usize;
        assert_eq!(actual_limit, 30);

        // Verify we can query with the capped limit
        let pools: Vec<_> = POOLS
            .range(deps.as_ref().storage, None, None, Order::Ascending)
            .take(actual_limit)
            .collect::<StdResult<Vec<_>>>()
            .unwrap();
        assert_eq!(pools.len(), 30);
    }

    // ==================== Augmented Fee Tests ====================

    #[test]
    fn test_authorized_creators_multiple() {
        let mut deps = mock_dependencies();
        let tokenlaunch_addr = "tokenlaunch_module";

        // Setup config with two authorized creators
        let config = Config {
            authorized_creators: vec![
                Addr::unchecked(BONDING_CURVE),
                Addr::unchecked(tokenlaunch_addr),
            ],
            swap_fee_bps: 100,
            max_wallet_bps: 300,
        };
        CONFIG.save(deps.as_mut().storage, &config).unwrap();

        let loaded = CONFIG.load(deps.as_ref().storage).unwrap();
        assert_eq!(loaded.authorized_creators.len(), 2);
        assert!(loaded.authorized_creators.contains(&Addr::unchecked(BONDING_CURVE)));
        assert!(loaded.authorized_creators.contains(&Addr::unchecked(tokenlaunch_addr)));
    }

    #[test]
    fn test_augmented_fee_calculation() {
        let mut deps = mock_dependencies();
        setup_contract_state(deps.as_mut());

        // Setup pool with augmented fee
        let token_addr = Addr::unchecked(TOKEN);
        let pool = Pool {
            token_address: token_addr.clone(),
            bwick_reserve: 10000,
            token_reserve: 10000,
            lp_token_address: Addr::unchecked("contract"),
            lp_total_supply: 10000,
            locked_lp_supply: 0,
            locked_lp_unlock_at: None,
            locked_lp_recipient: None,
            augmented_fee: Some(AugmentedFeeConfig {
                augmented_fee_bps: 100, // 1% augmented fee
                lp_target_ubwick: 100000,
                active: true,
            }),
        };
        POOLS.save(deps.as_mut().storage, &token_addr, &pool).unwrap();

        // Calculate swap with augmented fee
        // Base fee: 100 bps (1%), Augmented fee: 100 bps (1%), Total: 200 bps (2%)
        let input_amount = 1000u128;
        let total_fee_bps = 200u16; // base + augmented
        let (output, total_fee) = calculate_swap_output(10000, 10000, input_amount, total_fee_bps).unwrap();

        // Expected: fee = 1000 * 200 / 10000 = 20
        // input_after_fee = 1000 - 20 = 980
        // output = 10000 * 980 / (10000 + 980) = 9800000 / 10980 = 892
        assert_eq!(total_fee, 20);
        assert_eq!(output, 892);
    }

    #[test]
    fn test_augmented_fee_auto_disable() {
        let mut deps = mock_dependencies();
        setup_contract_state(deps.as_mut());

        // Setup pool with augmented fee close to target
        let token_addr = Addr::unchecked(TOKEN);
        let pool = Pool {
            token_address: token_addr.clone(),
            bwick_reserve: 9900, // pool value = 9900 * 2 = 19800
            token_reserve: 10000,
            lp_token_address: Addr::unchecked("contract"),
            lp_total_supply: 10000,
            locked_lp_supply: 0,
            locked_lp_unlock_at: None,
            locked_lp_recipient: None,
            augmented_fee: Some(AugmentedFeeConfig {
                augmented_fee_bps: 100,
                lp_target_ubwick: 20000, // target is 20000
                active: true,
            }),
        };
        POOLS.save(deps.as_mut().storage, &token_addr, &pool).unwrap();

        // Simulate swap that pushes pool value over target
        // Add 100 BWICK -> pool value becomes (9900 + 100) * 2 = 20000
        let mut pool = POOLS.load(deps.as_ref().storage, &token_addr).unwrap();
        pool.bwick_reserve += 100;

        // Check auto-disable condition
        let pool_value = pool.bwick_reserve * 2;
        assert!(pool_value >= pool.augmented_fee.as_ref().unwrap().lp_target_ubwick);

        // Auto-disable should trigger
        if let Some(ref mut aug_fee) = pool.augmented_fee {
            if aug_fee.active && pool_value >= aug_fee.lp_target_ubwick {
                aug_fee.active = false;
            }
        }

        assert!(!pool.augmented_fee.as_ref().unwrap().active);
    }

    #[test]
    fn test_pool_without_augmented_fee() {
        let mut deps = mock_dependencies();
        setup_contract_state(deps.as_mut());

        // Setup normal pool without augmented fee
        let token_addr = Addr::unchecked(TOKEN);
        let pool = Pool {
            token_address: token_addr.clone(),
            bwick_reserve: 10000,
            token_reserve: 10000,
            lp_token_address: Addr::unchecked("contract"),
            lp_total_supply: 10000,
            locked_lp_supply: 0,
            locked_lp_unlock_at: None,
            locked_lp_recipient: None,
            augmented_fee: None,
        };
        POOLS.save(deps.as_mut().storage, &token_addr, &pool).unwrap();

        // Verify swap uses only base fee (100 bps = 1%)
        let input_amount = 1000u128;
        let base_fee_bps = 100u16;
        let (output, fee) = calculate_swap_output(10000, 10000, input_amount, base_fee_bps).unwrap();

        // Expected: fee = 1000 * 100 / 10000 = 10
        // output = 10000 * 990 / (10000 + 990) = 9900000 / 10990 = 900
        assert_eq!(fee, 10);
        assert_eq!(output, 900);
    }

    #[test]
    fn test_augmented_fee_query_status() {
        let mut deps = mock_dependencies();
        setup_contract_state(deps.as_mut());

        // Pool with active augmented fee
        let token_addr = Addr::unchecked(TOKEN);
        let pool = Pool {
            token_address: token_addr.clone(),
            bwick_reserve: 5000, // pool value = 10000
            token_reserve: 10000,
            lp_token_address: Addr::unchecked("contract"),
            lp_total_supply: 10000,
            locked_lp_supply: 0,
            locked_lp_unlock_at: None,
            locked_lp_recipient: None,
            augmented_fee: Some(AugmentedFeeConfig {
                augmented_fee_bps: 200, // 2%
                lp_target_ubwick: 20000,
                active: true,
            }),
        };
        POOLS.save(deps.as_mut().storage, &token_addr, &pool).unwrap();

        let status = query_augmented_fee_status(deps.as_ref(), TOKEN.to_string()).unwrap();
        assert_eq!(status.active, true);
        assert_eq!(status.augmented_fee_bps, 200);
        assert_eq!(status.lp_target_ubwick, Uint128::from(20000u128));
        assert_eq!(status.current_pool_value_ubwick, Uint128::from(10000u128));
        assert_eq!(status.progress_percent, "50.00"); // 10000/20000 * 100 = 50%
    }

    // ==================== Fee Cap Validation Tests ====================

    #[test]
    fn test_instantiate_fee_cap_rejects_over_1000() {
        // Base swap_fee_bps > 1000 should be rejected
        // Test the validation logic directly
        let fee: u16 = 1001;
        assert!(fee > MAX_TOTAL_FEE_BPS);
    }

    #[test]
    fn test_instantiate_fee_cap_allows_1000() {
        // Base swap_fee_bps == 1000 (10%) should be allowed
        let fee: u16 = 1000;
        assert!(fee <= MAX_TOTAL_FEE_BPS);
    }

    #[test]
    fn test_instantiate_fee_cap_allows_zero() {
        // 0% fee should be allowed
        let fee: u16 = 0;
        assert!(fee <= MAX_TOTAL_FEE_BPS);
    }

    #[test]
    fn test_total_fee_cap_base_plus_augmented_at_limit() {
        // Base 900 bps + augmented 100 bps = 1000 bps (exactly at cap, should pass)
        let base_fee: u16 = 900;
        let augmented_fee: u16 = 100;
        let total = base_fee + augmented_fee;
        assert!(total <= MAX_TOTAL_FEE_BPS);
        assert_eq!(total, 1000);
    }

    #[test]
    fn test_total_fee_cap_base_plus_augmented_over_limit() {
        // Base 900 bps + augmented 101 bps = 1001 bps (over cap, should fail)
        let base_fee: u16 = 900;
        let augmented_fee: u16 = 101;
        let total = base_fee + augmented_fee;
        assert!(total > MAX_TOTAL_FEE_BPS);
    }

    #[test]
    fn test_total_fee_cap_typical_config() {
        // Typical config: 100 bps base + 100 bps augmented = 200 bps (well under cap)
        let base_fee: u16 = 100;
        let augmented_fee: u16 = 100;
        let total = base_fee + augmented_fee;
        assert!(total <= MAX_TOTAL_FEE_BPS);
        assert_eq!(total, 200);
    }

    #[test]
    fn test_total_fee_cap_max_augmented_with_base() {
        // Max individual augmented (500) + reasonable base (100) = 600 bps (under cap)
        let base_fee: u16 = 100;
        let augmented_fee: u16 = 500;
        let total = base_fee + augmented_fee;
        assert!(total <= MAX_TOTAL_FEE_BPS);
    }

    #[test]
    fn test_total_fee_cap_max_augmented_with_high_base() {
        // Max individual augmented (500) + high base (501) = 1001 bps (over cap)
        // Note: augmented_fee_bps <= 500 check passes, but total fee cap fails
        let base_fee: u16 = 501;
        let augmented_fee: u16 = 500;
        let total = base_fee + augmented_fee;
        assert!(total > MAX_TOTAL_FEE_BPS);
    }

    #[test]
    fn test_fee_cap_error_has_descriptive_fields() {
        // Verify the error type carries the right information
        let err = ContractError::TotalFeeCapExceeded {
            swap_fee_bps: 600,
            augmented_fee_bps: 500,
            total: 1100,
        };
        let msg = format!("{}", err);
        assert!(msg.contains("1000"));  // mentions the cap
        assert!(msg.contains("600"));   // mentions the base fee
        assert!(msg.contains("500"));   // mentions the augmented fee
        assert!(msg.contains("1100"));  // mentions the total
    }

}