cw-cyber/contracts/litium-stake/src/contract.rs

use cosmwasm_std::{
    entry_point, to_json_binary, Addr, Binary, CosmosMsg, Deps, DepsMut, Empty, Env, MessageInfo,
    Response, StdResult, Uint128, Uint256, WasmMsg,
};
use cw2::{get_contract_version, set_contract_version};
use cw20::{Cw20ExecuteMsg, Cw20ReceiveMsg};

use crate::error::ContractError;
use crate::msg::{
    ConfigResponse, ExecuteMsg, InstantiateMsg, QueryMsg, StakeInfoResponse, StakingStatsResponse,
    TestingOverrides, TotalPendingRewardsResponse, TotalStakedResponse,
};
use crate::state::{
    StakeConfig, StakerInfo, UnbondingBatch, CONFIG, STAKERS, STAKING_RESERVE,
    STAKING_REWARD_INDEX, STAKING_TOTAL_STAKED, TOTAL_ACCRUED_REWARDS, TOTAL_CLAIMED_REWARDS,
};

const CONTRACT_NAME: &str = "crates.io:litium-stake";
const CONTRACT_VERSION: &str = env!("CARGO_PKG_VERSION");

/// Default unbonding: 21 days per spec
pub const DEFAULT_UNBONDING_PERIOD_SECONDS: u64 = 1_814_400;
/// Minimum stake: 1 LI = 1_000_000 atomic units.
const MIN_STAKE_ATOMIC: u128 = 1_000_000;
const STAKING_INDEX_SCALE: u128 = 1_000_000_000_000u128;

#[derive(serde::Serialize)]
#[serde(rename_all = "snake_case")]
enum LitiumCoreExecuteMsg {
    Mint { to: String, amount: Uint128 },
}

#[cfg_attr(not(feature = "library"), entry_point)]
pub fn instantiate(
    deps: DepsMut,
    _env: Env,
    info: MessageInfo,
    msg: InstantiateMsg,
) -> Result<Response, ContractError> {
    set_contract_version(deps.storage, CONTRACT_NAME, CONTRACT_VERSION)?;

    let admin = msg
        .admin
        .map(|a| deps.api.addr_validate(&a))
        .transpose()?
        .unwrap_or(info.sender.clone());

    let config = StakeConfig {
        core_contract: deps.api.addr_validate(&msg.core_contract)?,
        mine_contract: deps.api.addr_validate(&msg.mine_contract)?,
        token_contract: deps.api.addr_validate(&msg.token_contract)?,
        unbonding_period_seconds: msg
            .unbonding_period_seconds
            .unwrap_or(DEFAULT_UNBONDING_PERIOD_SECONDS),
        admin,
        paused: false,
    };
    CONFIG.save(deps.storage, &config)?;
    STAKING_RESERVE.save(deps.storage, &Uint128::zero())?;
    STAKING_TOTAL_STAKED.save(deps.storage, &Uint128::zero())?;
    STAKING_REWARD_INDEX.save(deps.storage, &Uint256::zero())?;
    TOTAL_ACCRUED_REWARDS.save(deps.storage, &Uint128::zero())?;
    TOTAL_CLAIMED_REWARDS.save(deps.storage, &Uint128::zero())?;

    Ok(Response::new()
        .add_attribute("action", "instantiate")
        .add_attribute("admin", config.admin.to_string()))
}

#[cfg_attr(not(feature = "library"), entry_point)]
pub fn execute(
    deps: DepsMut,
    env: Env,
    info: MessageInfo,
    msg: ExecuteMsg,
) -> Result<Response, ContractError> {
    match msg {
        ExecuteMsg::AccrueReward { amount } => execute_accrue_reward(deps, info, amount),
        ExecuteMsg::Receive(cw20_msg) => execute_receive(deps, info, cw20_msg),
        ExecuteMsg::Unstake { amount } => execute_unstake(deps, env, info, amount),
        ExecuteMsg::ClaimUnbonding {} => execute_claim_unbonding(deps, env, info),
        ExecuteMsg::ClaimStakingRewards {} => execute_claim_staking_rewards(deps, info),
        ExecuteMsg::UpdateConfig {
            core_contract,
            mine_contract,
            token_contract,
            unbonding_period_seconds,
            admin,
        } => execute_update_config(
            deps,
            info,
            core_contract,
            mine_contract,
            token_contract,
            unbonding_period_seconds,
            admin,
        ),
        ExecuteMsg::ApplyTestingOverrides { overrides } => {
            execute_apply_testing_overrides(deps, info, overrides)
        }
        ExecuteMsg::ResetState {} => execute_reset_state(deps, info),
        ExecuteMsg::Pause {} => execute_pause(deps, info),
        ExecuteMsg::Unpause {} => execute_unpause(deps, info),
    }
}

fn execute_accrue_reward(
    deps: DepsMut,
    info: MessageInfo,
    amount: Uint128,
) -> Result<Response, ContractError> {
    let config = CONFIG.load(deps.storage)?;
    if config.paused {
        return Err(ContractError::Paused {});
    }
    if info.sender != config.mine_contract {
        return Err(ContractError::NotMineContract {});
    }

    accrue_staking_reward(deps.storage, amount)?;

    Ok(Response::new()
        .add_attribute("action", "accrue_reward")
        .add_attribute("amount", amount.to_string()))
}

/// Receive CW-20 tokens to stake. Called by litium-core CW-20 contract.
fn execute_receive(
    deps: DepsMut,
    info: MessageInfo,
    cw20_msg: Cw20ReceiveMsg,
) -> Result<Response, ContractError> {
    let config = CONFIG.load(deps.storage)?;
    if config.paused {
        return Err(ContractError::Paused {});
    }

    // Only accept tokens from the token_contract (litium-core CW-20)
    if info.sender != config.token_contract {
        return Err(ContractError::UnexpectedFunds {});
    }

    let amount = cw20_msg.amount;
    if amount < Uint128::from(MIN_STAKE_ATOMIC) {
        return Err(ContractError::InvalidStakeAmount {});
    }

    let staker_addr = deps.api.addr_validate(&cw20_msg.sender)?;

    accrue_staking_reward(deps.storage, Uint128::zero())?;
    let global_index = STAKING_REWARD_INDEX.load(deps.storage)?;
    let mut staker = load_staker(deps.storage, &staker_addr, global_index)?;
    settle_staker_rewards(&mut staker, global_index);
    staker.staked_amount += amount;
    staker.reward_index_snapshot = global_index;
    STAKERS.save(deps.storage, &staker_addr, &staker)?;

    STAKING_TOTAL_STAKED.update(deps.storage, |v| -> StdResult<_> { Ok(v + amount) })?;

    Ok(Response::new()
        .add_attribute("action", "stake")
        .add_attribute("address", staker_addr.to_string())
        .add_attribute("amount", amount.to_string()))
}

fn execute_unstake(
    deps: DepsMut,
    env: Env,
    info: MessageInfo,
    amount: Uint128,
) -> Result<Response, ContractError> {
    let config = CONFIG.load(deps.storage)?;
    if config.paused {
        return Err(ContractError::Paused {});
    }
    if amount.is_zero() {
        return Err(ContractError::InvalidStakeAmount {});
    }

    accrue_staking_reward(deps.storage, Uint128::zero())?;
    let global_index = STAKING_REWARD_INDEX.load(deps.storage)?;
    let mut staker = load_staker(deps.storage, &info.sender, global_index)?;
    settle_staker_rewards(&mut staker, global_index);
    if staker.staked_amount < amount {
        return Err(ContractError::InsufficientStake {});
    }
    let new_staked = staker.staked_amount.saturating_sub(amount);
    if !new_staked.is_zero() && new_staked < Uint128::from(MIN_STAKE_ATOMIC) {
        return Err(ContractError::InvalidStakeAmount {});
    }
    staker.staked_amount = new_staked;
    let unlock_at = env
        .block
        .time
        .seconds()
        .saturating_add(config.unbonding_period_seconds.max(1));
    staker
        .unbonding_batches
        .push(UnbondingBatch { amount, unlock_at });
    staker.reward_index_snapshot = global_index;
    STAKERS.save(deps.storage, &info.sender, &staker)?;

    STAKING_TOTAL_STAKED.update(deps.storage, |v| -> StdResult<_> {
        Ok(v.saturating_sub(amount))
    })?;

    let total_pending: Uint128 = staker.unbonding_batches.iter().map(|b| b.amount).sum();

    Ok(Response::new()
        .add_attribute("action", "unstake")
        .add_attribute("address", info.sender.to_string())
        .add_attribute("amount", amount.to_string())
        .add_attribute("pending_unbonding", total_pending.to_string())
        .add_attribute("unlock_at", unlock_at.to_string()))
}

fn execute_claim_unbonding(
    deps: DepsMut,
    env: Env,
    info: MessageInfo,
) -> Result<Response, ContractError> {
    let config = CONFIG.load(deps.storage)?;
    if config.paused {
        return Err(ContractError::Paused {});
    }

    let global_index = STAKING_REWARD_INDEX.load(deps.storage)?;
    let mut staker = load_staker(deps.storage, &info.sender, global_index)?;
    let now = env.block.time.seconds();

    // Partition batches: mature ones are claimed, immature ones stay
    let mut claimed_amount = Uint128::zero();
    let mut remaining = Vec::new();
    for batch in staker.unbonding_batches.drain(..) {
        if now >= batch.unlock_at {
            claimed_amount += batch.amount;
        } else {
            remaining.push(batch);
        }
    }
    staker.unbonding_batches = remaining;

    if claimed_amount.is_zero() {
        let next_unlock = staker
            .unbonding_batches
            .iter()
            .map(|b| b.unlock_at)
            .min()
            .unwrap_or(0);
        return Ok(Response::new()
            .add_attribute("action", "claim_unbonding")
            .add_attribute("address", info.sender.to_string())
            .add_attribute("claimed", "0")
            .add_attribute("next_unlock", next_unlock.to_string()));
    }

    STAKERS.save(deps.storage, &info.sender, &staker)?;
    let amount = claimed_amount;

    // Transfer CW-20 tokens from this contract to user (no burn — stake is authorized caller)
    let transfer_msg = CosmosMsg::Wasm(WasmMsg::Execute {
        contract_addr: config.token_contract.to_string(),
        msg: cosmwasm_std::to_json_binary(&Cw20ExecuteMsg::Transfer {
            recipient: info.sender.to_string(),
            amount,
        })?,
        funds: vec![],
    });

    Ok(Response::new()
        .add_message(transfer_msg)
        .add_attribute("action", "claim_unbonding")
        .add_attribute("address", info.sender.to_string())
        .add_attribute("claimed", amount.to_string()))
}

fn execute_claim_staking_rewards(
    deps: DepsMut,
    info: MessageInfo,
) -> Result<Response, ContractError> {
    let config = CONFIG.load(deps.storage)?;
    if config.paused {
        return Err(ContractError::Paused {});
    }

    accrue_staking_reward(deps.storage, Uint128::zero())?;
    let global_index = STAKING_REWARD_INDEX.load(deps.storage)?;
    let mut staker = load_staker(deps.storage, &info.sender, global_index)?;
    settle_staker_rewards(&mut staker, global_index);
    let claimed = staker.claimable_rewards;
    staker.claimable_rewards = Uint128::zero();
    staker.reward_index_snapshot = global_index;
    STAKERS.save(deps.storage, &info.sender, &staker)?;

    let mut resp = Response::new()
        .add_attribute("action", "claim_staking_rewards")
        .add_attribute("address", info.sender.to_string())
        .add_attribute("claimed", claimed.to_string());

    if !claimed.is_zero() {
        TOTAL_CLAIMED_REWARDS.update(deps.storage, |v| -> StdResult<_> { Ok(v + claimed) })?;
        // Mint via litium-core
        let mint_msg = CosmosMsg::Wasm(WasmMsg::Execute {
            contract_addr: config.core_contract.to_string(),
            msg: cosmwasm_std::to_json_binary(&LitiumCoreExecuteMsg::Mint {
                to: info.sender.to_string(),
                amount: claimed,
            })?,
            funds: vec![],
        });
        resp = resp.add_message(mint_msg);
    }
    Ok(resp)
}

fn execute_update_config(
    deps: DepsMut,
    info: MessageInfo,
    core_contract: Option<String>,
    mine_contract: Option<String>,
    token_contract: Option<String>,
    unbonding_period_seconds: Option<u64>,
    admin: Option<String>,
) -> Result<Response, ContractError> {
    let mut config = CONFIG.load(deps.storage)?;
    if info.sender != config.admin {
        return Err(ContractError::Unauthorized {});
    }
    if let Some(c) = core_contract {
        config.core_contract = deps.api.addr_validate(&c)?;
    }
    if let Some(m) = mine_contract {
        config.mine_contract = deps.api.addr_validate(&m)?;
    }
    if let Some(t) = token_contract {
        config.token_contract = deps.api.addr_validate(&t)?;
    }
    if let Some(u) = unbonding_period_seconds {
        config.unbonding_period_seconds = u.max(1);
    }
    if let Some(a) = admin {
        config.admin = deps.api.addr_validate(&a)?;
    }
    CONFIG.save(deps.storage, &config)?;
    Ok(Response::new().add_attribute("action", "update_config"))
}

fn execute_apply_testing_overrides(
    deps: DepsMut,
    info: MessageInfo,
    overrides: TestingOverrides,
) -> Result<Response, ContractError> {
    let mut config = CONFIG.load(deps.storage)?;
    if info.sender != config.admin {
        return Err(ContractError::Unauthorized {});
    }

    if let Some(paused) = overrides.paused {
        config.paused = paused;
    }
    if let Some(unbonding) = overrides.unbonding_period_seconds {
        config.unbonding_period_seconds = unbonding.max(1);
    }
    CONFIG.save(deps.storage, &config)?;

    if let Some(v) = overrides.staking_reserve {
        STAKING_RESERVE.save(deps.storage, &v)?;
    }
    if let Some(v) = overrides.staking_total_staked {
        STAKING_TOTAL_STAKED.save(deps.storage, &v)?;
    }
    if let Some(v) = overrides.staking_reward_index {
        STAKING_REWARD_INDEX.save(deps.storage, &v)?;
    }
    if let Some(v) = overrides.total_accrued_rewards {
        TOTAL_ACCRUED_REWARDS.save(deps.storage, &v)?;
    }
    if let Some(v) = overrides.total_claimed_rewards {
        TOTAL_CLAIMED_REWARDS.save(deps.storage, &v)?;
    }

    Ok(Response::new().add_attribute("action", "apply_testing_overrides"))
}

/// Full state reset for daily testing. Zeroes all staking state.
fn execute_reset_state(deps: DepsMut, info: MessageInfo) -> Result<Response, ContractError> {
    let config = CONFIG.load(deps.storage)?;
    if info.sender != config.admin {
        return Err(ContractError::Unauthorized {});
    }

    // Zero global staking state
    STAKING_RESERVE.save(deps.storage, &Uint128::zero())?;
    STAKING_TOTAL_STAKED.save(deps.storage, &Uint128::zero())?;
    STAKING_REWARD_INDEX.save(deps.storage, &Uint256::zero())?;
    TOTAL_ACCRUED_REWARDS.save(deps.storage, &Uint128::zero())?;
    TOTAL_CLAIMED_REWARDS.save(deps.storage, &Uint128::zero())?;

    // Clear all staker entries
    let stakers: Vec<_> = STAKERS
        .keys(deps.storage, None, None, cosmwasm_std::Order::Ascending)
        .collect::<Result<Vec<_>, _>>()?;
    for addr in &stakers {
        STAKERS.remove(deps.storage, addr);
    }

    Ok(Response::new()
        .add_attribute("action", "reset_state")
        .add_attribute("stakers_cleared", stakers.len().to_string()))
}

fn execute_pause(deps: DepsMut, info: MessageInfo) -> Result<Response, ContractError> {
    let mut config = CONFIG.load(deps.storage)?;
    if info.sender != config.admin {
        return Err(ContractError::Unauthorized {});
    }
    config.paused = true;
    CONFIG.save(deps.storage, &config)?;
    Ok(Response::new().add_attribute("action", "pause"))
}

fn execute_unpause(deps: DepsMut, info: MessageInfo) -> Result<Response, ContractError> {
    let mut config = CONFIG.load(deps.storage)?;
    if info.sender != config.admin {
        return Err(ContractError::Unauthorized {});
    }
    config.paused = false;
    CONFIG.save(deps.storage, &config)?;
    Ok(Response::new().add_attribute("action", "unpause"))
}

// ============================================================
// Migrate
// ============================================================

#[cfg_attr(not(feature = "library"), entry_point)]
pub fn migrate(deps: DepsMut, _env: Env, _msg: Empty) -> Result<Response, ContractError> {
    let stored = get_contract_version(deps.storage)?;
    let stored_ver = stored
        .version
        .parse::<semver::Version>()
        .map_err(|_| ContractError::MigrationError {})?;
    let current_ver = CONTRACT_VERSION
        .parse::<semver::Version>()
        .map_err(|_| ContractError::MigrationError {})?;
    if stored_ver >= current_ver {
        return Err(ContractError::MigrationError {});
    }
    // v0.1.1 → v0.1.2: migrate TOTAL_PENDING_REWARDS → monotonic counters
    if TOTAL_ACCRUED_REWARDS.may_load(deps.storage)?.is_none() {
        let old_pending: Uint128 = deps
            .storage
            .get(b"total_pending_rewards")
            .and_then(|v| cosmwasm_std::from_json(&v).ok())
            .unwrap_or_default();
        TOTAL_ACCRUED_REWARDS.save(deps.storage, &old_pending)?;
        TOTAL_CLAIMED_REWARDS.save(deps.storage, &Uint128::zero())?;
        deps.storage.remove(b"total_pending_rewards");
    }

    set_contract_version(deps.storage, CONTRACT_NAME, CONTRACT_VERSION)?;
    Ok(Response::default())
}

// ============================================================
// Staking Reward Index
// ============================================================

fn accrue_staking_reward(
    storage: &mut dyn cosmwasm_std::Storage,
    newly_accrued: Uint128,
) -> Result<(), ContractError> {
    if !newly_accrued.is_zero() {
        TOTAL_ACCRUED_REWARDS.update(storage, |v| -> StdResult<_> { Ok(v + newly_accrued) })?;
    }
    let mut reserve = STAKING_RESERVE.load(storage)?;
    reserve += newly_accrued;

    let total_staked = STAKING_TOTAL_STAKED.load(storage)?;
    if total_staked.is_zero() || reserve.is_zero() {
        STAKING_RESERVE.save(storage, &reserve)?;
        return Ok(());
    }

    let mut index = STAKING_REWARD_INDEX.load(storage)?;
    let delta_index = Uint256::from(reserve)
        .saturating_mul(Uint256::from(STAKING_INDEX_SCALE))
        .checked_div(Uint256::from(total_staked))
        .unwrap_or_else(|_| Uint256::zero());

    if !delta_index.is_zero() {
        let distributed_u256 = Uint256::from(total_staked)
            .saturating_mul(delta_index)
            .checked_div(Uint256::from(STAKING_INDEX_SCALE))
            .unwrap_or_else(|_| Uint256::zero());
        let distributed = u256_to_u128_floor(distributed_u256);
        reserve = reserve.saturating_sub(distributed);
        index += delta_index;
        STAKING_REWARD_INDEX.save(storage, &index)?;
    }

    STAKING_RESERVE.save(storage, &reserve)?;
    Ok(())
}

fn settle_staker_rewards(staker: &mut StakerInfo, global_index: Uint256) {
    if staker.staked_amount.is_zero() || global_index <= staker.reward_index_snapshot {
        staker.reward_index_snapshot = global_index;
        return;
    }
    let delta = global_index - staker.reward_index_snapshot;
    let accrued = u256_to_u128_floor(
        Uint256::from(staker.staked_amount)
            .saturating_mul(delta)
            .checked_div(Uint256::from(STAKING_INDEX_SCALE))
            .unwrap_or_else(|_| Uint256::zero()),
    );
    staker.claimable_rewards += accrued;
    staker.reward_index_snapshot = global_index;
}

fn u256_to_u128_floor(v: Uint256) -> Uint128 {
    if v > Uint256::from(Uint128::MAX) {
        Uint128::MAX
    } else {
        Uint128::try_from(v).unwrap_or(Uint128::zero())
    }
}

fn load_staker(
    storage: &dyn cosmwasm_std::Storage,
    addr: &Addr,
    current_index: Uint256,
) -> Result<StakerInfo, ContractError> {
    Ok(STAKERS.may_load(storage, addr)?.unwrap_or(StakerInfo {
        staked_amount: Uint128::zero(),
        unbonding_batches: vec![],
        claimable_rewards: Uint128::zero(),
        reward_index_snapshot: current_index,
    }))
}

// ============================================================
// Queries
// ============================================================

#[cfg_attr(not(feature = "library"), entry_point)]
pub fn query(deps: Deps, _env: Env, msg: QueryMsg) -> StdResult<Binary> {
    match msg {
        QueryMsg::Config {} => to_json_binary(&query_config(deps)?),
        QueryMsg::TotalStaked {} => to_json_binary(&query_total_staked(deps)?),
        QueryMsg::StakeInfo { address } => to_json_binary(&query_stake_info(deps, address)?),
        QueryMsg::StakingStats {} => to_json_binary(&query_staking_stats(deps)?),
        QueryMsg::TotalPendingRewards {} => to_json_binary(&query_total_pending_rewards(deps)?),
    }
}

fn query_config(deps: Deps) -> StdResult<ConfigResponse> {
    let config = CONFIG.load(deps.storage)?;
    Ok(ConfigResponse {
        core_contract: config.core_contract.to_string(),
        mine_contract: config.mine_contract.to_string(),
        token_contract: config.token_contract.to_string(),
        unbonding_period_seconds: config.unbonding_period_seconds,
        admin: config.admin.to_string(),
        paused: config.paused,
    })
}

fn query_total_staked(deps: Deps) -> StdResult<TotalStakedResponse> {
    let total_staked = STAKING_TOTAL_STAKED.load(deps.storage)?;
    Ok(TotalStakedResponse { total_staked })
}

fn query_stake_info(deps: Deps, address: String) -> StdResult<StakeInfoResponse> {
    let addr = deps.api.addr_validate(&address)?;
    let global_index = STAKING_REWARD_INDEX
        .may_load(deps.storage)?
        .unwrap_or_else(Uint256::zero);
    let mut staker = STAKERS
        .may_load(deps.storage, &addr)?
        .unwrap_or(StakerInfo {
            staked_amount: Uint128::zero(),
            unbonding_batches: vec![],
            claimable_rewards: Uint128::zero(),
            reward_index_snapshot: global_index,
        });
    settle_staker_rewards(&mut staker, global_index);

    let total_pending: Uint128 = staker.unbonding_batches.iter().map(|b| b.amount).sum();
    let latest_until = staker
        .unbonding_batches
        .iter()
        .map(|b| b.unlock_at)
        .max()
        .unwrap_or(0);

    Ok(StakeInfoResponse {
        address,
        staked_amount: staker.staked_amount,
        pending_unbonding: total_pending,
        pending_unbonding_until: latest_until,
        claimable_rewards: staker.claimable_rewards,
    })
}

fn query_staking_stats(deps: Deps) -> StdResult<StakingStatsResponse> {
    let reserve = STAKING_RESERVE.load(deps.storage)?;
    let total_staked = STAKING_TOTAL_STAKED.load(deps.storage)?;
    let reward_index = STAKING_REWARD_INDEX.load(deps.storage)?;
    Ok(StakingStatsResponse {
        reserve,
        total_staked,
        reward_index: reward_index.to_string(),
    })
}

fn query_total_pending_rewards(deps: Deps) -> StdResult<TotalPendingRewardsResponse> {
    let accrued = TOTAL_ACCRUED_REWARDS
        .may_load(deps.storage)?
        .unwrap_or_else(Uint128::zero);
    let claimed = TOTAL_CLAIMED_REWARDS
        .may_load(deps.storage)?
        .unwrap_or_else(Uint128::zero);
    Ok(TotalPendingRewardsResponse {
        total_pending_rewards: accrued.saturating_sub(claimed),
    })
}